"""AgentBrain — the synchronous agent loop with tool-call chaining.

The ENTIRE agent loop runs inside @work(thread=True) in a background
thread. Nothing here is async. All UI updates go through
app.call_from_thread().

Flow:
  1. Build system prompt with fresh log snapshot from InstallState.
  2. Send conversation to Kimi K2.5 via OpenRouter (sync streaming).
  3. Accumulate text tokens internally (no per-token UI updates).
  4. If response contains tool_calls → execute tools, append results,
     loop back to step 2 (max MAX_TOOL_ITERATIONS).
  5. When response is text-only (no tool_calls), we're done.
"""

from __future__ import annotations

import json
import logging
import threading
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from textual.app import App

    from ..state import InstallState

from ..constants import MAX_CONVERSATION_MESSAGES, MAX_TOOL_ITERATIONS, OPENROUTER_MODEL
from ..messages import AgentReply, AgentToolCall

# AgentTyping is a simple message to trigger the typing indicator
from ..messages import AgentTyping

from .openrouter import OpenRouterClient
from .prompts import build_system_prompt
from .tools import TOOL_DEFINITIONS, execute_tool

logger = logging.getLogger(__name__)


class AgentBrain:
    """Orchestrates the Kimi K2.5 agent loop (synchronous, thread-bound).

    Created once in the App. Its methods are called from the agent's
    background thread only.
    """

    def __init__(
        self, state: InstallState, api_key: str, mcp_url: str = "",
        mcp_key: str = "", model: str = OPENROUTER_MODEL,
    ) -> None:
        self.state = state
        self.api_key = api_key
        self.model = model
        self.client = OpenRouterClient(api_key, model=model)
        self.conversation: list[dict] = []
        self.mcp_tools: list[dict] = []

        # MCP knowledge base (lazy-initialized in background thread)
        self._mcp_url = mcp_url
        self._mcp_key = mcp_key
        self._mcp_initialized = False
        self._mcp_registry: object | None = None  # MCPToolRegistry

        # Track which steps the agent is already intervening on to
        # prevent cascading interventions from repeated step failures
        self._active_interventions: set[str] = set()

        # Lock protects self.conversation from concurrent access.
        # Even with exclusive @work grouping in app.py, Textual's
        # exclusive cancellation is cooperative — a thread in the
        # middle of an LLM call won't be interrupted instantly.
        self._lock = threading.Lock()

    # ------------------------------------------------------------------
    # Public entry points (called from @work(thread=True))
    # ------------------------------------------------------------------

    def ensure_mcp(self, app: App) -> None:
        """Lazy MCP initialization (must be called from background thread).

        Connects to the MCP knowledge base server and loads available
        tools. Safe to call multiple times — only runs once.
        """
        if self._mcp_initialized or not self._mcp_url:
            return
        self._mcp_initialized = True

        from ..mcp.client import MCPClient
        from ..mcp.registry import MCPToolRegistry

        try:
            client = MCPClient(self._mcp_url, api_key=self._mcp_key)
            if client.connect():
                registry = MCPToolRegistry(client)
                self.mcp_tools = registry.load()
                self._mcp_registry = registry
                count = len(self.mcp_tools)
                app.call_from_thread(
                    app.post_message,
                    AgentReply(
                        f"Connected to knowledge base — {count} tools available."
                    ),
                )
                logger.info("MCP: %d tools loaded", count)
            else:
                raise ConnectionError("client.connect() returned False")
        except Exception as exc:
            app.call_from_thread(
                app.post_message,
                AgentReply(
                    f"Could not connect to MCP knowledge base ({exc}) — "
                    "running without it."
                ),
            )
            logger.warning("MCP: connection failed: %s", exc)

    def handle_user_message(self, user_text: str, app: App) -> None:
        """Process a user chat message through the full agent loop."""
        self.ensure_mcp(app)

        with self._lock:
            user_msg = {"role": "user", "content": user_text}
            self.conversation.append(user_msg)
            self._trim_conversation()
            messages = self._build_messages()

        self._agent_loop(messages, app)

    def handle_step_failure(
        self, step_name: str, error: str, app: App
    ) -> None:
        """Proactively diagnose a failed installation step.

        Prevents cascading interventions — if the agent is already
        working on this step (or already tried and hit the iteration
        limit), it won't start a new investigation from scratch.
        Instead it gives a short status update.
        """
        self.ensure_mcp(app)

        # Prevent duplicate interventions for the same step
        if step_name in self._active_interventions:
            app.call_from_thread(
                app.post_message,
                AgentReply(
                    f"Step '{step_name}' failed again. "
                    f"I already investigated this — send me a message "
                    f"and I'll continue troubleshooting."
                ),
            )
            return

        self._active_interventions.add(step_name)

        with self._lock:
            synthetic = (
                f"[SYSTEM] Installation step '{step_name}' failed: {error}. "
                f"Diagnose the issue from the recent logs and suggest a fix."
            )
            synthetic_msg = {"role": "user", "content": synthetic}
            self.conversation.append(synthetic_msg)
            self._trim_conversation()
            messages = self._build_messages()

        self._agent_loop(messages, app)

        # Clear the intervention flag so user can manually re-trigger
        self._active_interventions.discard(step_name)

    # ------------------------------------------------------------------
    # Internal: message building
    # ------------------------------------------------------------------

    def _build_messages(self) -> list[dict]:
        """Build the full message list with a fresh system prompt."""
        # Collect MCP tool names for the prompt
        mcp_tool_names: list[str] = []
        if self._mcp_registry is not None:
            mcp_tool_names = self._mcp_registry.tool_names()

        system_prompt = build_system_prompt(
            current_step_index=self.state.current_step_index,
            total_steps=self.state.total_steps,
            current_step_name=self.state.current_step_name,
            is_running=self.state.is_running,
            is_paused=self.state.is_paused,
            is_complete=self.state.is_complete,
            has_error=self.state.has_error,
            log_snapshot=self.state.snapshot_logs(),
            mcp_tool_names=mcp_tool_names,
        )

        messages = [{"role": "system", "content": system_prompt}]
        messages.extend(self.conversation)
        return messages

    def _all_tools(self) -> list[dict]:
        """Local tools + any MCP tools registered in Phase 4."""
        return TOOL_DEFINITIONS + self.mcp_tools

    def _trim_conversation(self) -> None:
        """Cap conversation history to prevent context overflow."""
        if len(self.conversation) > MAX_CONVERSATION_MESSAGES:
            # Keep the most recent messages
            self.conversation = self.conversation[-MAX_CONVERSATION_MESSAGES:]

    # ------------------------------------------------------------------
    # Internal: the core agent loop
    # ------------------------------------------------------------------

    def _agent_loop(self, messages: list[dict], app: App) -> None:
        """Core loop: stream response, handle tool calls, repeat.

        Runs entirely in the agent's background thread.
        All UI updates go through app.call_from_thread().
        """
        for iteration in range(MAX_TOOL_ITERATIONS):
            # Show typing indicator in chat
            app.call_from_thread(
                app.post_message, AgentTyping()
            )

            try:
                accumulated_text, tool_calls = self._stream_response(
                    messages, app
                )
            except Exception as exc:
                logger.exception("Agent stream error")
                error_msg = f"[Error talking to Kimi K2.5: {exc}]"
                app.call_from_thread(
                    app.post_message, AgentReply(error_msg)
                )
                with self._lock:
                    self.conversation.append(
                        {"role": "assistant", "content": error_msg}
                    )
                return

            # Save assistant message to conversation history
            if accumulated_text and not tool_calls:
                with self._lock:
                    self.conversation.append(
                        {"role": "assistant", "content": accumulated_text}
                    )

            # No tool calls → we're done
            if not tool_calls:
                return

            # Build the assistant message with tool_calls
            assistant_msg: dict = {
                "role": "assistant",
                "tool_calls": tool_calls,
            }
            if accumulated_text:
                assistant_msg["content"] = accumulated_text

            with self._lock:
                self.conversation.append(assistant_msg)
            messages.append(assistant_msg)

            # Execute each tool call and collect results
            for tc in tool_calls:
                fn_name = tc["function"]["name"]
                try:
                    fn_args = json.loads(tc["function"]["arguments"])
                except json.JSONDecodeError:
                    fn_args = {}

                # Show tool indicator in chat
                app.call_from_thread(
                    app.post_message, AgentToolCall(fn_name, fn_args)
                )

                # Route to MCP or local executor (sync — we're in a thread)
                if (
                    self._mcp_registry is not None
                    and self._mcp_registry.is_mcp_tool(fn_name)
                ):
                    result = self._mcp_registry.call(fn_name, fn_args)
                else:
                    result = execute_tool(fn_name, fn_args, self.state, app)

                tool_result_msg = {
                    "role": "tool",
                    "tool_call_id": tc["id"],
                    "content": result,
                }
                with self._lock:
                    self.conversation.append(tool_result_msg)
                messages.append(tool_result_msg)

            # Loop continues — model sees tool results and responds

        # Safety: hit max iterations — tell the user they can continue
        app.call_from_thread(
            app.post_message,
            AgentReply(
                "I've used all my tool calls for this round. "
                "Send me a message to continue troubleshooting — "
                "I still have full context of everything I've done."
            ),
        )

    # ------------------------------------------------------------------
    # Internal: streaming response parser
    # ------------------------------------------------------------------

    def _stream_response(
        self, messages: list[dict], app: App
    ) -> tuple[str, list[dict]]:
        """Stream one API response, returning (text, tool_calls).

        Streams text tokens to the chat UI in real-time via RichLog.
        Accumulates tool_call deltas into complete tool_call objects.
        """
        accumulated_text = ""
        tool_calls_acc: list[dict] = []
        # Track tool calls by index for delta accumulation
        tool_calls_by_idx: dict[int, dict] = {}

        for chunk in self.client.stream_chat(
            messages, tools=self._all_tools()
        ):
            choices = chunk.get("choices", [])
            if not choices:
                continue

            delta = choices[0].get("delta", {})

            # --- Text content accumulation (no per-token UI updates) ---
            content = delta.get("content")
            if content:
                accumulated_text += content

            # --- Tool call delta accumulation ---
            tc_deltas = delta.get("tool_calls", [])
            for tc_delta in tc_deltas:
                idx = tc_delta.get("index", 0)

                # New tool call (has an id)
                if tc_delta.get("id"):
                    tool_calls_by_idx[idx] = {
                        "id": tc_delta["id"],
                        "type": "function",
                        "function": {
                            "name": tc_delta.get("function", {}).get(
                                "name", ""
                            ),
                            "arguments": "",
                        },
                    }

                # Append argument fragments
                arg_fragment = (
                    tc_delta.get("function", {}).get("arguments", "")
                )
                if arg_fragment and idx in tool_calls_by_idx:
                    tool_calls_by_idx[idx]["function"][
                        "arguments"
                    ] += arg_fragment

        # Flatten tool calls dict to list (ordered by index)
        if tool_calls_by_idx:
            tool_calls_acc = [
                tool_calls_by_idx[k]
                for k in sorted(tool_calls_by_idx.keys())
            ]

        # Post the final complete text to the chat UI
        if accumulated_text:
            app.call_from_thread(
                app.post_message, AgentReply(accumulated_text)
            )

        return accumulated_text, tool_calls_acc

    # ------------------------------------------------------------------
    # Cleanup
    # ------------------------------------------------------------------

    def close(self) -> None:
        """Release HTTP resources."""
        self.client.close()
        if self._mcp_registry is not None:
            self._mcp_registry.client.close()
