from __future__ import annotations

import os
import uuid
from pathlib import Path

from dotenv import load_dotenv

from klynx import create_agent, setup_model


load_dotenv()


def _extract_plan_request(user_input: str) -> tuple[str, bool]:
    """显式规划入口。默认直接执行，只有 /plan|plan:|spec: 才进入规划。"""
    stripped = str(user_input or "").strip()
    lowered = stripped.lower()
    if lowered in {"/plan", "plan", "spec"}:
        return "", True

    for prefix in ("/plan ", "plan:", "spec:"):
        if lowered.startswith(prefix):
            return stripped[len(prefix) :].strip(), True

    return stripped, False


def _build_agent(working_dir: str):
    model = setup_model(
        "moonshot",
        "kimi-k2.5",
        os.getenv("KIMI_API_KEY"),
        api_base="https://api.moonshot.cn/v1",
    )
    agent = create_agent(
        working_dir=working_dir,
        model=model,
        memory_dir=working_dir,
        load_project_docs=False,
    )
    # agent.add_tools("group:interactive")

    mcp_config = Path(working_dir) / ".klynx" / "mcp_servers.json"
    if mcp_config.exists():
        agent.add_mcp(str(mcp_config))
    return agent


def _copy_planning_context(agent, source_thread_id: str, target_thread_id: str) -> None:
    current_state = agent.app.get_state(
        {"configurable": {"thread_id": source_thread_id}}
    )
    if not current_state or not current_state.values:
        return

    state_to_copy = {
        "messages": current_state.values.get("messages", []),
        "overall_goal": current_state.values.get("overall_goal", ""),
        "context_summary": current_state.values.get("context_summary", ""),
    }
    agent.app.update_state(
        {"configurable": {"thread_id": target_thread_id}}, state_to_copy
    )


def do_planning(agent, user_input: str, working_dir: str, thread_id: str) -> str:
    """生成 Spec/spec.md，用户确认后返回 spec 内容。"""
    spec_path = Path(working_dir) / "Spec" / "spec.md"
    planner_task = f"""
用户有一个新的项目需求: "{user_input}"

请你作为架构师和需求规划师，分析这个需求，并在当前目录下执行：
1. 检查是否存在 `Spec` 文件夹，不存在则创建。
2. 在 `Spec` 文件夹下创建 `spec.md`。
3. 在 `spec.md` 中写下详细的实现规划、功能列表和技术细节。

完成后告诉我你已经写好了。
""".strip()

    print("\n>>> [规划阶段] 生成设计文档...")
    agent.run_terminal_agent_stream(task=planner_task, thread_id=thread_id)

    if not spec_path.exists():
        print(f"\n[Error] 未生成 {spec_path}")
        return ""

    print("\n" + "#" * 60)
    print(f"[人工确认] 已生成: {spec_path}")
    print("输入 'y' 继续执行，输入 'n' 取消。")

    while True:
        user_confirm = input("是否继续执行？(y/n) > ").strip().lower()
        if user_confirm in {"", "y", "yes"}:
            return spec_path.read_text(encoding="utf-8")
        if user_confirm in {"n", "no"}:
            print("[系统] 规划已取消，等待下一条输入。\n")
            return ""


def do_execution(agent, task: str, thread_id: str) -> dict:
    """执行任务。"""
    print("\n>>> [执行阶段] Agent 开始工作...")
    return agent.run_terminal_agent_stream(task=task, thread_id=thread_id)


def _print_context(agent, thread_id: str) -> None:
    state_values = agent.get_context(thread_id)
    if not state_values:
        print(f"\n[上下文] Thread: {thread_id} (当前为空)\n")
        return

    msgs = state_values.get("messages", [])
    approx_tokens = 0
    try:
        from klynx.agent.context_manager import TokenCounter

        approx_tokens = TokenCounter.count_message_tokens(msgs) if msgs else 0
    except Exception:
        approx_tokens = (
            sum(
                len(m.content)
                for m in msgs
                if hasattr(m, "content") and isinstance(m.content, str)
            )
            // 2
        )

    goal = state_values.get("overall_goal", "")
    print(f"\n[上下文] Thread: {thread_id}")
    print(f"  - 历史消息: {len(msgs)} 条")
    print(f"  - 估算用量: ~{approx_tokens} Tokens")
    print(f"  - 当前目标: {goal if goal else '(无)'}\n")


def main():
    working_dir = os.getcwd()

    print("=" * 60)
    print("Klynx Agent — 精简版计划/执行示例")
    print(f"工作目录: {working_dir}")
    print("=" * 60)

    try:
        agent = _build_agent(working_dir)
    except Exception as exc:
        print(f"[Error] 模型或 Agent 初始化失败: {exc}")
        return

    print("\n" + "=" * 60)
    print("      输入 '/plan 需求描述' 显式进入 Spec 规划流程")
    print("=" * 60 + "\n")

    thread_id = uuid.uuid4().hex[:8]
    round_count = 0
    total_tokens = 0

    while True:
        try:
            user_input = input(f"[轮次 {round_count + 1}] > ").strip()

            task_input, use_planning = _extract_plan_request(user_input)
            if not task_input.strip():
                if use_planning:
                    print("\n[系统] /plan 后需要补充需求描述\n")
                continue

            result = {}
            if use_planning:
                print("\n[规划] 使用显式 /plan 入口生成 Spec...")
                temp_plan_thread_id = f"plan_{uuid.uuid4().hex[:8]}"
                _copy_planning_context(agent, thread_id, temp_plan_thread_id)
                spec_content = do_planning(
                    agent, task_input, working_dir, temp_plan_thread_id
                )
                if spec_content:
                    executor_task = f"""
你是一个执行工程师。以下是我们在 `Spec/spec.md` 中确认的执行计划。

<spec_content>
{spec_content}
</spec_content>

请严格按照这份计划直接开始实现。
""".strip()
                    result = do_execution(
                        agent, task=executor_task, thread_id=thread_id
                    )
            else:
                result = do_execution(agent, task=task_input, thread_id=thread_id)

            round_count += 1
            if result:
                round_tokens = int(result.get("total_tokens", 0) or 0)
                total_tokens += round_tokens
                print(f"\n[第 {round_count} 轮完成]")
                print(f"本轮 Token: {round_tokens} | 累计 Token: {total_tokens}\n")

        except KeyboardInterrupt:
            print(f"\n\n[系统] 用户中断，退出程序")
            print(f"[统计] 共 {round_count} 轮对话，累计 Token: {total_tokens}")
            break
        except EOFError:
            break


if __name__ == "__main__":
    main()
