package com.aicodeassistant.controller;

import com.aicodeassistant.exception.PermissionModeMismatchException;

import com.aicodeassistant.engine.CompactService;
import com.aicodeassistant.engine.QueryConfig;
import com.aicodeassistant.engine.QueryEngine;
import com.aicodeassistant.engine.QueryLoopState;
import com.aicodeassistant.engine.QueryMessageHandler;
import com.aicodeassistant.engine.TokenCounter;
import com.aicodeassistant.exception.RequestValidationException;
import com.aicodeassistant.exception.SessionNotFoundException;
import com.aicodeassistant.llm.LlmProviderRegistry;
import com.aicodeassistant.llm.ModelCapabilities;
import com.aicodeassistant.llm.ModelRegistry;
import com.aicodeassistant.llm.ThinkingConfig;
import com.aicodeassistant.model.ContentBlock;
import com.aicodeassistant.model.Message;
import com.aicodeassistant.model.PermissionMode;
import com.aicodeassistant.permission.PermissionModeManager;
import com.aicodeassistant.model.Usage;
import com.aicodeassistant.prompt.EffectiveSystemPromptBuilder;
import com.aicodeassistant.prompt.SystemPromptConfig;
import com.aicodeassistant.service.ProjectWorkspaceService;
import com.aicodeassistant.session.SessionData;
import com.aicodeassistant.session.SessionManager;
import com.aicodeassistant.session.SessionMessagePersistence;
import com.aicodeassistant.session.SessionExecutionGate;
import com.aicodeassistant.session.SessionExecutionBusyException;
import com.aicodeassistant.tool.Tool;
import com.aicodeassistant.tool.ToolRegistry;
import com.aicodeassistant.tool.ToolUseContext;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.time.Instant;
import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * QueryController — 查询 API (CLI/SDK 专用)。
 * <p>
 * 端点:
 * <ul>
 *   <li>POST /api/query — 同步单次查询</li>
 *   <li>POST /api/query/stream — SSE 流式查询</li>
 *   <li>POST /api/query/conversation — 多轮会话查询</li>
 * </ul>
 * <p>
 * 与 WebSocket API (§6.2) 的分工:
 * - WebSocket STOMP: Web 前端长连接，双向实时交互
 * - REST /api/query: CLI/脚本单次调用，请求-响应模式
 * - SSE /api/query/stream: CLI 流式输出，实时展示
 *
 */
@RestController
@RequestMapping("/api/query")
public class QueryController {

    private static final Logger log = LoggerFactory.getLogger(QueryController.class);

    private final QueryEngine queryEngine;
    private final ToolRegistry toolRegistry;
    @org.springframework.beans.factory.annotation.Autowired
    private com.aicodeassistant.engine.HandoffContextService handoffContext;
    private final SessionManager sessionManager;
    private final LlmProviderRegistry providerRegistry;
    private final TokenCounter tokenCounter;
    private final ObjectMapper objectMapper;
    private final EffectiveSystemPromptBuilder systemPromptBuilder;
    private final PermissionModeManager permissionModeManager;
    private final ModelRegistry modelRegistry;
    private final ProjectWorkspaceService projectWorkspaces;
    private final SessionExecutionGate executionGate;

    public QueryController(QueryEngine queryEngine,
                           ToolRegistry toolRegistry,
                           SessionManager sessionManager,
                           LlmProviderRegistry providerRegistry,
                           TokenCounter tokenCounter,
                           ObjectMapper objectMapper,
                           EffectiveSystemPromptBuilder systemPromptBuilder,
                           PermissionModeManager permissionModeManager,
                           ModelRegistry modelRegistry,
                           ProjectWorkspaceService projectWorkspaces,
                           SessionExecutionGate executionGate) {
        this.queryEngine = queryEngine;
        this.toolRegistry = toolRegistry;
        this.sessionManager = sessionManager;
        this.providerRegistry = providerRegistry;
        this.tokenCounter = tokenCounter;
        this.objectMapper = objectMapper;
        this.systemPromptBuilder = systemPromptBuilder;
        this.permissionModeManager = permissionModeManager;
        this.modelRegistry = modelRegistry;
        this.projectWorkspaces = projectWorkspaces;
        this.executionGate = executionGate;
    }

