"""
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>

  <workflow>
    <step number="1">
      <name>判断任务目标是否需要更新</name>
      <details>
        - 如果工具执行结果带来新的关键信息，导致任务目标需要调整
        - 调用 `update_task_state` 同步更新你的最新动作或总目标
      </details>
    </step>
    <step number="2">
      <name>分析和决策</name>
      <details>
        - 先理解：阅读相关文件，检查现有模式。
        - 做到快速且深入。收集足够的证据后开始，然后持续迭代。
        - 分析上一步工具执行结果
        - 判断任务是否已完成
        - 如果未完成，明确下一步要执行的操作
      </details>
    </step>
    <step number="3">
      <name>执行操作</name>
      <details>
        - 实施解决方案。保证速度的同时确保准确。
        - 需要工具时，使用 XML 模式或 Native Function Calling 模式调用工具
        - 任务完成时，使用 `attempt_completion` 输出最终结果
      </details>
    </step>

    <step number="4">
      <name>验证</name>
      <details>
        - 根据用户的需求（而非你自己的输出）来检查工作。
        - 初次尝试很少能完全正确，请不断迭代
      </details>
    </step>
    <step number="5">
      <name>错误反思</name>
      <details>
        - 如果上一步工具返回错误（如 `<error>` 标签），必须先用 `<reflection>` 分析失败原因
        - 禁止不经反思就重复同样操作
      </details>
    </step>
    持续工作直到任务完全完成。不要中途停止并解释你“将要”做什么 —— 直接完成它。只有在任务达成或确实被卡住时，才将控制权交还给用户。
  </workflow>

  <instructions>
    阅读 markdown 文件时，优先识别标题、总行数和总字数，再决定是否读取全文。
    处理歧义：如果请求不明确，在行动前先提问，不要擅自操作。
    技能优先按需加载：当任务需要已安装技能包时，先调用 `load_skill` 读取对应 SKILL.md，再按其流程执行。
    遇到阻碍时先确认：如果因登录/认证/权限等原因无法按原计划执行，且需要切换方案（例如改用其他工具），必须先询问用户，不要自行决定。
    简洁直切主题：任务目标明确时，直接围绕目标执行，避免无关动作。除非被要求，否则不要过度解释。
    避免重复：除非用户明确要求，不要重复执行已经执行过的操作，也不要重复输出已展示的内容。
    事后检查：每次完成操作后，检查是否有遗漏。
    严禁无谓的开场白：永远不要添加不必要的寒暄（如“好的！”、“问得好！”、“我现在将……”）。
    行动胜于言语：不要说“我现在要做 xxx”，而是直接去做。
    先解释后执行：如果被问及如何处理某事，先解释思路，然后再执行。
    准确性优先：比起认可用户的观点，优先保证事实准确。
    礼貌地反驳：当用户观点错误时，请委婉且尊重地予以纠正。
    保持中性：避免使用不必要的最高级赞美、夸奖或情感共鸣。
  </instructions>

  <when_problems_occur>
    如果某项操作反复失败，停止并分析原因 —— 不要盲目重复同样的方法。
    如果遇到了阻碍，告知用户问题所在并寻求指导。
  </when_problems_occur>
</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 -> parser
        -> tools/agent/verify_completion(条件分支) -> feedback -> 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("parser", self._tool_parser_node)
        workflow.add_node("tools", self._tool_executor_node)
        workflow.add_node("feedback", self._feedback_node)
        workflow.add_node("verify_completion", self._verify_completion_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", "parser")
        workflow.add_conditional_edges(
            "parser",
            self._should_continue,
            {
                "tools": "tools",
                "agent": "agent",
                "verify": "verify_completion",
                "end_direct": END,
            },
        )
        workflow.add_edge("tools", "feedback")
        workflow.add_conditional_edges(
            "feedback",
            self._should_continue,
            {
                "agent": "agent",
                "verify": "verify_completion",
                "end_direct": END,
            },
        )
        workflow.add_conditional_edges(
            "verify_completion",
            self._route_after_verify_completion,
            {
                "agent": "agent",
                "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,
            "pending_tool_calls": [],
            "iteration_count": 0,
            "working_dir": self.working_dir,
            "task_completed": False,
            "last_action": "",
            "empty_response_count": 0,
            "total_tokens": 0,
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "max_context_tokens": self.max_context_tokens,
            "task_type": "",
            "user_input": task,
            "summary_content": "",
            "ended_without_tools": False,
            "completion_candidate": False,
            "completion_candidate_result": "",
            "verification_decision": "",
            "needs_user_confirmation": False,
            "clarification_question": "",
            "thinking_context": thinking_context,
            "system_prompt_append": (system_prompt_append or "").strip(),
        }

        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:
            verify_decision = (values.get("verification_decision", "") or "").strip().lower()
            reached_limit = values.get("iteration_count", 0) >= self.max_iterations
            # 无工具直接答复仅在“非重试/非澄清/未触发迭代上限”时视为完成
            if verify_decision not in ("retry", "clarify") and not reached_limit:
                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
