"""AgentBrain — conversational agent loop with tool-call chaining.

v1.2: Smarter conversational agent with research sub-agent capability.
The agent does NOT drive the install — that's handled by
install_runner.py. The agent is available during and after install for
questions, troubleshooting, and post-install setup.

New in v1.2:
  - Research sub-agent: `research` tool spawns a focused mini loop
  - Better prompt with problem-solving strategy + learning protocol
  - MiniMax 2.5 model with retry logic in OpenRouterClient

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 current install state (for context).
  2. Send conversation to LLM 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

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

logger = logging.getLogger(__name__)


class AgentBrain:
    """Orchestrates the LLM 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

        # Lock protects self.conversation from concurrent access
        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)."""
        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)

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

    def _build_messages(self) -> list[dict]:
        """Build the full message list with a fresh system 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(
            state=self.state,
            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."""
        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:
            self.conversation = self.conversation[-MAX_CONVERSATION_MESSAGES:]

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

    def _agent_loop(
        self, messages: list[dict], app: App,
        max_iterations: int | None = None,
    ) -> None:
        """Core loop: stream response, handle tool calls, repeat."""
        limit = max_iterations or MAX_TOOL_ITERATIONS
        for iteration in range(limit):
            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 LLM: {exc}]"
                app.call_from_thread(
                    app.post_message, AgentReply(error_msg)
                )
                with self._lock:
                    self.conversation.append(
                        {"role": "assistant", "content": error_msg}
                    )
                return

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

            if not tool_calls:
                return

            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)

            for tc in tool_calls:
                fn_name = tc["function"]["name"]
                try:
                    fn_args = json.loads(tc["function"]["arguments"])
                except json.JSONDecodeError:
                    fn_args = {}

                app.call_from_thread(
                    app.post_message, AgentToolCall(fn_name, fn_args)
                )

                if fn_name == "research":
                    # Research sub-agent — needs access to MCP + client
                    question = fn_args.get("question", "")
                    result = run_research(
                        question=question,
                        client=self.client,
                        mcp_registry=self._mcp_registry,
                        state=self.state,
                        app=app,
                    )
                    if not isinstance(result, str):
                        result = json.dumps(result)
                elif (
                    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)

        # Hit iteration limit
        app.call_from_thread(
            app.post_message,
            AgentReply(
                "I've used all my tool calls for this round. "
                "Send me a message to continue — "
                "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)."""
        accumulated_text = ""
        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", {})

            content = delta.get("content")
            if content:
                accumulated_text += content

            tc_deltas = delta.get("tool_calls", [])
            for tc_delta in tc_deltas:
                idx = tc_delta.get("index", 0)

                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": "",
                        },
                    }

                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

        tool_calls_acc = []
        if tool_calls_by_idx:
            tool_calls_acc = [
                tool_calls_by_idx[k]
                for k in sorted(tool_calls_by_idx.keys())
            ]

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