"""Agentic loop engine.

Implements the multi-turn tool-calling loop:
1. Build prompts/messages
2. Call LLM adapter with discovered tools
3. If tool calls are returned, execute MCP tools (validated, timed, retried)
4. Append tool results back to messages via adapter
5. Repeat until done or max-iteration safety limits are reached
"""

from __future__ import annotations

import asyncio
import json
import logging
import uuid
from dataclasses import dataclass
from typing import Any

from mcpbridge.adapters.base import BaseLLMAdapter, ToolCall, ToolResult
from mcpbridge.exceptions import (
    MCPProtocolError,
    MCPToolCallError,
    MCPToolNotFoundError,
    MCPSchemaValidationError,
    LoopTimeoutError,
)
from mcpbridge.registry import DynamicToolRegistry
from mcpbridge.transport.base import BaseTransport
from mcpbridge.utils.retry import async_retry
from mcpbridge.utils.schema_validator import validate_tool_args

logger = logging.getLogger("mcpbridge")


@dataclass(slots=True)
class LoopResult:
    """Result from a full agentic loop run."""

    text: str
    tool_calls_made: list[dict[str, Any]]
    iterations: int
    total_input_tokens: int
    total_output_tokens: int
    finish_reason: str  # "done" | "max_iterations" | "error"


async def run_loop(
    adapter: BaseLLMAdapter,
    prompt_builder: Any,
    registry: DynamicToolRegistry,
    transport: BaseTransport | dict[str, BaseTransport],
    user_query: str,
    history: list[dict[str, Any]],
    loop_config: Any,
) -> LoopResult:
    """Run the loop using the adapter + transport + tool registry."""
    if history is None:
        history = []
    # Keep a reference so we can persist the final provider-specific messages
    # back into the caller's history list.
    history_ref = history

    iteration = 0
    tool_calls_log: list[dict[str, Any]] = []
    total_input_tokens = 0
    total_output_tokens = 0

    while True:
        system_prompt, messages = prompt_builder.build_messages(
            history=history_ref,
            user_query=user_query,
        )

        provider_name = adapter.provider_name
        tools = registry.for_llm(provider_name)

        # Token accounting (best-effort).
        try:
            total_input_tokens += adapter.count_tokens(messages, system_prompt or "")
            if total_input_tokens > int(loop_config.max_tokens_total):
                raise LoopTimeoutError("Exceeded max_tokens_total in loop")
        except Exception:
            # Never break the loop due to accounting issues.
            pass

        response = await adapter.chat(messages, tools, system=system_prompt)

        if adapter.is_done(response):
            text = adapter.extract_text(response) or ""
            # Ensure assistant final message is appended to history using the adapter's
            # provider-specific message format.
            final_messages = adapter.append_tool_results(messages, response, [])
            history_ref[:] = final_messages
            return LoopResult(
                text=text,
                tool_calls_made=tool_calls_log,
                iterations=iteration,
                total_input_tokens=total_input_tokens,
                total_output_tokens=total_output_tokens + _rough_tokens(text),
                finish_reason="done",
            )

        tool_calls = adapter.extract_tool_calls(response)

        if loop_config.on_tool_call:
            for tc in tool_calls:
                await _maybe_await(loop_config.on_tool_call(tc.name, tc.arguments))

        # Execute tool calls.
        if bool(loop_config.parallel_tool_calls) and len(tool_calls) > 1:
            results = await asyncio.gather(
                *[
                    execute_tool(
                        tool_call=tc,
                        transport=transport,
                        registry=registry,
                        loop_config=loop_config,
                    )
                    for tc in tool_calls
                ],
                return_exceptions=False,
            )
        else:
            results = []
            for tc in tool_calls:
                results.append(
                    await execute_tool(
                        tool_call=tc,
                        transport=transport,
                        registry=registry,
                        loop_config=loop_config,
                    )
                )

        for tc, tr in zip(tool_calls, results):
            tool_calls_log.append(
                {
                    "id": tc.id,
                    "name": tc.name,
                    "arguments": tc.arguments,
                    "result": tr.content,
                    "is_error": tr.is_error,
                }
            )
            if loop_config.on_tool_result:
                await _maybe_await(loop_config.on_tool_result(tc.name, tr))

        # Append tool results via adapter, so adapter can shape messages accordingly.
        messages = adapter.append_tool_results(messages, response, results)

        # Update shared history for next iteration (in-place).
        history_ref[:] = messages
        iteration += 1

        if iteration >= int(loop_config.max_iterations):
            # Safety behavior: a provider model can get stuck repeatedly
            # emitting tool calls instead of producing a final assistant text.
            # Rather than failing the whole service, return the best-effort
            # payload so upstream can decide what to do.
            text = ""
            try:
                text = adapter.extract_text(response) or ""
            except Exception:
                text = ""

            # If the model produced no final text, fall back to the most
            # recent tool result content (common when the prompt expects
            # JSON-only tool output).
            if not text and tool_calls_log:
                last = tool_calls_log[-1]
                if isinstance(last.get("result"), str):
                    text = last["result"]

            return LoopResult(
                text=text,
                tool_calls_made=tool_calls_log,
                iterations=iteration,
                total_input_tokens=total_input_tokens,
                total_output_tokens=total_output_tokens + _rough_tokens(text),
                finish_reason="max_iterations",
            )

        if loop_config.on_iteration:
            await _maybe_await(loop_config.on_iteration(iteration, messages))


