package cjagent

import std.collection.*
import std.env.*
import std.fs.Path
import stdx.encoding.json.*

/*
 * Agent 主循环：把「用户意图 → 模型思考 → 工具执行 → 观察结果」串成闭环。
 *
 * 每收到一句用户输入，就反复请求模型：
 *   - 模型返回工具调用 → 逐个执行并把结果回灌，继续下一轮；
 *   - 模型返回纯文本   → 视为最终答复，本轮结束。
 * 直到给出答复或达到 maxSteps 上限。
 */

class Agent {
    let cfg: Config
    let llm: Llm
    let tools: Array<Tool>
    let byName: HashMap<String, Tool>

    init(cfg: Config) {
        this.cfg = cfg
        this.llm = Llm(cfg)
        this.tools = buildTools(cfg)
        let index = HashMap<String, Tool>()
        for (t in tools) {
            index.add(t.name, t)
        }
        this.byName = index
    }

    func run(): Unit {
        printBanner()
        let messages = ArrayList<Message>()
        messages.add(Message.system(systemPrompt(cfg.root)))
        match (cfg.task) {
            case Some(task) =>
                println("📋 任务：${task}")
                handleTurn(messages, task)
            case None => repl(messages)
        }
    }

    // 交互式对话循环。
    func repl(messages: ArrayList<Message>): Unit {
        println("进入交互模式，输入任务开始；/exit 或 Ctrl+Z 回车退出。")
        while (true) {
            print("\n🧑 你> ")
            match (getStdIn().readln()) {
                case Some(line) =>
                    let text = line.trimAscii()
                    if (text.isEmpty()) {
                        continue
                    }
                    if (text == "/exit" || text == "/quit") {
                        break
                    }
                    handleTurn(messages, text)
                case None => break
            }
        }
        println("\n再见。")
    }

    // 处理一轮用户输入：多步工具调用直至模型给出最终答复。
    func handleTurn(messages: ArrayList<Message>, userInput: String): Unit {
        messages.add(Message.user(userInput))
        var step = 0
        while (step < cfg.maxSteps) {
            step += 1
            let (content, toolCalls) = try {
                llm.complete(messages, tools)
            } catch (e: Exception) {
                println("[警告] 调用模型失败：${e.message}")
                return
            }

            if (toolCalls.size == 0) {
                if (!content.isEmpty()) {
                    println("\n🤖 ${content}")
                }
                messages.add(Message.assistant(content))
                return
            }

            // 模型请求调用工具：记录 assistant 消息（含 tool_calls），再逐个执行。
            if (!content.isEmpty()) {
                println("\n🤖 ${content}")
            }
            messages.add(Message.assistant(content, tcs: toolCalls))
            for (tc in toolCalls) {
                messages.add(Message.tool(tc.id, dispatch(tc)))
            }
        }
        println("[警告] 已达最大步数（${cfg.maxSteps}），中止本轮。")
    }

    // 执行单个工具调用并打印简洁轨迹。
    func dispatch(tc: ToolCall): String {
        println("  » ${tc.name} ${preview(tc.arguments, 120)}")
        let tool = match (byName.get(tc.name)) {
            case Some(t) => t
            case None => return "错误：未知工具 ${tc.name}"
        }
        let args = try {
            JsonValue.fromStr(tc.arguments).asObject()
        } catch (_: Exception) {
            return "错误：无法解析工具参数：${tc.arguments}"
        }
        let result = tool.handler(args)
        println("     ↳ ${preview(firstLine(result), 160)}")
        return result
    }

    func printBanner(): Unit {
        println("┌─ 极简仓颉 Coding Agent ─────────────────────")
        println("│ 沙箱目录: ${cfg.root}")
        println("│ 模型:     ${cfg.model}")
        println("│ 端点:     ${cfg.endpoint}")
        println("│ 工具:     read_file / write_file / list_dir / run_command")
        println("└─────────────────────────────────────────────")
    }
}

func systemPrompt(root: Path): String {
    return """
你是一个运行在命令行中的编码助手（coding agent），在一个受限的工作目录里帮用户完成开发任务。

工作目录（沙箱根）：${root}
你只能在该目录内读写文件、执行命令；所有路径都相对于它，无法也不得访问其外层。

可用工具：
- read_file(path)            读取文件文本
- write_file(path, content)  写入/覆盖文件（父目录自动创建）
- list_dir(path)             列目录，path 可省略表示根目录
- run_command(command)       在工作目录内执行 shell 命令（编译、运行、测试）

工作准则：
1. 先用 list_dir / read_file 摸清现状，再动手改。
2. 只做必要的最小改动，代码保持整洁规范。
3. 需要验证时用 run_command 编译或运行。
4. 任务完成后用简洁中文说明你做了什么，并停止调用工具。
"""
}

// 取首行，便于轨迹展示。
func firstLine(s: String): String {
    match (s.indexOf("\n")) {
        case Some(i) => s[0..i]
        case None => s
    }
}

// 单行化并截断，用于打印工具调用轨迹。
func preview(s: String, limit: Int64): String {
    truncate(s.replace("\n", " "), limit)
}