package tech.easyflow.agent.runtime.agui;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * 以共享调度器为本节点上的 AG-UI 重连订阅者增量重放 Redis 事件。
 *
 * <p>所有订阅共享一个短周期任务，不为每个浏览器连接创建阻塞线程。</p>
 */
@Component
public class AgentAguiRunSubscriptionService {

    private static final Logger LOG = LoggerFactory.getLogger(AgentAguiRunSubscriptionService.class);
    private static final long KEEP_ALIVE_INTERVAL_MILLIS = 15_000L;

    private final AgentAguiRunStore store;
    private final Map<String, CopyOnWriteArrayList<Subscriber>> subscribers = new ConcurrentHashMap<>();

    /**
     * 创建 AG-UI 重连订阅服务。
     *
     * @param store Redis 运行存储
     */
    public AgentAguiRunSubscriptionService(AgentAguiRunStore store) {
        this.store = store;
    }

    /**
     * 注册一个从指定游标继续的 SSE 订阅。
     *
     * @param descriptor 已授权运行描述
     * @param afterCursor 已消费的最后游标
     * @return 重连 SSE 发射器
     */
    public SseEmitter subscribe(AgentAguiRunDescriptor descriptor, long afterCursor) {
        return subscribe(descriptor, afterCursor, new SseEmitter(0L));
    }

    /**
     * 使用指定发射器注册订阅，供生命周期测试复用同一调度逻辑。
     *
     * @param descriptor 已授权运行描述
     * @param afterCursor 已消费的最后游标
     * @param emitter SSE 发射器
     * @return 已注册发射器
     */
    SseEmitter subscribe(AgentAguiRunDescriptor descriptor,
                         long afterCursor,
                         SseEmitter emitter) {
        Subscriber subscriber = new Subscriber(emitter, Math.max(0L, afterCursor));
        subscribers.computeIfAbsent(descriptor.runId(), ignored -> new CopyOnWriteArrayList<>())
                .add(subscriber);
        Runnable remove = () -> remove(descriptor.runId(), subscriber);
        emitter.onCompletion(remove);
        emitter.onTimeout(remove);
        emitter.onError(error -> remove.run());
        return emitter;
    }

    /**
     * 周期性向全部本地订阅者重放新增事件并发送保活注释。
     */
    @Scheduled(fixedDelayString = "${easyflow.agent.runtime.agui-replay-poll-interval:100ms}")
    public void dispatch() {
        for (Map.Entry<String, CopyOnWriteArrayList<Subscriber>> entry : subscribers.entrySet()) {
            dispatch(entry.getKey(), entry.getValue());
        }
    }

    private void dispatch(String runId, CopyOnWriteArrayList<Subscriber> runSubscribers) {
        try {
            AgentAguiRunDescriptor descriptor = store.find(runId);
            if (descriptor == null) {
                runSubscribers.forEach(subscriber -> removeAndComplete(runId, subscriber));
                return;
            }
            Map<Long, List<Subscriber>> cursorGroups = new LinkedHashMap<>();
            for (Subscriber subscriber : runSubscribers) {
                cursorGroups.computeIfAbsent(subscriber.cursor, ignored -> new java.util.ArrayList<>())
                        .add(subscriber);
            }
            for (Map.Entry<Long, List<Subscriber>> cursorGroup : cursorGroups.entrySet()) {
                long cursor = cursorGroup.getKey();
                List<String> events = store.readAfter(runId, cursor);
                for (Subscriber subscriber : cursorGroup.getValue()) {
                    dispatchSubscriber(runId, subscriber, cursor, events, descriptor);
                }
            }
        } catch (RuntimeException exception) {
            LOG.error("AG-UI 重连订阅调度失败: runId={}, subscriberCount={}",
                    runId, runSubscribers.size(), exception);
        }
    }

    private void dispatchSubscriber(String runId,
                                    Subscriber subscriber,
                                    long afterCursor,
                                    List<String> events,
                                    AgentAguiRunDescriptor descriptor) {
        try {
            for (int index = 0; index < events.size(); index++) {
                long cursor = afterCursor + index + 1L;
                subscriber.emitter.send(SseEmitter.event().id(Long.toString(cursor)).data(events.get(index)));
                subscriber.cursor = cursor;
                subscriber.lastWriteAt = System.currentTimeMillis();
            }
            if (descriptor.status().isTerminal() && subscriber.cursor >= descriptor.lastCursor()) {
                removeAndComplete(runId, subscriber);
                return;
            }
            long now = System.currentTimeMillis();
            if (events.isEmpty() && now - subscriber.lastWriteAt >= KEEP_ALIVE_INTERVAL_MILLIS) {
                subscriber.emitter.send(SseEmitter.event().comment("keepalive"));
                subscriber.lastWriteAt = now;
            }
        } catch (IOException | IllegalStateException exception) {
            LOG.debug("移除已断开的 AG-UI 重连订阅: runId={}, cursor={}",
                    runId, subscriber.cursor, exception);
            remove(runId, subscriber);
        }
    }

    private void removeAndComplete(String runId, Subscriber subscriber) {
        remove(runId, subscriber);
        subscriber.emitter.complete();
    }

    private void remove(String runId, Subscriber subscriber) {
        CopyOnWriteArrayList<Subscriber> values = subscribers.get(runId);
        if (values == null) {
            return;
        }
        values.remove(subscriber);
        if (values.isEmpty()) {
            subscribers.remove(runId, values);
        }
    }

    private static final class Subscriber {
        private final SseEmitter emitter;
        private long cursor;
        private long lastWriteAt = System.currentTimeMillis();

        private Subscriber(SseEmitter emitter, long cursor) {
            this.emitter = emitter;
            this.cursor = cursor;
        }
    }
}
