package com.easyagents.agent.runtime.hitl;

import com.easyagents.agent.runtime.AgentResumeRequest;
import com.easyagents.agent.runtime.AgentRuntimeException;
import reactor.core.publisher.Mono;

import java.time.Instant;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.CompletableFuture;

/**
 * 工具审批协调器。
 */
public class AgentToolApprovalCoordinator {

    /** MCP 工具类型。 */
    private static final String MCP_TOOL_TYPE = "MCP";

    /** 是否启用内存审批协调。 */
    private final boolean enabled;
    /** 恢复令牌到待审批项的索引。 */
    private final Map<String, PendingApproval> approvals = new LinkedHashMap<>();
    /** 审批批次ID到批次状态的索引。 */
    private final Map<String, ApprovalBatch> approvalBatches = new LinkedHashMap<>();
    /** 工具调用ID到恢复令牌的唯一索引。 */
    private final Map<String, String> tokensByToolCallId = new LinkedHashMap<>();
    /** 工具调用ID到一次性执行授权的索引。 */
    private final Map<String, ExecutionAuthorization> executionAuthorizations = new LinkedHashMap<>();
    /** 当前 Turn 已批准的可复用工具作用域。 */
    private final Set<String> reusableApprovalScopes = new LinkedHashSet<>();

    /**
     * 创建已启用的协调器。
     *
     * @return 已启用的协调器
     */
    public static AgentToolApprovalCoordinator enabled() {
        return new AgentToolApprovalCoordinator(true);
    }

    /**
     * 创建禁用的协调器。
     *
     * @return 禁用的协调器
     */
    public static AgentToolApprovalCoordinator disabled() {
        return new AgentToolApprovalCoordinator(false);
    }

    /**
     * 创建协调器。
     *
     * @param enabled 是否启用
     */
    public AgentToolApprovalCoordinator(boolean enabled) {
        this.enabled = enabled;
    }

    /**
     * 注册一个待审批请求。
     *
     * @param sessionId 会话ID
     * @param agentId 智能体ID
     * @param toolCallId 工具调用ID
     * @param toolName 工具名称
     * @param approvalPrompt 审批文案
     * @param toolInput 工具入参
     * @param metadata 元数据
     * @param expiresAt 过期时间
     * @return 待审批状态
     */
    public AgentPendingState register(String sessionId,
                                      String agentId,
                                      String toolCallId,
                                      String toolName,
                                      String approvalPrompt,
                                      Map<String, Object> toolInput,
                                      Map<String, Object> metadata,
                                      Instant expiresAt) {
        return register(sessionId, agentId, toolCallId, toolName, approvalPrompt, toolInput,
                metadata, expiresAt, null);
    }

