import 'dotenv/config';
import { createDeepAgent } from "deepagents";
import { AIMessageChunk, ToolMessage, todoListMiddleware } from "langchain";
import { HumanMessage } from "@langchain/core/messages";
import { ChatDeepSeek } from "@langchain/deepseek";
import config from '../../../config.js';
import logger from '../../../utils/logger.js';
import { createRestrictedBackend } from "../../infra/restricted_fs_backend.js";

// 密钥由 conf.json 统一提供；百度地图 skill 通过 env_get 读环境变量，此处桥接
process.env.TAVILY_API_KEY ??= config.agent.tavily?.apiKey;
process.env.BAIDU_MAP_AUTH_TOKEN ??= config.agent.baiduMap?.authToken;
import EscortAdminPrompts from "./prompts.js";
import { createPatientRecordSubagent } from "../../subagent/patient_record.js";
import { createResearchSubagent } from "../../subagent/research.js";
import { createHospitalInfoSubagent } from "../../subagent/hospital_info.js";
import {
    getEnvTool, webFetchTool, webSearchTool, getCalendarInfoTool,
    getLunarCalendarInfoTool, getYearHolidaysTool, getYearTermsTool, getLatLngTool,
    httpGetTool, httpPostTool, escortRecordQueryTool, escortRecordSetTool
} from "../../infra/tools/index.js";

export default class EscortAdminAgent {
    constructor() {
        // 按用户隔离会话：userId -> { agent, messages }
        this.sessions = new Map();
        this.maxSessions = 100;
    }

    clearMessages(userInfo) {
        this.sessions.delete(userInfo._id);
    }

    // msg: { agent: 'escort-admin', type: 'chat', ts: "2023-08-01 10:00:00", content: "你好" }
    async streamChat(userInfo, msgs, callback) {
        if (!msgs.length) {
            return;
        }

        const userId = userInfo._id;

        // LRU：按用户取会话，并更新使用顺序
        let session = this.sessions.get(userId);
        if (session) {
            this.sessions.delete(userId);
        } else {
            session = { agent: null, messages: [] };
        }
        this.sessions.set(userId, session);

        // 超出上限时淘汰最久未使用的会话，避免内存泄漏
        if (this.sessions.size > this.maxSessions) {
            const oldest = this.sessions.keys().next().value;
            this.sessions.delete(oldest);
        }

        msgs.forEach(msg => {
            if (msg.type === "clear") {
                session.messages = [];
                session.agent = null;
            } else {
                session.messages.push(new HumanMessage(`${msg.ts} - ${msg.content}`));
            }
        });

        if (session.messages.length === 0) {
            return;
        }

        const agent = session.agent ?? this._genAgent(userInfo);
        session.agent = agent;

        const INTERESTING_NODES = new Set(["model_request", "tools"]);
        for await (const [namespace, mode, data] of await agent.stream(
            { messages: session.messages },
            {
                recursion_limit: 50,
                streamMode: ["updates", "messages", "custom"], subgraphs: true,
                configurable: {
                    thread_id: userId
                }
            })) {
            const isSubagent = namespace.some(s => s.startsWith("tools:"));
            const source = isSubagent ? "subagent" : "main";
            if (mode === "updates") {
                // Main agent updates (empty namespace)
                if (namespace.length === 0) {
                    for (const [nodeName, data_] of Object.entries(data)) {
                        if (!INTERESTING_NODES.has(nodeName)) continue;
                        if (nodeName === "tools") {
                            // 工具结果必须并入历史，否则下一轮会出现 tool_calls 缺少对应 tool 结果的断链
                            for (const msg of data_.messages ?? []) {
                                if (msg.type === "tool") {
                                    session.messages.push(msg);
                                    logger.info(`Tool complete: ${msg.name}, Result: ${String(msg.content).slice(0, 200)}`);
                                }
                            }
                        } else if (nodeName === "model_request") {
                            session.messages.push(...data_.messages);
                        }
                    }
                } else {
                    // Subagent updates (non-empty namespace)
                    for (const nodeName of Object.keys(data)) {
                        if (!INTERESTING_NODES.has(nodeName)) continue;
                        logger.info(`[${namespace[0]}] step: ${nodeName}`);
                    }
                }
            } else if (mode === "messages") {
                const [message] = data;
                if (message.tool_call_chunks?.length) {
                    continue;
                }
                if (AIMessageChunk.isInstance(message)) {
                    if (message.text) {
                        callback(source, "ai", message.text, message.id);
                    }
                    if (message.additional_kwargs.reasoning_content) {
                        callback(source, "reasoning", message.additional_kwargs.reasoning_content, message.id);
                    }
                }
                if (ToolMessage.isInstance(message) && message.text) {
                    callback(source, "tool", message.text, message.id);
                }
            } else if (mode === "custom") {
                logger.info("custom: ", data);
            }
        }
    }

    _genAgent(userInfo) {
        const rootDir = process.cwd();
        // 受限后端：锁定沙箱根并禁止读取 conf.json/.env/logs 及用户数据目录
        const backend = createRestrictedBackend({ rootDir, dataDir: config.agent.dataDir });

        this.flashModel = new ChatDeepSeek({
            model: config.agent.deepseek.flashModel,
            apiKey: config.agent.deepseek.apiKey,
            temperature: 0.0
        });
        this.proModel = new ChatDeepSeek({
            model: config.agent.deepseek.proModel,
            apiKey: config.agent.deepseek.apiKey,
            temperature: 0.3
        });

        return createDeepAgent({
            name: "deep-agent",
            model: this.flashModel,
            systemPrompt: EscortAdminPrompts.buildSystemPrompt(userInfo),
            backend,
            middleware: [todoListMiddleware()], // deepagents 1.12 起 todo 计划工具改为显式开启，多步记录操作用
            tools: [getEnvTool, webFetchTool, webSearchTool, getLatLngTool, httpGetTool, httpPostTool,
                getCalendarInfoTool, getLunarCalendarInfoTool, getYearHolidaysTool, getYearTermsTool,
                escortRecordQueryTool, escortRecordSetTool],
            subagents: [createPatientRecordSubagent(this.flashModel), createResearchSubagent(this.proModel), createHospitalInfoSubagent(this.flashModel)]
        });
    }
}

const adminAgent = new EscortAdminAgent();
export { adminAgent };