    // ════════════════════════════════════════════
    // 端点 1: 同步查询 — 等待完成后返回完整结果
    // ════════════════════════════════════════════

    /**
     * POST /api/query — 同步单次查询。
     * <p>
     * 阻塞直到 LLM 完成所有轮次，返回完整结果。
     * 在 Virtual Thread 中执行，不阻塞平台线程。
     * 权限策略: 新会话默认 AUTO_APPROVE；已有会话沿用保存的权限。
     */
    @PostMapping
    public ResponseEntity<QueryResponse> query(@RequestBody QueryRequest request) {
        rejectUnsupportedMaxBudgetUsd(request.maxBudgetUsd());
        // 1. 创建或复用会话
        try (LockedSession locked = resolveSessionLocked(request)) {
        SessionData session = locked.session();
        String sessionId = session.sessionId();

        // 请求参数仅校验已有会话权限，不隐式修改会话选择。
        PermissionMode effectiveMode = requireMatchingMode(sessionId, request.permissionMode());
        log.debug("REST API /api/query: permissionMode={} (requested={})",
                effectiveMode, request.permissionMode());

        // 2. 组装工具池
        List<Tool> tools = assembleToolPool(request.allowedTools(), request.disallowedTools());

        // 3. 构建系统提示（使用新的 SystemPromptBuilder）
        SystemPromptConfig promptConfig = SystemPromptConfig.defaults()
                .withCustom(request.systemPrompt())
                .withAppend(request.appendSystemPrompt())
                .withSessionId(sessionId);
        String systemPrompt = systemPromptBuilder.buildEffectiveSystemPrompt(
                promptConfig, tools, request.model(),
                Path.of(session.workingDir()));

        // 4. 组装用户消息
        String userMessage = buildUserMessage(request.prompt(), request.context());

        // 5. 构建 QueryConfig（★ 别名解析统一指向 qwen3.8-max-0902）
        String rawModel = request.model() != null ? request.model() : providerRegistry.getDefaultModel();
        String model = providerRegistry.resolveModelAlias(rawModel);
        int maxTurns = request.maxTurns() != null ? request.maxTurns() : QueryConfig.DEFAULT_MAX_TURNS;
        int contextWindow = getContextWindow(model);

        QueryConfig config = QueryConfig.withDefaults(
                model, systemPrompt, tools,
                tools.stream().map(Tool::toToolDefinition).toList(),
                QueryConfig.getRecommendedMaxTokens(modelRegistry, model),
                contextWindow,
                new ThinkingConfig.Adaptive(),
                maxTurns, "rest-api"
        );

        // 6. 初始化循环状态 — 如果是已有会话，加载历史消息
        List<Message> historyMessages =
                new ArrayList<>(session.messages());
        ToolUseContext toolCtx = ToolUseContext.of(
                session.workingDir(), sessionId);
        QueryLoopState state = new QueryLoopState(historyMessages, toolCtx);
        SessionMessagePersistence.attach(
                state, sessionManager, sessionId, "REST /api/query");
        // 添加用户消息
        Message.UserMessage requestMessage = new Message.UserMessage(
                UUID.randomUUID().toString(), Instant.now(),
                List.of(new ContentBlock.TextBlock(userMessage)),
                null, null);
        state.addMessage(requestMessage);

        // 7. 执行查询
        ResultCollectingHandler handler = new ResultCollectingHandler();
        if (com.aicodeassistant.engine.HandoffContextService.isMerged(session.config()))
            config = handoffContext.configure(config, state, session.config(), request.allowedTools(), request.disallowedTools());
        QueryEngine.QueryResult result = queryEngine.execute(config, state, handler);

        // 8. 提取最终文本
        String finalText = extractCurrentRunText(result.messages(), state, result.isSuccess());

        return ResponseEntity.ok(new QueryResponse(
                sessionId,
                finalText,
                result.totalUsage(),
                0.0, // costUsd — 后续 CostTracker 实现
                handler.getToolCalls(),
                result.stopReason(),
                result.isSuccess() ? null : result.error()
        ));
        }
    }

