"""
Klynx agent orchestrations.

Define one subclass per concrete agent to customize behavior.

`LOCKED_SYSTEM_PROMPT` is framework-owned and should not be overridden for the
main agent loop. Prompt extension is append-only via constructor/invoke params.
"""

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langgraph.graph import END, StateGraph

from .graph import GraphKlynxAgent
from .state import AgentState


class KlynxGeneralAgent(GraphKlynxAgent):
    """Generic Klynx orchestration with the baseline system prompt."""

    LOCKED_SYSTEM_PROMPT = """<system_instructions>
  <role>你是 Klynx，一个以完成任务为目标的工程执行型智能体。</role>

  <decision_policy>
    1. 请求明确且无需外部操作：直接回答。
    2. 请求涉及文件/命令/浏览器/联网/技能等执行动作：使用工具推进。
    3. 默认处于执行态。除非用户明确要求计划、方案比较、设计讨论或纯问答，否则不要停在分析。
    4. 仅当信息无法从当前环境中自行发现，或存在多种高风险路径需要用户拍板时，才提问澄清。
    5. 若遇到登录/认证/权限阻碍且需要切换方案：必须先向用户确认，不得擅自改路。
  </decision_policy>

  <execution_policy>
    1. 主循环遵循 think -> act -> feedback。若本轮无工具调用，则自然收敛并给出可交付结果。
    2. 可以在必要时调用 `update_task_state` 修正 overall_goal / current_focus。
    3. 复杂任务可使用 `write_todos` 维护步骤；不要因为任务是多步的就默认先展开长计划。
    4. 默认工作流优先走 execute_command(rg) -> read_file(局部) -> apply_patch(最小修改) -> execute_command/语义验证。
    5. 每轮只保留 1 个主假设，最多 1 个备选；已证伪假设不得重开，除非出现相反新证据。
    6. 调用技能前，先用 `load_skill` 读取 SKILL.md，再按技能流程执行。
    7. 工具失败时先定位失败原因并缩小下一步动作；避免无变化重复调用同一失败操作。
    8. 一旦已有直接修改点，优先修改并验证，不要继续泛读文件或重复解释已知结论。
    9. 调试过程中创建的中间文件，结束后必须清除。
  </execution_policy>

  <completion_policy>
    1. 结束条件：无待执行工具调用、无待轮询命令、且当前有可交付文本。
    2. 当系统要求澄清时，向用户提出具体问题并等待回复。
    3. 执行任务应当直奔目标，以最快的速度和最优的方案完成任务。
    4. 完成任务后应该对完成过程进行检查和复盘，并将总结告知用户。
  </completion_policy>

  <response_style>
    简洁、准确、可执行；避免无意义寒暄与夸张措辞。
  </response_style>
</system_instructions>"""

    ASK_SYSTEM_PROMPT = (
        "You are Klynx Ask Assistant. Answer the user directly and concisely."
    )

    def __init__(self, *args, append_system_prompt: str = "", **kwargs):
        """
        Method 1: inject appended prompt via constructor.
        """
        super().__init__(*args, **kwargs)
        self._init_system_prompt_append = (append_system_prompt or "").strip()
        self._runtime_system_prompt_append = ""

    def _inject_system_prompt_node(self, state: AgentState) -> dict:
        """
        Inject user-appendable system prompt extension.

        The core agent prompt stays locked. Users can append extra instructions
        in two ways:
        1) constructor `append_system_prompt`
        2) per-call `invoke(..., system_prompt_append=...)`
        """
        init_append = (getattr(self, "_init_system_prompt_append", "") or "").strip()
        call_append = (state.get("system_prompt_append", "") or "").strip()

        merged_parts = [p for p in (init_append, call_append) if p]
        merged_append = "\n\n".join(merged_parts).strip()

        self._runtime_system_prompt_append = merged_append
        if merged_append:
            self._emit("info", "[Prompt] 已注入追加系统提示词")
        return {"system_prompt_append": merged_append}

    def _build_graph(self) -> StateGraph:
        """
        构建 LangGraph 状态图。
        流程：init -> inject_system_prompt -> load_context -> agent -> act -> feedback
        -> agent/end(条件分支)。
        """
        workflow = StateGraph(AgentState)

        workflow.add_node("init", self._init_node)
        workflow.add_node("inject_system_prompt", self._inject_system_prompt_node)
        workflow.add_node("load_context", self._load_context_node)
        workflow.add_node("agent", self._model_inference_node)
        workflow.add_node("act", self._act_node)
        workflow.add_node("feedback", self._feedback_node)

        workflow.set_entry_point("init")
        workflow.add_edge("init", "inject_system_prompt")
        workflow.add_edge("inject_system_prompt", "load_context")
        workflow.add_edge("load_context", "agent")
        workflow.add_edge("agent", "act")
        workflow.add_edge("act", "feedback")
        workflow.add_conditional_edges(
            "feedback",
            self._should_continue,
            {
                "agent": "agent",
                "tools": "act",
                "end_direct": END,
            },
        )
        return workflow

    def ask(self, message: str, system_prompt: str = None, thread_id: str = "default"):
        """
        Direct ask API: call LLM directly without agent-node prompt injection.
        """
        if self.model is None:
            yield {"type": "error", "content": "Model is not configured."}
            return

        default_system = (getattr(self, "ASK_SYSTEM_PROMPT", "") or "").strip()
        if not default_system:
            default_system = (
                "You are Klynx Ask Assistant. Answer the user directly and concisely."
            )
        sys_prompt = (system_prompt or default_system).strip()

        history_msgs = []
        if thread_id:
            config = {"configurable": {"thread_id": thread_id}}
            current_state = self.app.get_state(config)
            if current_state and current_state.values:
                history_msgs = current_state.values.get("messages", [])

        messages = [SystemMessage(content=sys_prompt)]
        if history_msgs:
            for msg in history_msgs[-20:]:
                if isinstance(msg, (HumanMessage, AIMessage)):
                    messages.append(msg)
        messages.append(HumanMessage(content=message))

        has_stream = hasattr(self.model, "stream") and callable(self.model.stream)
        full_content = ""
        full_reasoning = ""
        total_tokens = 0

        try:
            if has_stream:
                for chunk in self.model.stream(messages):
                    reasoning_part = chunk.get("reasoning_content", "")
                    content_part = chunk.get("content", "")
                    if reasoning_part:
                        full_reasoning += reasoning_part
                        yield {"type": "reasoning_token", "content": reasoning_part}
                    if content_part:
                        full_content += content_part
                        yield {"type": "token", "content": content_part}
            else:
                response = self.model.invoke(messages)
                full_content = response.content or ""
                full_reasoning = self._extract_reasoning_content(response) or ""
                usage = getattr(response, "usage", None)
                if usage and isinstance(usage, dict):
                    total_tokens = usage.get("total_tokens", 0)
                if full_reasoning:
                    yield {"type": "reasoning", "content": full_reasoning}
                if full_content:
                    yield {"type": "answer", "content": full_content}
        except Exception as e:
            yield {"type": "error", "content": f"ask call failed: {str(e)}"}
            return

        yield {
            "type": "done",
            "content": "",
            "answer": full_content,
            "reasoning": full_reasoning,
            "total_tokens": total_tokens,
        }

    def invoke(
        self,
        task: str,
        thread_id: str = "default",
        thinking_context: bool = False,
        system_prompt_append: str = "",
    ):
        """
        执行用户任务，流式返回事件（生成器）。
        """
        config = {
            "configurable": {"thread_id": thread_id},
            "recursion_limit": 2000,
        }

        messages = [HumanMessage(content=task)]
        initial_state = {
            "messages": messages,
            "thread_id": thread_id,
            "env_snapshot": "",
            "pending_tool_calls": [],
            "iteration_count": 0,
            "working_dir": self.working_dir,
            "task_completed": False,
            "last_action": "",
            "empty_response_count": 0,
            "stall_rounds": 0,
            "stall_round_threshold": 3,
            "tui_stall_rounds": 0,
            "tui_stall_threshold": int(getattr(self, "tui_stall_threshold", 3) or 3),
            "last_tui_screen_hash": "",
            "tui_last_status_tokens": [],
            "total_tokens": 0,
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "max_context_tokens": self.max_context_tokens,
            "task_type": "",
            "overall_goal": task,
            "current_focus": task,
            "initial_user_request": task,
            "has_new_user_input": True,
            "current_task": "",
            "context_summary": "",
            "context_summarized": False,
            "context": "",
            "reasoning_content": "",
            "klynx_docs": "",
            "project_rules": "",
            "progress_summary": "",
            "user_input": task,
            "ended_without_tools": False,
            "command_executions": [],
            "needs_user_confirmation": False,
            "clarification_question": "",
            "loaded_skill_names": [],
            "skill_context": "",
            "tui_guide_loaded": False,
            "active_tool_groups": list(getattr(self, "active_tool_groups", []) or []),
            "active_tool_names": sorted(list(getattr(self, "tools", {}).keys())),
            "active_skill_names": [],
            "active_skill_paths": [],
            "skill_registry_digest": str(
                getattr(self, "skill_registry_digest", "") or ""
            ),
            "thinking_context": thinking_context,
            "system_prompt_append": (system_prompt_append or "").strip(),
            "tool_protocol_mode": str(getattr(self, "tool_protocol_mode", "native")),
            "dedupe_tools": True,
            "tool_dedupe_window": 3,
            "tool_call_history": [],
            "tool_dedupe_hits": 0,
            "repeated_read_hits": 0,
            "read_coverage": {},
            "file_views": {},
            "active_file_view_paths": [],
            "last_read_chunks": [],
            "last_read_fingerprint": "",
            "evidence_index": [],
            "evidence_digest": "",
            "search_hits_index": [],
            "last_search_backend": "",
            "file_candidates": [],
            "trusted_modified_files": [],
            "last_patch_summaries": [],
            "last_mutation": {},
            "recent_mutations": [],
            "pending_verification_targets": [],
            "mutation_truth_digest": "",
            "recent_terminal_events": [],
            "recent_tui_events": [],
            "recent_exec_sessions": [],
            "tui_views": {},
            "active_tui_view_names": [],
            "last_tui_snapshots": [],
            "tui_verification_targets": [
                {
                    "goal": "selection_navigation",
                    "assertion": "A navigation key should move the highlighted selection in a predictable way.",
                    "pass_condition": "selected option changes or highlighted row changes consistently.",
                },
                {
                    "goal": "maze_enter_game",
                    "assertion": "Entering the maze should transition from lobby/menu to the maze gameplay screen.",
                    "pass_condition": "screen contains maze/gameplay markers instead of menu-only content.",
                },
                {
                    "goal": "maze_key_response",
                    "assertion": "After a movement key, at least one gameplay signal should change.",
                    "pass_condition": "step count, player position, or move/block/win status changes.",
                },
                {
                    "goal": "maze_readability",
                    "assertion": "Maze screen should expose distinct markers and readable controls/info.",
                    "pass_condition": "player/goal/wall markers are distinguishable and controls/info remain visible.",
                },
            ],
            "last_tui_verification": {},
            "recent_tui_verifications": [],
            "tui_verification_digest": "",
            "command_verification_targets": [],
            "last_command_verification": {},
            "recent_command_verifications": [],
            "command_verification_digest": "",
            "active_terminal_op": {},
            "active_tui_focus": {},
            "active_exec_session": {},
            "last_terminal_delta": {},
            "last_terminal_failure": {},
            "last_exec_output": {},
            "last_tui_diff": {},
            "last_tui_region": {},
            "last_tui_anchor_match": {},
            "should_plan": False,
            "task_plan": [],
            "current_step_id": "",
            "completed_steps": [],
            "blocked_reason": "",
            "convergence_mode": "normal",
            "convergence_reason": "",
            "next_step_requirements": [],
            "max_tools_per_step": int(getattr(self, "max_tools_per_step", 20) or 20),
            "max_reads_per_file_per_step": int(
                getattr(self, "max_reads_per_file_per_step", 6) or 6
            ),
            "max_retry_per_tool_per_step": int(
                getattr(self, "max_retry_per_tool_per_step", 2) or 2
            ),
            "step_execution_stats": {},
            "tool_artifacts": [],
            "summary_events": [],
            "subtask_history": [],
        }

        self._event_buffer.clear()

        if not hasattr(self, "_cancel_event"):
            import threading

            self._cancel_event = threading.Event()
        else:
            self._cancel_event.clear()

        import threading
        import time

        graph_done = threading.Event()
        graph_error = [None]

        def _run_graph():
            try:
                for _ in self.app.stream(initial_state, config=config):
                    if (
                        getattr(self, "_cancel_event", None)
                        and self._cancel_event.is_set()
                    ):
                        self._emit(
                            "warning", "[System] Task interrupted by user (Ctrl+C)."
                        )
                        break
            except Exception as e:
                import traceback

                traceback.print_exc()
                graph_error[0] = e
            finally:
                graph_done.set()

        thread = threading.Thread(target=_run_graph, daemon=True)
        thread.start()

        while not graph_done.is_set() or self._event_buffer:
            if (
                getattr(self, "_cancel_event", None)
                and self._cancel_event.is_set()
                and not graph_done.is_set()
            ):
                break
            if self._event_buffer:
                yield self._event_buffer.pop(0)
            else:
                time.sleep(0.01)

        if graph_error[0]:
            self._emit("error", f"[System Error] {str(graph_error[0])}")

        while self._event_buffer:
            yield self._event_buffer.pop(0)

        if getattr(self, "_cancel_event", None) and self._cancel_event.is_set():
            yield {"type": "warning", "content": "\n[系统] 已终止该轮次操作"}

            current_state = self.app.get_state(config)
            values = current_state.values if current_state else {}

            prompt_tokens = values.get("prompt_tokens", 0)
            if prompt_tokens == 0:
                messages = values.get("messages", [])
                if messages and hasattr(self, "token_counter"):
                    prompt_tokens = self.token_counter.count_message_tokens(messages)

            total_tokens = values.get("total_tokens", 0)
            if total_tokens == 0:
                total_tokens = prompt_tokens

            yield {
                "type": "done",
                "content": "cancelled",
                "iteration_count": values.get("iteration_count", 0),
                "task_completed": False,
                "total_tokens": total_tokens,
                "prompt_tokens": prompt_tokens,
                "completion_tokens": values.get("completion_tokens", 0),
            }
            return

        final_state = self.app.get_state(config)
        values = final_state.values if final_state else {}
        final_completed = bool(values.get("task_completed", False))
        needs_user_confirmation = bool(values.get("needs_user_confirmation", False))
        if (
            not final_completed
            and values.get("ended_without_tools", False)
            and not needs_user_confirmation
        ):
            final_completed = True
        yield {
            "type": "done",
            "content": "",
            "iteration_count": values.get("iteration_count", 0),
            "task_completed": final_completed,
            "total_tokens": values.get("total_tokens", 0),
            "prompt_tokens": values.get("prompt_tokens", 0),
            "completion_tokens": values.get("completion_tokens", 0),
        }


class KlynxAgent(KlynxGeneralAgent):
    """Public default Klynx agent."""

    pass