    /**
     * 注册一个属于指定审批批次的待审批请求。
     *
     * <p>同一批次内的全部工具调用均批准后，协调器才会签发执行授权。任意一项拒绝或
     * 过期都会关闭整个批次，避免未批准调用跟随已批准调用一起恢复执行。</p>
     *
     * @param sessionId 会话ID
     * @param agentId 智能体ID
     * @param toolCallId 工具调用ID
     * @param toolName 工具名称
     * @param approvalPrompt 审批文案
     * @param toolInput 工具入参
     * @param metadata 元数据
     * @param expiresAt 过期时间
     * @param approvalBatchId 审批批次ID；为空时创建单调用批次
     * @return 待审批状态
     */
    public synchronized AgentPendingState register(String sessionId,
                                                   String agentId,
                                                   String toolCallId,
                                                   String toolName,
                                                   String approvalPrompt,
                                                   Map<String, Object> toolInput,
                                                   Map<String, Object> metadata,
                                                   Instant expiresAt,
                                                   String approvalBatchId) {
        if (enabled && (toolCallId == null || toolCallId.isBlank())) {
            throw new AgentRuntimeException("Approval-required tool call must include toolCallId.");
        }
        if (enabled && toolCallId != null && !toolCallId.isBlank()) {
            String existingToken = tokensByToolCallId.get(toolCallId);
            PendingApproval existing = existingToken == null ? null : approvals.get(existingToken);
            if (existing != null && !isExpired(existing.state)) {
                if (!Objects.equals(existing.state.getToolName(), toolName)
                        || !Objects.equals(existing.state.getToolInput(), toolInput)) {
                    throw new AgentRuntimeException(
                            "Duplicate toolCallId is bound to a different tool call: " + toolCallId);
                }
                return existing.state;
            }
            if (existing != null) {
                closeBatch(existing.batchId, "审批请求已过期。");
            }
        }
        AgentPendingState state = new AgentPendingState();
        state.setSessionId(sessionId);
        state.setAgentId(agentId);
        state.setToolCallId(toolCallId);
        state.setToolName(toolName);
        state.setApprovalPrompt(approvalPrompt);
        state.setToolInput(toolInput);
        state.setExpiresAt(expiresAt);
        String token = state.getResumeToken().getValue();
        String effectiveBatchId = approvalBatchId == null || approvalBatchId.isBlank()
                ? token
                : approvalBatchId;
        Map<String, Object> effectiveMetadata = metadata == null
                ? new LinkedHashMap<>()
                : new LinkedHashMap<>(metadata);
        effectiveMetadata.put("approvalBatchId", effectiveBatchId);
        state.setMetadata(effectiveMetadata);
        if (enabled) {
            PendingApproval pendingApproval = new PendingApproval(
                    state, effectiveBatchId, new CompletableFuture<>());
            approvals.put(token, pendingApproval);
            ApprovalBatch batch = approvalBatches.computeIfAbsent(effectiveBatchId, ApprovalBatch::new);
            batch.tokens.add(token);
            batch.members.put(token, pendingApproval);
            if (toolCallId != null && !toolCallId.isBlank()) {
                tokensByToolCallId.put(toolCallId, token);
            }
        }
        return state;
    }

    /**
     * 等待指定审批令牌的响应。
     *
     * @param resumeToken 恢复令牌
     * @return 恢复请求
     */
    public synchronized Mono<AgentResumeRequest> await(AgentResumeToken resumeToken) {
        if (!enabled) {
            AgentResumeRequest request = new AgentResumeRequest();
            request.setResumeToken(resumeToken);
            request.setApproved(true);
            return Mono.just(request);
        }
        if (resumeToken == null || resumeToken.getValue() == null || resumeToken.getValue().isBlank()) {
            return Mono.error(new AgentToolApprovalRejectedException("缺少审批令牌。"));
        }
        PendingApproval pendingApproval = approvals.get(resumeToken.getValue());
        if (pendingApproval == null) {
            return Mono.error(new AgentToolApprovalRejectedException("审批请求已失效。"));
        }
        if (isExpired(pendingApproval.state)) {
            closeBatch(pendingApproval.batchId, "审批请求已过期。");
            return Mono.error(new AgentToolApprovalRejectedException("审批请求已过期。"));
        }
        return Mono.fromFuture(pendingApproval.future);
    }