    // ════════════════════════════════════════════
    // 端点 2: SSE 流式查询 — 实时事件流
    // ════════════════════════════════════════════

    /**
     * POST /api/query/stream — SSE 流式查询。
     * <p>
     * 通过 Server-Sent Events 实时推送:
     * stream_delta / thinking_delta / tool_use / tool_result / message_complete / error
     */
    @PostMapping("/stream")
    public SseEmitter streamQuery(@RequestBody QueryRequest request) {
        rejectUnsupportedMaxBudgetUsd(request.maxBudgetUsd());
        long timeoutMs = (request.timeoutSeconds() != null
                ? request.timeoutSeconds() : 600) * 1000L;
        SseEmitter emitter = new SseEmitter(timeoutMs);
        // Validate admission before returning an SSE response so conflicts remain HTTP 409.
        LockedSession locked = resolveSessionLocked(request);

        // The session token supports transfer to the worker thread.
        try {
        Thread.ofVirtual().name("zhiku-query-stream").start(() -> {
            try {
                SessionData session = locked.session();
                String sessionId = session.sessionId();
                try (locked) {
                // Execution uses the saved session permission mode.
                // Admission was checked before returning the SSE response.
                List<Tool> tools = assembleToolPool(
                        request.allowedTools(), request.disallowedTools());
                SystemPromptConfig promptConfig = SystemPromptConfig.defaults()
                        .withCustom(request.systemPrompt())
                        .withAppend(request.appendSystemPrompt())
                        .withSessionId(sessionId);
                String systemPrompt = systemPromptBuilder.buildEffectiveSystemPrompt(
                        promptConfig, tools, request.model(),
                        Path.of(session.workingDir()));
                String userMessage = buildUserMessage(request.prompt(), request.context());
                String rawModel = request.model() != null
                        ? request.model() : providerRegistry.getDefaultModel();
                String model = providerRegistry.resolveModelAlias(rawModel);
                int maxTurns = request.maxTurns() != null ? request.maxTurns() : QueryConfig.DEFAULT_MAX_TURNS;
                int contextWindow = getContextWindow(model);

                QueryConfig config = QueryConfig.withDefaults(
                        model, systemPrompt, tools,
                        tools.stream().map(Tool::toToolDefinition).toList(),
                        QueryConfig.getRecommendedMaxTokens(modelRegistry, model),
                        contextWindow,
                        new ThinkingConfig.Adaptive(),
                        maxTurns, "rest-api-stream"
                );

                // 如果是已有会话，加载历史消息
                List<Message> historyMessages =
                        new ArrayList<>(session.messages());
                ToolUseContext toolCtx = ToolUseContext.of(
                        session.workingDir(), sessionId);
                QueryLoopState state = new QueryLoopState(historyMessages, toolCtx);
                SessionMessagePersistence.attach(
                        state, sessionManager, sessionId, "SSE /api/query/stream");
                Message.UserMessage requestMessage = new Message.UserMessage(
                        UUID.randomUUID().toString(), Instant.now(),
                        List.of(new ContentBlock.TextBlock(userMessage)),
                        null, null);
                state.addMessage(requestMessage);

                // 流式回调 → SSE 推送（使用独立 SseStreamHandler）
                SseStreamHandler handler = new SseStreamHandler(emitter, objectMapper);
                if (com.aicodeassistant.engine.HandoffContextService.isMerged(session.config()))
                    config = handoffContext.configure(config, state, session.config(), request.allowedTools(), request.disallowedTools());
                QueryEngine.QueryResult result = queryEngine.execute(config, state, handler);

                // 完成事件
                sendEvent(emitter, "message_complete", Map.of(
                        "sessionId", sessionId,
                        "usage", result.totalUsage(),
                        "stopReason", result.stopReason()));
                emitter.complete();
                }

            } catch (Exception e) {
                log.error("Stream query error", e);
                sendEvent(emitter, "error", Map.of(
                        "type", "error", "message", e.getMessage()));
                emitter.completeWithError(e);
            }
        });
        } catch (RuntimeException | Error failure) {
            locked.close();
            throw failure;
        }

        return emitter;
    }

