package tech.easyflow.agent.runtime;

import com.easyagents.agent.runtime.AgentResumeRequest;
import com.easyagents.agent.runtime.AgentRuntime;
import com.easyagents.agent.runtime.event.AgentRuntimeEvent;
import com.easyagents.agent.runtime.hitl.AgentResumeToken;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import reactor.core.Disposable;
import tech.easyflow.agent.distributed.AgentRuntimeRouteRegistry;
import tech.easyflow.agent.runtime.lock.AgentRunLock;
import tech.easyflow.common.web.exceptions.BusinessException;
import tech.easyflow.core.chat.protocol.sse.ChatSseEmitter;
import tech.easyflow.core.runtime.ChatAssistantAccumulator;
import tech.easyflow.core.runtime.ChatRuntimeContext;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;

/**
 * Agent 运行态注册表。
 */
@Component
public class AgentRunRegistry {

    private static final Logger LOG = LoggerFactory.getLogger(AgentRunRegistry.class);

    private final Map<String, AgentRunContext> runs = new ConcurrentHashMap<>();
    private final Map<String, String> sessionRuns = new ConcurrentHashMap<>();
    private final Map<String, String> resumeTokenIndex = new ConcurrentHashMap<>();
    private final Map<String, Set<String>> requestTokens = new ConcurrentHashMap<>();
    private final Map<String, RunOwner> owners = new ConcurrentHashMap<>();
    private AgentRuntimeRouteRegistry routeRegistry;

    /**
     * 设置 Agent 运行态 Redis 路由注册表。
     *
     * @param routeRegistry Redis 路由注册表
     */
    @Autowired(required = false)
    public void setRouteRegistry(AgentRuntimeRouteRegistry routeRegistry) {
        this.routeRegistry = routeRegistry;
    }

    /**
     * 注册运行态。
     *
     * @param context 运行态上下文
     */
    public void register(AgentRunContext context) {
        if (context == null || context.requestId() == null || context.requestId().isBlank()) {
            throw new BusinessException("Agent 运行请求 ID 不能为空");
        }
        if (context.sessionId() == null || context.sessionId().isBlank()) {
            throw new BusinessException("Agent 会话 ID 不能为空");
        }
        String existingRequestId = sessionRuns.putIfAbsent(context.sessionId(), context.requestId());
        if (existingRequestId != null && !existingRequestId.equals(context.requestId())) {
            throw new BusinessException("当前 Agent 会话已有运行中的请求，请稍后再试");
        }
        AgentRunContext existing = runs.putIfAbsent(context.requestId(), context);
        if (existing != null) {
            sessionRuns.remove(context.sessionId(), context.requestId());
            throw new BusinessException("当前 Agent 运行请求已存在");
        }
        owners.put(context.requestId(), context.owner());
        if (routeRegistry != null) {
            routeRegistry.registerRun(context.requestId());
        }
    }

    /**
     * 绑定运行订阅。
     *
     * @param requestId 请求 ID
     * @param subscription Reactor 订阅
     */
    public void bindSubscription(String requestId, Disposable subscription) {
        AgentRunContext context = runs.get(requestId);
        if (context != null) {
            context.setSubscription(subscription);
        }
    }

    /**
     * 获取运行态。
     *
     * @param requestId 请求 ID
     * @return 运行态
     */
    public AgentRunContext get(String requestId) {
        return requestId == null ? null : runs.get(requestId);
    }

    /**
     * 取消并移除指定会话当前活跃运行。
     *
     * <p>草稿试运行清理会删除 AgentScope session。若同一会话仍有 SSE 运行中，
     * 必须先取消 runtime 订阅并关闭 SSE，避免旧运行继续向已清空的会话写入状态。</p>
     *
     * @param sessionId 会话 ID
     */
    public void cancelSession(String sessionId) {
        cancelSession(sessionId, null);
    }

    /**
     * 取消并移除指定会话当前活跃运行。
     *
     * @param sessionId 会话 ID
     * @param userId 当前用户 ID，非空时会校验运行归属
     */
    public void cancelSession(String sessionId, String userId) {
        if (sessionId == null || sessionId.isBlank()) {
            return;
        }
        String requestId = sessionRuns.get(sessionId);
        if (requestId == null) {
            return;
        }
        assertOwner(requestId, userId);
        AgentRunContext context = runs.get(requestId);
        if (context != null) {
            context.cancelAndComplete();
        }
        remove(requestId);
    }

