import argparse
import glob
import json
import os
import subprocess
import urllib.request

BASE_URL = os.environ.get("AI_BUILDER_BASE_URL", "https://space.ai-builders.com/backend/v1")
API_KEY = os.environ.get("AI_BUILDER_TOKEN", "")
WORKDIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "workspace"))


def safe_path(rel_path: str) -> str:
    target = os.path.abspath(os.path.join(WORKDIR, rel_path))
    if not (target == WORKDIR or target.startswith(WORKDIR + os.sep)):
        raise ValueError(f"路径越界（只允许访问 workspace/ 目录内文件）: {rel_path}")
    return target


def read_file(path: str) -> str:
    try:
        full = safe_path(path)
        with open(full, "r", encoding="utf-8") as f:
            lines = f.readlines()
        return "".join(f"{i+1}: {line}" for i, line in enumerate(lines))
    except Exception as e:
        return f"读取失败: {e}。请检查文件路径是否存在于 workspace/ 下。"


def write_file(path: str, content: str) -> str:
    try:
        full = safe_path(path)
        os.makedirs(os.path.dirname(full), exist_ok=True)
        with open(full, "w", encoding="utf-8") as f:
            f.write(content)
        return f"已写入 {path}（共 {len(content)} 字节）。下一步请运行测试验证改动。"
    except Exception as e:
        return f"写入失败: {e}"


def run_bash(command: str) -> str:
    try:
        res = subprocess.run(
            command, shell=True, cwd=WORKDIR, capture_output=True, text=True, timeout=15
        )
        out = (res.stdout + res.stderr).strip()
        if len(out) > 2000:
            out = out[:2000] + "\n...[输出过长已截断]"
        return f"exit_code={res.returncode}\n{out}"
    except Exception as e:
        return f"命令执行异常: {e}"


def load_skill_index() -> str:
    entries = []
    for p in sorted(glob.glob(os.path.join(WORKDIR, ".agents/skills/*/SKILL.md"))):
        rel = os.path.relpath(p, WORKDIR)
        name, desc = os.path.basename(os.path.dirname(p)), ""
        with open(p, "r", encoding="utf-8") as f:
            for line in f:
                if line.startswith("name:"):
                    name = line.split(":", 1)[1].strip()
                elif line.startswith("description:"):
                    desc = line.split(":", 1)[1].strip()
        entries.append(f"- {name} ({rel}): {desc}")
    return "\n".join(entries)


TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "读取工作区内的文本文件，返回带行号的内容。也可以用来按需读取 .agents/skills/ 下的 SKILL.md。",
            "parameters": {
                "type": "object",
                "properties": {"path": {"type": "string", "description": "相对于工作区根目录的文件路径"}},
                "required": ["path"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "write_file",
            "description": "将完整文件内容写入工作区内的指定路径。",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string", "description": "相对于工作区根目录的文件路径"},
                    "content": {"type": "string", "description": "要写入的完整文本内容"},
                },
                "required": ["path", "content"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "bash",
            "description": "在工作区目录下执行 shell 命令（如运行 pytest 或 CLI 脚本），返回退出码与输出。",
            "parameters": {
                "type": "object",
                "properties": {"command": {"type": "string", "description": "要运行的 shell 命令"}},
                "required": ["command"],
            },
        },
    },
]


def call_llm(model: str, messages: list) -> dict:
    req = urllib.request.Request(
        f"{BASE_URL}/chat/completions",
        data=json.dumps({"model": model, "messages": messages, "tools": TOOLS, "temperature": 0.0}).encode("utf-8"),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {API_KEY}"},
    )
    with urllib.request.urlopen(req, timeout=60) as resp:
        return json.loads(resp.read().decode("utf-8"))


def run_agent(task: str, model: str = "deepseek", max_turns: int = 10):
    skill_index = load_skill_index()
    sys_prompt = (
        "你是一个精简的编码 Agent，工作在受控工作区中。\n"
        "遇到特定领域任务时，请先用 read_file 阅读下方目录中对应的 SKILL.md，严格遵循其中的检查单与 CLI 规范；"
        "修改代码后必须用 bash 跑测试验证，直到测试全部通过再给出最终汇报。\n\n"
        f"可用 Skills 目录（仅含摘要，按需读取全文）：\n{skill_index}"
    )
    messages = [{"role": "system", "content": sys_prompt}, {"role": "user", "content": task}]
    for turn in range(1, max_turns + 1):
        resp = call_llm(model, messages)
        msg = resp["choices"][0]["message"]
        messages.append(msg)
        tool_calls = msg.get("tool_calls") or []
        if not tool_calls:
            print(f"\n[第 {turn} 轮 · 完成]\n{msg.get('content', '').strip()}")
            return
        for tc in tool_calls:
            fn = tc["function"]["name"]
            args = json.loads(tc["function"]["arguments"])
            print(f"[第 {turn} 轮 · 调用工具] {fn}({json.dumps(args, ensure_ascii=False)})")
            if fn == "read_file":
                out = read_file(args["path"])
            elif fn == "write_file":
                out = write_file(args["path"], args["content"])
            elif fn == "bash":
                out = run_bash(args["command"])
            else:
                out = f"未知工具: {fn}"
            first_line = out.splitlines()[0] if out else ""
            print(f"  -> {first_line}")
            messages.append({"role": "tool", "tool_call_id": tc["id"], "content": out})


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="deepseek")
    parser.add_argument(
        "--task",
        default="请检查 settlement.py 里的订单结算函数为什么没通过测试，按照团队的结算核对规范修好它，并跑通全部校验。",
    )
    args = parser.parse_args()
    run_agent(args.task, args.model)