    /**
     * 处理恢复请求并生成批次级审批决议。
     *
     * @param request 恢复请求
     * @return 审批决议
     */
    public synchronized AgentToolApprovalResolution resolve(AgentResumeRequest request) {
        validateResumeRequest(request);
        if (!enabled) {
            AgentPendingState state = new AgentPendingState();
            state.setResumeToken(request.getResumeToken());
            return new AgentToolApprovalResolution(
                    AgentToolApprovalResolution.Status.READY, state, List.of(), null);
        }
        String token = request.getResumeToken().getValue();
        PendingApproval pendingApproval = approvals.get(token);
        if (pendingApproval == null || pendingApproval.decision != ApprovalDecision.PENDING) {
            throw new AgentRuntimeException("Agent resume token is invalid, expired, or already consumed.");
        }
        if (isExpired(pendingApproval.state)) {
            closeBatch(pendingApproval.batchId, "审批请求已过期。");
            return new AgentToolApprovalResolution(
                    AgentToolApprovalResolution.Status.EXPIRED,
                    pendingApproval.state,
                    List.of(),
                    "审批请求已过期。");
        }

        pendingApproval.decision = request.isApproved()
                ? ApprovalDecision.APPROVED
                : ApprovalDecision.REJECTED;
        approvals.remove(token);
        removeToolCallIndex(pendingApproval.state, token);
        pendingApproval.future.complete(request);

        if (!request.isApproved()) {
            String reason = request.getRejectReason() == null || request.getRejectReason().isBlank()
                    ? "工具执行已被拒绝。"
                    : request.getRejectReason();
            closeBatch(pendingApproval.batchId, reason);
            return new AgentToolApprovalResolution(
                    AgentToolApprovalResolution.Status.REJECTED,
                    pendingApproval.state,
                    List.of(),
                    reason);
        }

        ApprovalBatch batch = approvalBatches.get(pendingApproval.batchId);
        if (batch == null) {
            authorize(pendingApproval);
            return new AgentToolApprovalResolution(
                    AgentToolApprovalResolution.Status.READY,
                    pendingApproval.state,
                    List.of(),
                    null);
        }
        List<AgentPendingState> remainingStates = pendingStatesInBatch(batch);
        AgentPendingState expiredState = remainingStates.stream()
                .filter(this::isExpired)
                .findFirst()
                .orElse(null);
        if (expiredState != null) {
            closeBatch(batch.batchId, "审批请求已过期。");
            return new AgentToolApprovalResolution(
                    AgentToolApprovalResolution.Status.EXPIRED,
                    pendingApproval.state,
                    List.of(),
                    "审批请求已过期。");
        }
        if (!remainingStates.isEmpty()) {
            return new AgentToolApprovalResolution(
                    AgentToolApprovalResolution.Status.WAITING,
                    pendingApproval.state,
                    remainingStates,
                    null);
        }

        for (String batchToken : batch.tokens) {
            PendingApproval member = batch.members.get(batchToken);
            if (member != null && member.decision == ApprovalDecision.APPROVED) {
                authorize(member);
            }
        }
        approvalBatches.remove(batch.batchId);
        return new AgentToolApprovalResolution(
                AgentToolApprovalResolution.Status.READY,
                pendingApproval.state,
                List.of(),
                null);
    }

    /**
     * 消费恢复请求对应的待审批状态。
     *
     * <p>该兼容入口供工具适配器内部的 await/consume 流程使用。批次中仍有未决调用时
     * 会拒绝提前消费；单调用批准后会移除无需经过 PreActing 的执行凭证。</p>
     *
     * @param request 恢复请求
     * @return 待审批状态
     */
    public synchronized AgentPendingState consume(AgentResumeRequest request) {
        AgentToolApprovalResolution resolution = resolve(request);
        if (resolution.getStatus() == AgentToolApprovalResolution.Status.WAITING) {
            throw new AgentRuntimeException("Approval batch still has pending tool calls.");
        }
        AgentPendingState state = resolution.getResolvedState();
        if (resolution.getStatus() == AgentToolApprovalResolution.Status.READY
                && state != null
                && state.getToolCallId() != null) {
            // 兼容工具适配器内部 await/consume 流程，该流程会直接调用工具，不经过 PreActing 二次校验。
            executionAuthorizations.remove(state.getToolCallId());
        }
        return state;
    }

    /**
     * 根据服务端持久化审批结果签发受信任的执行授权。
     *
     * <p>恢复元数据应提供单个 {@code toolCallId/toolName/toolInput}，或通过
     * {@code approvedToolCalls} 提供上述字段组成的列表。MCP 调用可额外携带受信任的
     * {@code toolType/mcpId}，用于签发当前 Turn 的复用作用域。该入口仅供已经完成
     * 持久化令牌校验和一次性消费的服务端集成层使用。</p>
     *
     * @param request 受信任恢复请求
     */
    public synchronized void authorizeTrustedExecution(AgentResumeRequest request) {
        validateResumeRequest(request);
        if (!enabled || !request.isApproved()) {
            return;
        }
        Map<String, Object> metadata = request.getMetadata() == null
                ? Map.of()
                : request.getMetadata();
        Object approvedToolCalls = metadata.get("approvedToolCalls");
        Map<String, ExecutionAuthorization> trustedAuthorizations = new LinkedHashMap<>();
        Set<String> trustedReusableScopes = new LinkedHashSet<>();
        int authorizationCount = 0;
        if (approvedToolCalls instanceof List<?> calls) {
            for (Object call : calls) {
                if (call instanceof Map<?, ?> callMap) {
                    authorizeTrustedCall(callMap, trustedAuthorizations, trustedReusableScopes);
                    authorizationCount++;
                }
            }
        } else if (metadata.containsKey("toolCallId")) {
            authorizeTrustedCall(metadata, trustedAuthorizations, trustedReusableScopes);
            authorizationCount++;
        }
        if (authorizationCount == 0) {
            throw new AgentRuntimeException(
                    "Trusted resume metadata must include approved toolCallId, toolName, and toolInput.");
        }
        executionAuthorizations.putAll(trustedAuthorizations);
        reusableApprovalScopes.addAll(trustedReusableScopes);
    }