    /**
     * 记录等待审批的恢复令牌。
     *
     * @param requestId 请求 ID
     * @param resumeToken 恢复令牌
     */
    public void registerResumeToken(String requestId, String resumeToken) {
        if (requestId != null && resumeToken != null && !resumeToken.isBlank()) {
            resumeTokenIndex.put(resumeToken, requestId);
            requestTokens.computeIfAbsent(requestId, ignored -> ConcurrentHashMap.newKeySet()).add(resumeToken);
            if (routeRegistry != null) {
                routeRegistry.registerResumeToken(requestId, resumeToken);
            }
        }
    }

    /**
     * 运行结束后移除运行态。
     *
     * @param requestId 请求 ID
     */
    public synchronized void remove(String requestId) {
        if (requestId == null) {
            return;
        }
        AgentRunContext context = runs.remove(requestId);
        if (context != null) {
            sessionRuns.remove(context.sessionId(), requestId);
            context.releaseLock();
            context.closeRuntime();
        }
        owners.remove(requestId);
        Set<String> tokens = requestTokens.remove(requestId);
        if (tokens != null) {
            tokens.forEach(token -> {
                resumeTokenIndex.remove(token);
                if (routeRegistry != null) {
                    routeRegistry.removeResumeToken(token);
                }
            });
        }
        if (routeRegistry != null) {
            routeRegistry.removeRun(requestId);
        }
    }

    /**
     * 批准工具执行。
     *
     * @param requestId 请求 ID
     * @param resumeToken 恢复令牌
     * @param userId 当前用户 ID
     */
    public void approve(String requestId, String resumeToken, String userId) {
        submit(requestId, resumeToken, userId, true, null);
    }

    /**
     * 批准工具执行，并在恢复 runtime 前执行持久化消费动作。
     *
     * @param requestId 请求 ID
     * @param resumeToken 恢复令牌
     * @param userId 当前用户 ID
     * @param beforeResume runtime resume 前的持久化消费动作
     */
    public void approve(String requestId, String resumeToken, String userId, Runnable beforeResume) {
        submit(requestId, resumeToken, userId, true, null, beforeResume);
    }

    /**
     * 拒绝工具执行。
     *
     * @param requestId 请求 ID
     * @param resumeToken 恢复令牌
     * @param userId 当前用户 ID
     * @param reason 拒绝原因
     */
    public void reject(String requestId, String resumeToken, String userId, String reason) {
        submit(requestId, resumeToken, userId, false, reason);
    }

    /**
     * 拒绝工具执行，并在恢复 runtime 前执行持久化消费动作。
     *
     * @param requestId 请求 ID
     * @param resumeToken 恢复令牌
     * @param userId 当前用户 ID
     * @param reason 拒绝原因
     * @param beforeResume runtime resume 前的持久化消费动作
     */
    public void reject(String requestId, String resumeToken, String userId, String reason, Runnable beforeResume) {
        submit(requestId, resumeToken, userId, false, reason, beforeResume);
    }

    /**
     * 当前进程是否存在指定 token 对应的活跃运行态。
     *
     * @param requestId 请求 ID，可为空
     * @param resumeToken 恢复令牌
     * @return true 表示当前节点可直接恢复
     */
    public boolean containsResumeTarget(String requestId, String resumeToken) {
        try {
            String resolvedRequestId = resolveRequestId(requestId, resumeToken);
            return runs.containsKey(resolvedRequestId);
        } catch (BusinessException ignored) {
            return false;
        }
    }

    /**
     * 当前恢复目标是否为草稿试运行。
     *
     * @param requestId 请求 ID，可为空
     * @param resumeToken 恢复令牌
     * @return true 表示目标为草稿试运行
     */
    public boolean isDraftResumeTarget(String requestId, String resumeToken) {
        try {
            String resolvedRequestId = resolveRequestId(requestId, resumeToken);
            AgentRunContext context = runs.get(resolvedRequestId);
            return context != null && !context.persistChatlog();
        } catch (BusinessException ignored) {
            return false;
        }
    }

    private void submit(String requestId, String resumeToken, String userId, boolean approved, String reason) {
        submit(requestId, resumeToken, userId, approved, reason, null);
    }