async def execute_tool(
    tool_call: ToolCall,
    transport: BaseTransport | dict[str, BaseTransport],
    registry: DynamicToolRegistry,
    loop_config: Any,
) -> ToolResult:
    """Execute a single MCP tool call with validation/timeouts/retries."""

    async def _do_call() -> ToolResult:
        try:
            tool = registry.get_tool(tool_call.name)
        except Exception:
            err_msg = f"Tool `{tool_call.name}` not found in registry"
            if loop_config.error_strategy == "raise":
                raise MCPToolNotFoundError(err_msg)
            return ToolResult(
                tool_call_id=tool_call.id,
                name=tool_call.name,
                content=err_msg,
                is_error=True,
            )

        valid, schema_err = validate_tool_args(tool_call.arguments, tool.input_schema)
        if not valid:
            if loop_config.error_strategy == "raise":
                raise MCPSchemaValidationError(schema_err or "Invalid tool arguments")
            return ToolResult(
                tool_call_id=tool_call.id,
                name=tool_call.name,
                content=f"Tool argument schema validation failed: {schema_err}",
                is_error=True,
            )

        params = {"name": tool.name, "arguments": tool_call.arguments}
        request = {"jsonrpc": "2.0", "id": str(uuid.uuid4()), "method": "tools/call", "params": params}

        transport_for_tool: BaseTransport
        if isinstance(transport, dict):
            transport_for_tool = transport.get(tool.server_name)  # type: ignore[assignment]
            if transport_for_tool is None:
                err_msg = f"No transport for MCP server `{tool.server_name}`"
                if loop_config.error_strategy == "raise":
                    raise MCPToolCallError(err_msg)
                return ToolResult(tool_call_id=tool_call.id, name=tool_call.name, content=err_msg, is_error=True)
        else:
            transport_for_tool = transport

        try:
            resp = await asyncio.wait_for(
                transport_for_tool.send(request),
                timeout=int(loop_config.tool_timeout),
            )
        except asyncio.TimeoutError:
            err_msg = f"Tool `{tool_call.name}` timed out after {int(loop_config.tool_timeout)}s"
            if loop_config.error_strategy == "raise":
                raise TimeoutError(err_msg)
            return ToolResult(
                tool_call_id=tool_call.id,
                name=tool_call.name,
                content=err_msg,
                is_error=True,
            )

        if not isinstance(resp, dict):
            err_msg = f"Invalid MCP response for tool `{tool_call.name}`"
            if loop_config.error_strategy == "raise":
                raise MCPProtocolError(err_msg)
            return ToolResult(tool_call_id=tool_call.id, name=tool_call.name, content=err_msg, is_error=True)

        if resp.get("error"):
            err_msg = f"MCP tools/call error for `{tool_call.name}`: {resp.get('error')}"
            if loop_config.error_strategy == "raise":
                raise MCPToolCallError(err_msg)
            return ToolResult(tool_call_id=tool_call.id, name=tool_call.name, content=err_msg, is_error=True)

        result = resp.get("result") or {}
        contents = result.get("content") or []
        is_error = bool(result.get("isError") or result.get("is_error") or result.get("error"))

        texts: list[str] = []
        for block in contents:
            if not isinstance(block, dict):
                continue
            b_type = block.get("type")
            if b_type == "text":
                # MCP text block
                text = block.get("text")
                if text is None:
                    text = block.get("content")
                if isinstance(text, str):
                    texts.append(text)
            elif b_type == "image":
                # Convert image content to a text descriptor.
                is_error = is_error or bool(block.get("isError"))
                texts.append("[Image content received]")
                logger.warning("Tool `%s` returned image; sending a text placeholder.", tool_call.name)
            else:
                if "text" in block and isinstance(block.get("text"), str):
                    texts.append(block["text"])

        content_str = "\n".join([t for t in texts if t]).strip()
        return ToolResult(tool_call_id=tool_call.id, name=tool_call.name, content=content_str, is_error=is_error)

    if bool(loop_config.retry_on_tool_error):
        # Retry once for errored tool calls.
        attempt = 0
        last: ToolResult | None = None
        while attempt < 2:
            attempt += 1
            last = await _do_call()
            if not last.is_error:
                return last
            if loop_config.error_strategy == "skip":
                return last
            if attempt >= 2:
                return last
        assert last is not None
        return last

    return await _do_call()


def _rough_tokens(text: str) -> int:
    return max(0, len(text) // 4)


async def _maybe_await(x: Any) -> Any:
    if asyncio.iscoroutine(x):
        return await x
    return x