    /**
     * 消费指定工具调用的一次性执行授权。
     *
     * @param toolCallId 工具调用ID
     * @param toolName 工具名称
     * @param toolInput 工具入参
     */
    public synchronized void consumeExecutionAuthorization(String toolCallId,
                                                           String toolName,
                                                           Map<String, Object> toolInput) {
        if (!enabled) {
            return;
        }
        if (toolCallId == null || toolCallId.isBlank()) {
            throw new AgentToolApprovalRejectedException("待执行工具缺少 toolCallId，无法校验审批结果。");
        }
        ExecutionAuthorization authorization = executionAuthorizations.remove(toolCallId);
        if (authorization == null) {
            throw new AgentToolApprovalRejectedException("工具调用未获得批准或批准已被消费。");
        }
        if (!Objects.equals(authorization.toolName, toolName)
                || !Objects.equals(authorization.toolInput, toolInput)) {
            throw new AgentToolApprovalRejectedException("工具调用与已批准内容不一致。");
        }
    }

    /**
     * 判断工具元数据对应的复用作用域是否已在当前 Turn 获得批准。
     *
     * @param metadata 服务端工具元数据
     * @return 当前 Turn 已批准时为 true
     */
    public synchronized boolean isReusableApprovalGranted(Map<String, Object> metadata) {
        String approvalScope = reusableApprovalScope(metadata);
        return approvalScope != null && reusableApprovalScopes.contains(approvalScope);
    }

    /**
     * 解析可在当前 Turn 复用的审批作用域。
     *
     * <p>MCP 使用稳定 {@code mcpId} 生成作用域；受控 Shell 脚本仅接受动态审批策略写入的
     * 内容摘要作用域。缺少受信任标识时返回 null，使调用方继续执行逐调用审批。</p>
     *
     * @param metadata 服务端工具元数据
     * @return 可复用审批作用域；不可复用时返回 null
     */
    public String reusableApprovalScope(Map<String, Object> metadata) {
        if (metadata == null || metadata.isEmpty()) {
            return null;
        }
        String explicitScope = stringValue(metadata.get("approvalScope"));
        if (Boolean.TRUE.equals(metadata.get("operateTool"))
                && "SHELL".equalsIgnoreCase(stringValue(metadata.get("operateToolType")))
                && explicitScope != null
                && explicitScope.startsWith("SHELL_SCRIPT:")) {
            return explicitScope;
        }
        String toolType = stringValue(metadata.get("toolType"));
        String mcpId = stringValue(metadata.get("mcpId"));
        if (!MCP_TOOL_TYPE.equalsIgnoreCase(toolType) || mcpId == null) {
            return null;
        }
        return MCP_TOOL_TYPE + ":" + mcpId;
    }

    /**
     * 清理尚未消费的一次性执行授权。
     */
    public synchronized void clearExecutionAuthorizations() {
        executionAuthorizations.clear();
    }

    /**
     * 清理当前 Turn 的可复用工具审批作用域。
     */
    public synchronized void clearReusableApprovalScopes() {
        reusableApprovalScopes.clear();
    }

    /**
     * 获取指定会话当前仍待处理的审批状态。
     *
     * @param sessionId 会话ID
     * @return 待审批状态快照
     */
    public synchronized List<AgentPendingState> pendingStates(String sessionId) {
        return approvals.values().stream()
                .filter(pending -> pending.decision == ApprovalDecision.PENDING)
                .map(pending -> pending.state)
                .filter(state -> sessionId == null || Objects.equals(sessionId, state.getSessionId()))
                .toList();
    }

