"""Research sub-agent for Agent Johnny 5.

Spawns a focused mini agent loop to deeply investigate a question.
The sub-agent has access to search tools (MCP KB, web, shell, read_file)
but NOT write tools or research (no recursion).

Runs in the same background thread as the parent agent — blocking.
Returns findings as a formatted string for the parent to act on.
"""

from __future__ import annotations

import json
import logging
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from textual.app import App

from ..constants import AGENT_HTTP_TIMEOUT
from .openrouter import OpenRouterClient
from .prompts import RESEARCH_SYSTEM_PROMPT
from .tools import execute_tool

logger = logging.getLogger(__name__)

# Max iterations for the research sub-agent (keeps it focused)
_MAX_RESEARCH_ITERATIONS = 15

# Tools the research sub-agent can use (subset — no write, no research)
_RESEARCH_TOOL_DEFS: list[dict] = [
    {
        "type": "function",
        "function": {
            "name": "run_shell_command",
            "description": (
                "Execute a shell command on the server for investigation. "
                "Use for checking logs, service status, config files."
            ),
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "The shell command to execute",
                    },
                    "timeout": {
                        "type": "integer",
                        "description": "Timeout in seconds (default 120)",
                    },
                },
                "required": ["command"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "Read the full contents of a file on the server.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "Absolute file path to read",
                    },
                },
                "required": ["path"],
            },
        },
    },
]


def run_research(
    question: str,
    client: OpenRouterClient,
    mcp_registry: object | None,
    state: object,
    app: "App",
) -> str:
    """Run a focused research sub-agent and return findings.

    Args:
        question: The research question to investigate.
        client: The parent agent's OpenRouter client (reused).
        mcp_registry: The parent agent's MCP registry (for KB/web search).
        state: The InstallState (for tool execution context).
        app: The Textual app (for tool execution context).

    Returns:
        Formatted findings string, or error message on failure.
    """
    logger.info("Research sub-agent starting: %s", question[:100])

    # Build tool list: local read-only tools + MCP search tools
    tools = list(_RESEARCH_TOOL_DEFS)
    if mcp_registry is not None:
        # Add MCP tools that are useful for research
        _research_mcp_names = {
            "mcp_search_knowledge",
            "mcp_get_full_guide",
            "mcp_search_web",
            "mcp_diagnose_issue",
        }
        for tool_def in mcp_registry.openai_tools():
            fn_name = tool_def.get("function", {}).get("name", "")
            if fn_name in _research_mcp_names:
                tools.append(tool_def)

    # Build messages
    messages: list[dict] = [
        {"role": "system", "content": RESEARCH_SYSTEM_PROMPT},
        {"role": "user", "content": question},
    ]

    # Mini agent loop
    findings: list[str] = []

    for iteration in range(_MAX_RESEARCH_ITERATIONS):
        try:
            accumulated_text, tool_calls = _stream_response(
                client, messages, tools
            )
        except Exception as exc:
            logger.exception("Research sub-agent stream error")
            findings.append(f"[Research error: {exc}]")
            break

        if accumulated_text:
            findings.append(accumulated_text)

        if not tool_calls:
            break  # Sub-agent is done

        # Process tool calls
        assistant_msg: dict = {"role": "assistant", "tool_calls": tool_calls}
        if accumulated_text:
            assistant_msg["content"] = accumulated_text
        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 = {}

            logger.info("Research tool call: %s(%s)", fn_name, fn_args)

            # Execute the tool
            if (
                mcp_registry is not None
                and mcp_registry.is_mcp_tool(fn_name)
            ):
                result = mcp_registry.call(fn_name, fn_args)
            else:
                result = execute_tool(fn_name, fn_args, state, app)

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

    # Combine findings
    if findings:
        combined = "\n\n".join(findings)
        logger.info(
            "Research sub-agent done: %d findings, %d chars",
            len(findings), len(combined),
        )
        return combined
    else:
        return "Research sub-agent returned no findings."


def _stream_response(
    client: OpenRouterClient,
    messages: list[dict],
    tools: list[dict],
) -> tuple[str, list[dict]]:
    """Stream one API response, returning (text, tool_calls)."""
    accumulated_text = ""
    tool_calls_by_idx: dict[int, dict] = {}

    for chunk in client.stream_chat(messages, tools=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())
        ]

    return accumulated_text, tool_calls_acc
