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()