    /**
     * 取消所有挂起审批。
     *
     * @param reason 取消原因
     */
    public synchronized void cancelAll(String reason) {
        if (!enabled) {
            return;
        }
        for (String batchId : new ArrayList<>(approvalBatches.keySet())) {
            closeBatch(batchId, reason);
        }
        approvals.clear();
        tokensByToolCallId.clear();
        executionAuthorizations.clear();
        reusableApprovalScopes.clear();
    }

    /**
     * 是否已启用。
     *
     * @return 启用标记
     */
    public boolean isEnabled() {
        return enabled;
    }

    /**
     * 校验恢复请求中的令牌字段。
     *
     * @param request 恢复请求
     */
    private void validateResumeRequest(AgentResumeRequest request) {
        if (request == null || request.getResumeToken() == null
                || request.getResumeToken().getValue() == null
                || request.getResumeToken().getValue().isBlank()) {
            throw new AgentRuntimeException("Agent resume token is required.");
        }
    }

    /**
     * 判断待审批状态是否已经过期。
     *
     * @param state 待审批状态
     * @return 已过期时为 true
     */
    private boolean isExpired(AgentPendingState state) {
        return state != null
                && state.getExpiresAt() != null
                && !state.getExpiresAt().isAfter(Instant.now());
    }

    /**
     * 获取批次内尚未决策的审批状态。
     *
     * @param batch 审批批次
     * @return 未决审批状态
     */
    private List<AgentPendingState> pendingStatesInBatch(ApprovalBatch batch) {
        List<AgentPendingState> states = new ArrayList<>();
        for (String token : batch.tokens) {
            PendingApproval member = batch.members.get(token);
            if (member != null && member.decision == ApprovalDecision.PENDING) {
                states.add(member.state);
            }
        }
        return states;
    }

    /**
     * 为已批准状态签发一次性执行授权。
     *
     * @param pendingApproval 已批准状态
     */
    private void authorize(PendingApproval pendingApproval) {
        AgentPendingState state = pendingApproval.state;
        String approvalScope = reusableApprovalScope(state.getMetadata());
        if (approvalScope != null) {
            reusableApprovalScopes.add(approvalScope);
            return;
        }
        if (state.getToolCallId() == null || state.getToolCallId().isBlank()) {
            throw new AgentRuntimeException("Approved tool call is missing toolCallId.");
        }
        executionAuthorizations.put(state.getToolCallId(), new ExecutionAuthorization(
                state.getToolName(),
                state.getToolInput()));
    }

    /**
     * 为服务端持久化审批结果签发一次性执行授权。
     *
     * @param callMap 已批准调用元数据
     * @param trustedAuthorizations 本次恢复待签发的一次性授权集合
     * @param trustedReusableScopes 本次恢复待签发的可复用作用域集合
     */
    private void authorizeTrustedCall(Map<?, ?> callMap,
                                      Map<String, ExecutionAuthorization> trustedAuthorizations,
                                      Set<String> trustedReusableScopes) {
        String toolCallId = stringValue(callMap.get("toolCallId"));
        String toolName = stringValue(callMap.get("toolName"));
        if (toolCallId == null || toolName == null) {
            throw new AgentRuntimeException(
                    "Trusted resume metadata must include non-empty toolCallId and toolName.");
        }
        Map<String, Object> toolInput = stringKeyMap(callMap.get("toolInput"));
        Map<String, Object> scopeMetadata = new LinkedHashMap<>();
        scopeMetadata.put("toolType", callMap.get("toolType"));
        scopeMetadata.put("mcpId", callMap.get("mcpId"));
        scopeMetadata.put("operateTool", callMap.get("operateTool"));
        scopeMetadata.put("operateToolType", callMap.get("operateToolType"));
        scopeMetadata.put("approvalScope", callMap.get("approvalScope"));
        String reusableScope = reusableApprovalScope(scopeMetadata);
        if (reusableScope != null) {
            trustedReusableScopes.add(reusableScope);
            return;
        }
        ExecutionAuthorization authorization = new ExecutionAuthorization(toolName, toolInput);
        ExecutionAuthorization previous = trustedAuthorizations.put(toolCallId, authorization);
        if (previous != null
                && (!Objects.equals(previous.toolName, toolName)
                || !Objects.equals(previous.toolInput, toolInput))) {
            throw new AgentRuntimeException(
                    "Trusted resume contains conflicting tool calls for toolCallId: " + toolCallId);
        }
    }

