"""Tool call loop — core agent conversation loop.

Implements the cycle: user → model → [tool calls → model]* → final response.
Tool call errors are returned to the model (not raised), per the project philosophy.
The runtime only intervenes for safety limits and infrastructure failures.
"""

from __future__ import annotations

import asyncio
import time
from dataclasses import dataclass, field

from hatchdx import HdxError
from hatchdx.agent.providers.base import AgentResponse, ModelProvider, ToolCall, ToolResult
from hatchdx.agent.runtime.servers import ManagedServer, ServerManager
from hatchdx.agent.runtime.tools import CollectedTools, ToolRoute


# ---------------------------------------------------------------------------
# Errors
# ---------------------------------------------------------------------------


class ToolCallLoopError(HdxError):
    """Exceeded max_tool_calls limit or other loop-level failure."""


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------


@dataclass
class ToolCallEvent:
    """A single tool call execution record."""

    tool_name: str
    server_name: str
    arguments: dict
    result: str
    is_error: bool
    latency_ms: float
    was_retry: bool = False


@dataclass
class LoopTurn:
    """One model turn in the conversation loop."""

    response: AgentResponse
    tool_events: list[ToolCallEvent] = field(default_factory=list)


@dataclass
class LoopResult:
    """Complete result of running the agent loop."""

    turns: list[LoopTurn] = field(default_factory=list)
    final_text: str | None = None
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    total_tool_calls: int = 0

    @property
    def all_tool_events(self) -> list[ToolCallEvent]:
        events = []
        for turn in self.turns:
            events.extend(turn.tool_events)
        return events


# ---------------------------------------------------------------------------
# Timeout parsing
# ---------------------------------------------------------------------------


def _parse_timeout(timeout_str: str) -> float:
    """Parse a timeout string like '30s', '5m', '1.5s' into seconds."""
    s = timeout_str.strip().lower()
    if s.endswith("ms"):
        return float(s[:-2]) / 1000.0
    elif s.endswith("m"):
        return float(s[:-1]) * 60.0
    elif s.endswith("s"):
        return float(s[:-1])
    else:
        return float(s)


# ---------------------------------------------------------------------------
# Tool call execution
# ---------------------------------------------------------------------------


async def _execute_tool_call(
    tool_call: ToolCall,
    route: ToolRoute,
    server_manager: ServerManager,
    timeout: float,
) -> ToolCallEvent:
    """Execute a single tool call against the correct server.

    Returns a ToolCallEvent (never raises — errors are captured in the event).
    """
    start = time.monotonic()

    # Find the managed server
    server: ManagedServer | None = None
    for s in server_manager.servers:
        if s.config.name == route.server_name:
            server = s
            break

    if server is None or server.simulator is None:
        elapsed = (time.monotonic() - start) * 1000
        return ToolCallEvent(
            tool_name=route.exposed_name,
            server_name=route.server_name,
            arguments=tool_call.arguments,
            result=f"Server '{route.server_name}' is not running",
            is_error=True,
            latency_ms=elapsed,
        )

    try:
        result = await asyncio.wait_for(
            server.simulator.call_tool(route.original_name, tool_call.arguments),
            timeout=timeout,
        )

        elapsed = (time.monotonic() - start) * 1000

        # Extract text content from MCP result
        content_parts = []
        for block in result.content:
            if isinstance(block, dict):
                if block.get("type") == "text":
                    content_parts.append(block.get("text", ""))
                else:
                    content_parts.append(str(block))
            else:
                content_parts.append(str(block))

        content = "\n".join(content_parts) if content_parts else "(empty result)"

        return ToolCallEvent(
            tool_name=route.exposed_name,
            server_name=route.server_name,
            arguments=tool_call.arguments,
            result=content,
            is_error=result.is_error,
            latency_ms=elapsed,
        )

    except asyncio.TimeoutError:
        elapsed = (time.monotonic() - start) * 1000
        return ToolCallEvent(
            tool_name=route.exposed_name,
            server_name=route.server_name,
            arguments=tool_call.arguments,
            result=f"Tool call timed out after {timeout:.0f}s",
            is_error=True,
            latency_ms=elapsed,
        )
    except Exception as e:
        elapsed = (time.monotonic() - start) * 1000
        return ToolCallEvent(
            tool_name=route.exposed_name,
            server_name=route.server_name,
            arguments=tool_call.arguments,
            result=f"Tool call failed: {e}",
            is_error=True,
            latency_ms=elapsed,
        )


