package tech.easyflow.core.chat.protocol.sse;

import com.alibaba.fastjson.JSON;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.web.context.request.async.AsyncRequestNotUsableException;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import tech.easyflow.common.util.SpringContextUtil;
import tech.easyflow.common.util.StringUtil;
import tech.easyflow.core.chat.protocol.ChatEnvelope;

import java.io.IOException;
import java.time.Duration;
import java.util.concurrent.atomic.AtomicBoolean;

public class ChatSseEmitter {

    private static final Logger LOG = LoggerFactory.getLogger(ChatSseEmitter.class);
    private static final long DEFAULT_TIMEOUT = Duration.ofMinutes(10).toMillis();
    private static final String PROP_SSE_TIMEOUT_MS = "easyflow.chat.sse-timeout-ms";
    private static final String PROP_SSE_TIMEOUT_MS_ALT = "easyflow.chat.sse.timeout-ms";

    private final SseEmitter emitter;
    private final AtomicBoolean closed = new AtomicBoolean(false);

    public ChatSseEmitter() {
        this(resolveTimeoutMillis());
    }

    public ChatSseEmitter(long timeoutMillis) {
        this.emitter = new SseEmitter(timeoutMillis);
        registerLifecycleCallbacks();
    }

    public SseEmitter getEmitter() {
        return emitter;
    }

    /** 发送普通 ChatEnvelope（event: message） */
    public boolean send(ChatEnvelope<?> envelope) {
        return send("message", envelope);
    }

    /** 发送 error 事件 */
    public boolean sendError(ChatEnvelope<?> envelope) {
        return send("error", envelope);
    }

    /** 发送 done 事件并关闭 */
    public boolean sendDone(ChatEnvelope<?> envelope) {
        boolean sent = send("done", envelope);
        complete();
        return sent;
    }

    /** 🔥 新增：发送并立即关闭 */
    public boolean sendAndClose(ChatEnvelope<?> envelope) {
        boolean sent = send("message", envelope);
        if (!sent) {
            return false;
        }
        ThreadPoolTaskExecutor threadPoolTaskExecutor = SpringContextUtil.getBean("sseThreadPool");
        threadPoolTaskExecutor.execute(() -> {
            try {
                Thread.sleep(500);
                complete();
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                LOG.error("ChatSseEmitter sendAndClose interrupted, message={}, exception={}", e.getMessage(), e.toString(), e);
            }
        });
        return true;
    }

    /** 通知前端保存该消息 */
    public boolean sendMessageNeedSave(ChatEnvelope<?> envelope) {
        return send("needSaveMessage", envelope);
    }

    /** SSE 底层发送 */
    private boolean send(String event, ChatEnvelope<?> envelope) {
        if (closed.get()) {
            return false;
        }
        try {
            String json = JSON.toJSONString(envelope);
            emitter.send(
                    SseEmitter.event()
                            .name(event)
                            .data(json)
            );
            return true;
        } catch (AsyncRequestNotUsableException e) {
            markDisconnected(event, e);
            return false;
        } catch (IOException e) {
            markDisconnected(event, e);
            return false;
        } catch (IllegalStateException e) {
            closed.compareAndSet(false, true);
            LOG.warn("ChatSseEmitter send failed(event={}), message={}, exception={}", event, e.getMessage(), e.toString());
            return false;
        } catch (Exception e) {
            if (isClientDisconnected(e)) {
                markDisconnected(event, e);
                return false;
            }
            LOG.error("ChatSseEmitter send unexpected failed(event={}), message={}, exception={}", event, e.getMessage(), e.toString(), e);
            safeCompleteWithError(e);
            return false;
        }
    }

    public void complete() {
        if (closed.compareAndSet(false, true)) {
            emitter.complete();
        }
    }

    public void completeWithError(Throwable ex) {
        if (ex == null) {
            complete();
            return;
        }
        safeCompleteWithError(ex);
    }

    public boolean isClosed() {
        return closed.get();
    }

    private static long resolveTimeoutMillis() {
        Long fromProp = SpringContextUtil.getProperty(PROP_SSE_TIMEOUT_MS, Long.class, null);
        if (fromProp == null) {
            fromProp = SpringContextUtil.getProperty(PROP_SSE_TIMEOUT_MS_ALT, Long.class, null);
        }
        if (fromProp != null && fromProp > 0) {
            return fromProp;
        }
        String fromString = SpringContextUtil.getProperty(PROP_SSE_TIMEOUT_MS);
        if (StringUtil.noText(fromString)) {
            fromString = SpringContextUtil.getProperty(PROP_SSE_TIMEOUT_MS_ALT);
        }
        if (StringUtil.hasText(fromString)) {
            try {
                long parsed = Long.parseLong(fromString.trim());
                if (parsed > 0) {
                    return parsed;
                }
            } catch (Exception e) {
                LOG.error("Invalid sse timeout config: key={}, value={}, message={}, exception={}",
                        PROP_SSE_TIMEOUT_MS, fromString, e.getMessage(), e.toString(), e);
            }
        }
        return DEFAULT_TIMEOUT;
    }

    private void registerLifecycleCallbacks() {
        emitter.onCompletion(() -> closed.compareAndSet(false, true));
        emitter.onTimeout(() -> closed.compareAndSet(false, true));
        emitter.onError(ex -> closed.compareAndSet(false, true));
    }

    private void safeCompleteWithError(Throwable ex) {
        if (closed.compareAndSet(false, true)) {
            try {
                emitter.completeWithError(ex);
            } catch (Exception completeEx) {
                LOG.error("ChatSseEmitter completeWithError failed, message={}, exception={}",
                        completeEx.getMessage(), completeEx.toString(), completeEx);
            }
        }
    }

    private void markDisconnected(String event, Throwable ex) {
        closed.compareAndSet(false, true);
        LOG.warn("ChatSseEmitter client disconnected(event={}), message={}, exception={}",
                event, ex == null ? null : ex.getMessage(), ex == null ? null : ex.toString());
    }

    private boolean isClientDisconnected(Throwable ex) {
        Throwable current = ex;
        while (current != null) {
            if (current instanceof AsyncRequestNotUsableException || current instanceof IOException) {
                return true;
            }
            String message = current.getMessage();
            if (message != null) {
                String lowerMessage = message.toLowerCase();
                if (lowerMessage.contains("broken pipe")
                        || lowerMessage.contains("connection reset")
                        || lowerMessage.contains("disconnected client")
                        || lowerMessage.contains("response not usable")) {
                    return true;
                }
            }
            current = current.getCause();
        }
        return false;
    }
}