    private synchronized void submit(String requestId,
                                     String resumeToken,
                                     String userId,
                                     boolean approved,
                                     String reason,
                                     Runnable beforeResume) {
        String resolvedRequestId = resolveRequestId(requestId, resumeToken);
        AgentRunContext context = runs.get(resolvedRequestId);
        if (context == null) {
            throw new BusinessException("当前 Agent 运行已结束或不在本进程内");
        }
        assertOwner(resolvedRequestId, userId);
        assertResumeTokenBelongsToRequest(resolvedRequestId, resumeToken);
        if (beforeResume != null) {
            beforeResume.run();
        }
        Set<String> tokens = requestTokens.get(resolvedRequestId);
        if (tokens != null) {
            tokens.remove(resumeToken);
        }
        resumeTokenIndex.remove(resumeToken);
        if (routeRegistry != null) {
            routeRegistry.removeResumeToken(resumeToken);
        }
        AgentResumeToken token = new AgentResumeToken();
        token.setValue(resumeToken);
        AgentResumeRequest request = new AgentResumeRequest();
        request.setResumeToken(token);
        request.setApproved(approved);
        request.setRejectReason(reason);
        request.getMetadata().put("requestId", resolvedRequestId);
        request.getMetadata().put("operatorId", userId);
        context.resume(request);
    }

    private String resolveRequestId(String requestId, String resumeToken) {
        if (requestId != null && !requestId.isBlank()) {
            return requestId;
        }
        String resolved = resumeTokenIndex.get(resumeToken);
        if (resolved == null || resolved.isBlank()) {
            throw new BusinessException("当前 Agent 运行已结束或不在本进程内");
        }
        return resolved;
    }

    private void assertResumeTokenBelongsToRequest(String requestId, String resumeToken) {
        if (resumeToken == null || resumeToken.isBlank()) {
            throw new BusinessException("Agent 恢复令牌不能为空");
        }
        Set<String> tokens = requestTokens.get(requestId);
        if (tokens == null || !tokens.contains(resumeToken)) {
            throw new BusinessException("Agent 审批请求已失效");
        }
    }

    private void assertOwner(String requestId, String userId) {
        if (userId == null || userId.isBlank()) {
            return;
        }
        RunOwner owner = owners.get(requestId);
        if (owner == null) {
            throw new BusinessException("当前 Agent 运行已结束或不在本进程内");
        }
        if (owner.userId() != null && !owner.userId().equals(userId)) {
            throw new BusinessException("无权处理该 Agent 运行审批");
        }
    }

    /**
     * 当前运行归属信息。
     *
     * @param agentId Agent ID
     * @param sessionId 会话 ID
     * @param userId 用户 ID
     */
    public record RunOwner(String agentId, String sessionId, String userId) {
    }

    /**
     * 单机内存运行态。
     *
     */
    public static final class AgentRunContext {

        private final String requestId;
        private final String sessionId;
        private final AgentRuntime runtime;
        private final ChatSseEmitter chatSseEmitter;
        private final ChatRuntimeContext chatContext;
        private final StringBuilder answer;
        private final ChatAssistantAccumulator assistantAccumulator;
        private final AtomicBoolean finished;
        private final boolean persistChatlog;
        private final RunOwner owner;
        private final AgentRunLock.Handle lockHandle;
        private final Consumer<AgentRuntimeEvent> eventConsumer;
        private final Consumer<Throwable> errorConsumer;
        private final Runnable completionHandler;
        private final AtomicBoolean suspended = new AtomicBoolean(false);
        private final AtomicReference<Disposable> subscription = new AtomicReference<>();