    // ════════════════════════════════════════════
    // 端点 3: 多轮会话查询 — 关联已有会话
    // ════════════════════════════════════════════

    /**
     * POST /api/query/conversation — 在已有会话中追加查询。
     * <p>
     * 与 /api/query 的区别:
     * - 必须提供 sessionId
     * - 自动加载会话历史消息
     */
    @PostMapping("/conversation")
    public ResponseEntity<QueryResponse> conversationQuery(
            @RequestBody ConversationRequest request) {

        rejectUnsupportedMaxBudgetUsd(request.maxBudgetUsd());
        rejectClientWorkingDirectory(request.workingDirectory());
        try (SessionExecutionGate.Token ignored = acquireSession(request.sessionId())) {

        // 1. 加载会话
        SessionData session = sessionManager.loadSession(request.sessionId())
                .orElseThrow(() -> new SessionNotFoundException(request.sessionId()));
        projectWorkspaces.requireCurrentBinding(session.workingDir());

        // Execution uses the saved session permission mode.
        PermissionMode effectiveMode = requireMatchingMode(request.sessionId(), request.permissionMode());

        // 2. 准备工具和系统提示
        List<Tool> tools = assembleToolPool(
                request.allowedTools(), request.disallowedTools());
        SystemPromptConfig promptConfig = SystemPromptConfig.defaults()
                .withCustom(request.systemPrompt())
                .withAppend(request.appendSystemPrompt())
                .withSessionId(request.sessionId());
        String systemPrompt = systemPromptBuilder.buildEffectiveSystemPrompt(
                promptConfig, tools, request.model(),
                Path.of(session.workingDir()));
        String rawModel = request.model() != null ? request.model() : session.model();
        String model = providerRegistry.resolveModelAlias(rawModel);
        int maxTurns = request.maxTurns() != null ? request.maxTurns() : QueryConfig.DEFAULT_MAX_TURNS;
        int contextWindow = getContextWindow(model);

        QueryConfig config = QueryConfig.withDefaults(
                model, systemPrompt, tools,
                tools.stream().map(Tool::toToolDefinition).toList(),
                QueryConfig.getRecommendedMaxTokens(modelRegistry, model),
                contextWindow,
                new ThinkingConfig.Adaptive(),
                maxTurns, "rest-api-conversation"
        );

        // 3. 初始化状态 — 加载历史消息
        // Do not install a request-specific notifier or override the saved session mode.
        ToolUseContext toolCtx = ToolUseContext.of(
                session.workingDir(), request.sessionId())
                .withPermissionNotifier(null);  // 明确标注: REST无pusher
        log.debug("REST API conversation: permissionMode={}, notifier=null (by design)",
                effectiveMode);
        QueryLoopState state = new QueryLoopState(new ArrayList<>(session.messages()), toolCtx);
        SessionMessagePersistence.attach(
                state, sessionManager, request.sessionId(), "REST /api/query/conversation");

        // 追加新用户消息
        Message.UserMessage requestMessage = new Message.UserMessage(
                UUID.randomUUID().toString(), Instant.now(),
                List.of(new ContentBlock.TextBlock(request.prompt())),
                null, null);
        state.addMessage(requestMessage);

        // 4. 执行查询
        ResultCollectingHandler handler = new ResultCollectingHandler();
        if (com.aicodeassistant.engine.HandoffContextService.isMerged(session.config()))
            config = handoffContext.configure(config, state, session.config(), request.allowedTools(), request.disallowedTools());
        QueryEngine.QueryResult result = queryEngine.execute(config, state, handler);

        String finalText = extractCurrentRunText(result.messages(), state, result.isSuccess());

        return ResponseEntity.ok(new QueryResponse(
                request.sessionId(),
                finalText,
                result.totalUsage(),
                0.0,
                handler.getToolCalls(),
                result.stopReason(),
                result.isSuccess() ? null : result.error()
        ));
        }
    }

    // ═══ 私有方法 ═══