# ---------------------------------------------------------------------------
# Core loop
# ---------------------------------------------------------------------------


async def run_agent_loop(
    provider: ModelProvider,
    server_manager: ServerManager,
    collected_tools: CollectedTools,
    system_prompt: str,
    user_message: str,
    *,
    max_tokens: int = 4096,
    temperature: float = 0.0,
    max_tool_calls: int = 50,
    tool_call_timeout: str = "30s",
    retry_on_error: bool = True,
    max_retries: int = 1,
    on_turn: object | None = None,
) -> LoopResult:
    """Run the core agent conversation loop.

    Args:
        provider: The model provider to use.
        server_manager: Active server manager with running servers.
        collected_tools: Processed tools with routing table.
        system_prompt: The agent's system prompt.
        user_message: The user's input message.
        max_tokens: Max tokens per model response.
        temperature: Model temperature.
        max_tool_calls: Safety limit on total tool calls.
        tool_call_timeout: Per-call timeout string (e.g., "30s").
        retry_on_error: Whether to retry failed tool calls.
        max_retries: How many times to retry.
        on_turn: Optional async callback(LoopTurn) called after each model turn.

    Returns:
        LoopResult with all turns, final text, and aggregate stats.

    Raises:
        ToolCallLoopError: If max_tool_calls is exceeded.
    """
    timeout_secs = _parse_timeout(tool_call_timeout)

    # Build tool definitions in provider format
    tool_defs = [
        {
            "name": t.name,
            "description": t.description,
            "inputSchema": t.input_schema,
        }
        for t in collected_tools.tools
    ]

    messages: list[dict] = [{"role": "user", "content": user_message}]
    loop_result = LoopResult()
    tool_call_count = 0

    while True:
        # Call the model
        response = await provider.chat(
            messages=messages,
            tools=tool_defs,
            system=system_prompt,
            max_tokens=max_tokens,
            temperature=temperature,
        )

        # Track usage
        loop_result.total_input_tokens += response.usage.get("input", 0)
        loop_result.total_output_tokens += response.usage.get("output", 0)

        turn = LoopTurn(response=response)

        # If no tool calls, we're done
        if not response.tool_calls:
            loop_result.final_text = response.text
            loop_result.turns.append(turn)
            if on_turn and asyncio.iscoroutinefunction(on_turn):
                await on_turn(turn)
            break

        # Add assistant message to conversation
        messages.append({
            "role": "assistant",
            "content": response.text,
            "tool_calls": response.tool_calls,
        })

        # Execute each tool call
        for tool_call in response.tool_calls:
            tool_call_count += 1
            if tool_call_count > max_tool_calls:
                loop_result.turns.append(turn)
                loop_result.total_tool_calls = tool_call_count - 1
                raise ToolCallLoopError(
                    f"Exceeded max tool calls ({max_tool_calls}). "
                    "Increase settings.max_tool_calls in agent.yaml if needed."
                )

            route = collected_tools.get_route(tool_call.name)
            if route is None:
                # Model hallucinated a tool name
                event = ToolCallEvent(
                    tool_name=tool_call.name,
                    server_name="unknown",
                    arguments=tool_call.arguments,
                    result=f"Unknown tool: '{tool_call.name}'",
                    is_error=True,
                    latency_ms=0,
                )
                turn.tool_events.append(event)
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "content": event.result,
                    "is_error": True,
                })
                continue

            # Execute with retry
            event = await _execute_tool_call(
                tool_call, route, server_manager, timeout_secs
            )

            if event.is_error and retry_on_error:
                for attempt in range(max_retries):
                    retry_event = await _execute_tool_call(
                        tool_call, route, server_manager, timeout_secs
                    )
                    retry_event.was_retry = True
                    if not retry_event.is_error:
                        event = retry_event
                        break
                    event = retry_event

            turn.tool_events.append(event)

            # Add tool result to messages
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call.id,
                "content": event.result,
                "is_error": event.is_error,
            })

        loop_result.turns.append(turn)
        loop_result.total_tool_calls = tool_call_count

        if on_turn and asyncio.iscoroutinefunction(on_turn):
            await on_turn(turn)

        # Loop back — model sees tool results and decides next action

    return loop_result