        /**
         * 创建运行态。
         *
         * @param requestId 请求 ID
         * @param sessionId 会话 ID
         * @param runtime 有状态运行时
         * @param chatSseEmitter SSE 连接
         * @param chatContext 聊天上下文
         * @param answer 助手正文累计缓冲
         * @param assistantAccumulator 助手结构化累计器
         * @param finished 运行收口标记
         * @param persistChatlog 是否持久化聊天日志
         * @param owner 运行归属
         * @param eventConsumer 运行事件处理器
         * @param errorConsumer 错误处理器
         * @param completionHandler 完成处理器
         */
        public AgentRunContext(String requestId,
                               String sessionId,
                               AgentRuntime runtime,
                               ChatSseEmitter chatSseEmitter,
                               ChatRuntimeContext chatContext,
                               StringBuilder answer,
                               ChatAssistantAccumulator assistantAccumulator,
                               AtomicBoolean finished,
                               boolean persistChatlog,
                               RunOwner owner,
                               AgentRunLock.Handle lockHandle,
                               Consumer<AgentRuntimeEvent> eventConsumer,
                               Consumer<Throwable> errorConsumer,
                               Runnable completionHandler) {
            this.requestId = requestId;
            this.sessionId = sessionId;
            this.runtime = runtime;
            this.chatSseEmitter = chatSseEmitter;
            this.chatContext = chatContext;
            this.answer = answer;
            this.assistantAccumulator = assistantAccumulator;
            this.finished = finished;
            this.persistChatlog = persistChatlog;
            this.owner = owner;
            this.lockHandle = lockHandle;
            this.eventConsumer = eventConsumer;
            this.errorConsumer = errorConsumer;
            this.completionHandler = completionHandler;
        }

        /**
         * 获取请求 ID。
         *
         * @return 请求 ID
         */
        public String requestId() {
            return requestId;
        }

        /**
         * 获取会话 ID。
         *
         * @return 会话 ID
         */
        public String sessionId() {
            return sessionId;
        }

        /**
         * 获取运行归属。
         *
         * @return 运行归属
         */
        public RunOwner owner() {
            return owner;
        }

        /**
         * 获取运行事件处理器。
         *
         * @return 运行事件处理器
         */
        public Consumer<AgentRuntimeEvent> eventConsumer() {
            return eventConsumer;
        }

        /**
         * 获取错误处理器。
         *
         * @return 错误处理器
         */
        public Consumer<Throwable> errorConsumer() {
            return errorConsumer;
        }

        /**
         * 获取完成处理器。
         *
         * @return 完成处理器
         */
        public Runnable completionHandler() {
            return completionHandler;
        }

        /**
         * 标记当前运行已进入 HITL 挂起态。
         */
        public void markSuspended() {
            suspended.set(true);
        }

        /**
         * 当前运行是否处于 HITL 挂起态。
         *
         * @return true 表示等待审批恢复
         */
        public boolean isSuspended() {
            return suspended.get();
        }

        /**
         * 当前运行是否持久化聊天日志与运行态。
         *
         * @return true 表示正式聊天持久化运行
         */
        public boolean persistChatlog() {
            return persistChatlog;
        }

        /**
         * 绑定运行订阅。
         *
         * @param subscription Reactor 订阅
         */
        public void setSubscription(Disposable subscription) {
            if (subscription == null) {
                return;
            }
            Disposable previous = this.subscription.getAndSet(subscription);
            if (previous != null && !previous.isDisposed()) {
                previous.dispose();
            }
        }

        /**
         * 取消当前运行订阅。
         */
        public void cancel() {
            Disposable subscription = this.subscription.getAndSet(null);
            if (subscription != null && !subscription.isDisposed()) {
                subscription.dispose();
            }
        }

        /**
         * 取消当前运行并关闭 SSE。
         *
         * <p>该方法用于调用方主动清理会话的场景。它不通过 runtime 事件链发送取消事件，
         * 因为调用方此时已经明确要求丢弃当前草稿会话。</p>
         */
        public void cancelAndComplete() {
            cancel();
            if (finished.compareAndSet(false, true) && chatSseEmitter != null) {
                chatSseEmitter.complete();
            }
        }

        /**
         * 释放当前运行持有的分布式锁。
         */
        public void releaseLock() {
            if (lockHandle != null) {
                lockHandle.release();
            }
        }

        /**
         * 关闭底层运行时并释放资源。
         */
        public void closeRuntime() {
            try {
                runtime.close();
            } catch (Exception e) {
                LOG.warn("Close Agent runtime failed, requestId={}, sessionId={}, message={}",
                        requestId, sessionId, e.getMessage(), e);
            }
        }

        /**
         * 通过同一个 runtime 恢复挂起运行，事件继续写入原 SSE。
         *
         * @param request 恢复请求
         */
        public void resume(AgentResumeRequest request) {
            suspended.set(false);
            Disposable subscription = runtime.resume(request).subscribe(eventConsumer, errorConsumer, completionHandler);
            setSubscription(subscription);
        }
    }
}