    private SessionExecutionGate.Token acquireSession(String sessionId) {
        // 门闩 Key 对 sessionId 有 requireNonNull 断言；空会话 id 必须复用既有
        // SessionNotFoundException → 404 映射，不能让 NPE 漂移为 500。
        if (sessionId == null || sessionId.isBlank()) {
            throw new SessionNotFoundException(String.valueOf(sessionId));
        }
        SessionExecutionGate.Token token = executionGate.tryAcquire(
                sessionManager.dataSourceIdentity(), sessionId);
        if (token == null) throw new SessionExecutionBusyException(sessionId);
        return token;
    }

    private record LockedSession(SessionData session, SessionExecutionGate.Token token)
            implements AutoCloseable {
        @Override public void close() { token.close(); }
    }

    /** Existing sessions are locked before their history is read. */
    private LockedSession resolveSessionLocked(QueryRequest request) {
        SessionExecutionGate.Token token = null;
        try {
            if (request.sessionId() != null && !request.sessionId().isBlank()) {
                token = acquireSession(request.sessionId());
            }
            SessionData session = resolveSession(request);
            requireMatchingMode(session.sessionId(), request.permissionMode());
            if (token == null) token = acquireSession(session.sessionId());
            return new LockedSession(session, token);
        } catch (RuntimeException failure) {
            if (token != null) token.close();
            throw failure;
        }
    }

    private SessionData resolveSession(QueryRequest request) {
        rejectClientWorkingDirectory(request.workingDirectory());
        if (request.sessionId() != null
                && !request.sessionId().isBlank()) {
            if (request.projectId() != null
                    && !request.projectId().isBlank()) {
                throw new RequestValidationException(
                        "QUERY_PROJECT_WITH_SESSION_UNSUPPORTED",
                        "projectId can only be used when creating a Session");
            }
            SessionData existing = sessionManager.loadSession(request.sessionId())
                    .orElseThrow(() ->
                            new SessionNotFoundException(
                                    request.sessionId()));
            projectWorkspaces.requireCurrentBinding(
                    existing.workingDir());
            return existing;
        }
        String model = request.model() != null ? request.model() : providerRegistry.getDefaultModel();
        String workingDir = projectWorkspaces.resolveWorkspace(
                request.projectId()).toString();
        String sessionId = sessionManager.createSession(
                model, workingDir, request.permissionMode() != null
                        ? request.permissionMode() : PermissionMode.AUTO_APPROVE);
        return sessionManager.loadSession(sessionId)
                .orElseThrow(() ->
                        new SessionNotFoundException(sessionId));
    }

    private PermissionMode requireMatchingMode(String sessionId, PermissionMode requested) {
        PermissionMode actual = permissionModeManager.getMode(sessionId);
        if (requested != null && requested != actual) {
            throw new PermissionModeMismatchException();
        }
        return actual;
    }

    private void rejectClientWorkingDirectory(
            String workingDirectory) {
        if (workingDirectory != null
                && !workingDirectory.isBlank()) {
            throw new RequestValidationException(
                    "QUERY_WORKING_DIRECTORY_UNSUPPORTED",
                    "Use projectId or an existing sessionId instead "
                            + "of workingDirectory");
        }
    }

    private void rejectUnsupportedMaxBudgetUsd(Double maxBudgetUsd) {
        if (maxBudgetUsd != null) {
            throw new RequestValidationException(
                    "QUERY_BUDGET_USD_UNSUPPORTED",
                    "当前不支持金额预算；请省略该字段或传 null，系统不会执行金额上限。");
        }
    }

    private List<Tool> assembleToolPool(List<String> allowedTools, List<String> disallowedTools) {
        List<Tool> tools = toolRegistry.getEnabledTools();
        if (allowedTools != null && !allowedTools.isEmpty()) {
            Set<String> allowed = Set.copyOf(allowedTools);
            tools = tools.stream().filter(t -> allowed.contains(t.getName())).toList();
        }
        if (disallowedTools != null && !disallowedTools.isEmpty()) {
            Set<String> denied = Set.copyOf(disallowedTools);
            tools = tools.stream().filter(t -> !denied.contains(t.getName())).toList();
        }
        return tools;
    }