    /**
     * 将值转换为非空字符串。
     *
     * @param value 原始值
     * @return 非空字符串；无法转换时返回 null
     */
    private String stringValue(Object value) {
        if (value == null || String.valueOf(value).isBlank()) {
            return null;
        }
        return String.valueOf(value);
    }

    /**
     * 将任意键 Map 转换为字符串键 Map。
     *
     * @param value 原始值
     * @return 字符串键 Map
     */
    private Map<String, Object> stringKeyMap(Object value) {
        if (value == null) {
            return Map.of();
        }
        if (!(value instanceof Map<?, ?> source)) {
            throw new AgentRuntimeException("Trusted resume toolInput must be a map.");
        }
        Map<String, Object> result = new LinkedHashMap<>();
        for (Map.Entry<?, ?> entry : source.entrySet()) {
            if (entry.getKey() != null) {
                result.put(String.valueOf(entry.getKey()), entry.getValue());
            }
        }
        return result;
    }

    /**
     * 移除工具调用ID到审批令牌的索引。
     *
     * @param state 待审批状态
     * @param token 审批令牌
     */
    private void removeToolCallIndex(AgentPendingState state, String token) {
        if (state.getToolCallId() != null && !state.getToolCallId().isBlank()) {
            tokensByToolCallId.remove(state.getToolCallId(), token);
        }
    }

    /**
     * 关闭审批批次并拒绝尚未决策的审批项。
     *
     * @param batchId 审批批次ID
     * @param reason 关闭原因
     */
    private void closeBatch(String batchId, String reason) {
        ApprovalBatch batch = approvalBatches.remove(batchId);
        if (batch == null) {
            return;
        }
        for (String token : batch.tokens) {
            PendingApproval member = batch.members.get(token);
            if (member == null) {
                continue;
            }
            approvals.remove(token);
            removeToolCallIndex(member.state, token);
            if (member.decision == ApprovalDecision.PENDING) {
                member.decision = ApprovalDecision.REJECTED;
                AgentResumeRequest rejection = new AgentResumeRequest();
                rejection.setResumeToken(member.state.getResumeToken());
                rejection.setApproved(false);
                rejection.setRejectReason(reason);
                member.future.complete(rejection);
            }
        }
    }

    /**
     * 单个审批项的内部决策状态。
     */
    private enum ApprovalDecision {
        PENDING,
        APPROVED,
        REJECTED
    }

    /**
     * 待审批项及其异步等待句柄。
     */
    private static class PendingApproval {
        private final AgentPendingState state;
        private final String batchId;
        private final CompletableFuture<AgentResumeRequest> future;
        private ApprovalDecision decision = ApprovalDecision.PENDING;

        /**
         * 创建待审批项。
         *
         * @param state 待审批状态
         * @param batchId 审批批次ID
         * @param future 审批响应等待句柄
         */
        private PendingApproval(AgentPendingState state,
                                String batchId,
                                CompletableFuture<AgentResumeRequest> future) {
            this.state = Objects.requireNonNull(state, "state");
            this.batchId = Objects.requireNonNull(batchId, "batchId");
            this.future = Objects.requireNonNull(future, "future");
        }
    }

    /**
     * 同一轮推理产生的审批批次。
     */
    private static class ApprovalBatch {
        private final String batchId;
        private final Set<String> tokens = new LinkedHashSet<>();
        private final Map<String, PendingApproval> members = new LinkedHashMap<>();

        /**
         * 创建审批批次。
         *
         * @param batchId 审批批次ID
         */
        private ApprovalBatch(String batchId) {
            this.batchId = Objects.requireNonNull(batchId, "batchId");
        }
    }

    /**
     * 已批准工具调用的一次性执行凭证。
     */
    private static class ExecutionAuthorization {
        private final String toolName;
        private final Map<String, Object> toolInput;

        /**
         * 创建一次性执行授权。
         *
         * @param toolName 工具名称
         * @param toolInput 工具入参
         */
        private ExecutionAuthorization(String toolName, Map<String, Object> toolInput) {
            this.toolName = toolName;
            this.toolInput = toolInput == null ? Map.of() : new LinkedHashMap<>(toolInput);
        }
    }
}
