package com.agent.management.external;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import jakarta.annotation.PreDestroy;
import java.io.IOException;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

/**
 * SSE 心跳工具（v1.2 抽取自 WorkflowController / ChatController 的重复实现）。
 *
 * <p>所有需要向 SseEmitter 周期性发送 comment 帧防止 NAT/防火墙清连接的 Controller
 * 都应调用 {@link #start(SseEmitter)}，避免每个 Controller 各持一份调度器。</p>
 *
 * <p>emitter 完成/超时/出错时通过注册的回调自动停止该 emitter 的调度任务；
 * 容器关闭时由 {@link #shutdown()} 统一关闭线程池。</p>
 */
@Slf4j
@Component
public final class SseHeartbeats {

    /** 单例调度器：所有 SSE emitter 共享。daemon 线程，2 个并发足够（聊天 + 工作流）。 */
    private static final ScheduledExecutorService SCHEDULER = Executors.newScheduledThreadPool(2, r -> {
        Thread t = new Thread(r, "sse-heartbeat");
        t.setDaemon(true);
        return t;
    });

    private SseHeartbeats() {
    }

    /**
     * 启动心跳：每 30 秒发送一次 {@code :keep-alive} comment 帧。
     * 自动在 emitter 完成/超时/出错时停止调度。
     */
    public static void start(SseEmitter emitter) {
        ScheduledFuture<?>[] handle = new ScheduledFuture<?>[1];
        Runnable stop = () -> {
            if (handle[0] != null) {
                handle[0].cancel(false);
            }
        };
        handle[0] = SCHEDULER.scheduleAtFixedRate(() -> {
            try {
                emitter.send(SseEmitter.event().comment(":keep-alive"));
            } catch (IOException e) {
                // 连接已断开，停止心跳
                stop.run();
            }
        }, 30, 30, TimeUnit.SECONDS);
        emitter.onCompletion(stop);
        emitter.onTimeout(stop);
        emitter.onError(ex -> stop.run());
    }

    /** 容器关闭时清理线程池（由 Spring 通过 @PreDestroy 触发）。 */
    @PreDestroy
    public void shutdown() {
        SCHEDULER.shutdown();
        try {
            if (!SCHEDULER.awaitTermination(5, TimeUnit.SECONDS)) {
                SCHEDULER.shutdownNow();
            }
        } catch (InterruptedException e) {
            SCHEDULER.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }
}