    private String buildUserMessage(String prompt, QueryContext context) {
        StringBuilder sb = new StringBuilder();
        if (context != null && context.stdin() != null && !context.stdin().isBlank()) {
            sb.append("<stdin>\n").append(context.stdin()).append("\n</stdin>\n\n");
        }
        if (prompt != null) {
            sb.append(prompt);
        }
        return sb.toString();
    }

    private int getContextWindow(String model) {
        // 使用 ModelRegistry 的四级查询（自定义→Provider→内置→默认），
        // 避免 Provider 未注册千问模型能力时 fallback 到 DEFAULT(8192) 导致 effectiveWindow 为负
        return modelRegistry.getContextWindowForModel(model);
    }

    private String extractCurrentRunText(List<Message> messages, QueryLoopState state, boolean success) {
        Set<String> allowed = new HashSet<>(state.getCurrentRunAssistantIds());
        String preferred = success ? state.getCurrentRunFinalMessageId() : null;
        for (int i = messages.size() - 1; i >= 0; i--) {
            Message msg = messages.get(i);
            if (msg instanceof Message.AssistantMessage assistant
                    && allowed.contains(assistant.uuid())
                    && (preferred == null || preferred.equals(assistant.uuid()))
                    && assistant.content() != null) {
                StringBuilder sb = new StringBuilder();
                for (var block : assistant.content()) {
                    if (block instanceof ContentBlock.TextBlock text) {
                        sb.append(text.text());
                    }
                }
                if (!sb.isEmpty()) return sb.toString();
            }
        }
        return "";
    }

    private void sendEvent(SseEmitter emitter, String eventName, Object data) {
        try {
            emitter.send(SseEmitter.event()
                    .name(eventName)
                    .data(data, MediaType.APPLICATION_JSON));
        } catch (IOException e) {
            // 客户端断开连接，静默忽略
        }
    }

    // ═══ 内部回调实现 ═══

    /**
     * 结果收集处理器 — 收集工具调用信息。
     */
    private static class ResultCollectingHandler implements QueryMessageHandler {
        private final List<ToolCallSummary> toolCalls = new ArrayList<>();

        @Override public void onTextDelta(String text) {}

        @Override
        public void onToolUseStart(String toolUseId, String toolName) {
            toolCalls.add(new ToolCallSummary(toolName, null, null, false));
        }

        @Override
        public void onToolUseComplete(String toolUseId, ContentBlock.ToolUseBlock toolUse) {}

        @Override
        public void onToolResult(String toolUseId, ContentBlock.ToolResultBlock result) {
            // 更新最后一个工具调用的结果
            if (!toolCalls.isEmpty()) {
                ToolCallSummary last = toolCalls.getLast();
                toolCalls.set(toolCalls.size() - 1,
                        new ToolCallSummary(last.tool(), last.input(),
                                result.content(), result.isError()));
            }
        }

        @Override
        public void onAssistantMessage(Message.AssistantMessage message) {}

        public List<ToolCallSummary> getToolCalls() { return toolCalls; }
    }

    // ═══ DTO Records ═══

    public record QueryRequest(
            String prompt,
            String model,
            String systemPrompt,
            String appendSystemPrompt,
            PermissionMode permissionMode,
            Integer maxTurns,
            Double maxBudgetUsd,
            List<String> allowedTools,
            List<String> disallowedTools,
            String projectId,
            String sessionId,
            String workingDirectory,
            Integer timeoutSeconds,
            String outputFormat,
            QueryContext context
    ) {}

    public record QueryContext(
            String stdin,
            List<String> files
    ) {}

    public record ConversationRequest(
            String sessionId,
            String prompt,
            String model,
            String systemPrompt,
            String appendSystemPrompt,
            PermissionMode permissionMode,
            Integer maxTurns,
            Double maxBudgetUsd,
            List<String> allowedTools,
            List<String> disallowedTools,
            String workingDirectory,
            Integer timeoutSeconds
    ) {}

    public record QueryResponse(
            String sessionId,
            String result,
            Usage usage,
            double costUsd,
            List<ToolCallSummary> toolCalls,
            String stopReason,
            String error
    ) {}

    public record ToolCallSummary(
            String tool,
            Object input,
            String output,
            boolean isError
    ) {}
}