"""
Agent runner — sends completion requests to the LLM pool service.

Communication pattern:
  orchestrator.tasks.handle_message
      → run_agent (this module, async)
          → POST http://llm-pool:8004/complete  (httpx async)
              → LiteLLM Router (router.py in llm-pool)
                  → Ollama / Anthropic / OpenAI

Tool-call loop (Phase 2):
  After each LLM response, check for tool_calls in the response.
  If tool_calls present:
    - Execute each tool via execute_tool()
    - If requires_confirmation: stop loop, return confirmation message to user
    - Otherwise: append tool result as 'tool' role message, re-call LLM
  Loop until LLM returns plain text (no tool_calls) or max_iterations reached.
  Max iterations = 5 (prevents runaway tool chains).

Streaming (web channel only):
  run_agent_streaming() calls POST /complete/stream and yields token strings.
  Tool calls are NOT supported in the streaming path — only used for the final
  text response after the tool-call loop has resolved (or when no tools needed).
"""

from __future__ import annotations

import json
import logging
import time
import uuid
from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING, Any

import httpx

from orchestrator.agents.builder import build_messages, build_system_prompt
from shared.config import settings
from shared.models.message import KonstructMessage
from shared.models.tenant import Agent

if TYPE_CHECKING:
    from orchestrator.audit.logger import AuditLogger
    from orchestrator.tools.registry import ToolDefinition

logger = logging.getLogger(__name__)

_FALLBACK_RESPONSE = (
    "I'm having trouble processing your request right now. "
    "Please try again in a moment."
)

# Timeout for LLM pool HTTP requests — generous to allow slow local inference
_LLM_TIMEOUT = httpx.Timeout(timeout=120.0, connect=10.0)
# Streaming needs a longer read timeout — first token can take 30-60s with cloud models
_LLM_STREAM_TIMEOUT = httpx.Timeout(timeout=300.0, connect=10.0, read=300.0)

# Maximum number of tool-call iterations before breaking the loop
_MAX_TOOL_ITERATIONS = 5


async def run_agent(
    msg: KonstructMessage,
    agent: Agent,
    messages: list[dict] | None = None,
    audit_logger: "AuditLogger | None" = None,
    tool_registry: "dict[str, ToolDefinition] | None" = None,
) -> str:
    """
    Execute an agent against the LLM pool and return the response text.

    Implements a multi-turn tool-call loop:
      1. Call LLM with messages (and tool definitions if registry provided)
      2. If LLM returns tool_calls:
         a. Execute each tool via execute_tool()
         b. If tool requires confirmation: return confirmation message immediately
         c. Append tool results as 'tool' role messages
         d. Re-call LLM with updated messages
      3. Repeat until LLM returns plain text or _MAX_TOOL_ITERATIONS reached

    Every LLM call and tool invocation is logged to the audit trail if
    audit_logger is provided.

    Args:
        msg:           The inbound Konstruct message being processed.
        agent:         The ORM Agent instance that handles this message.
        messages:      Optional pre-built messages array (e.g. from
                       build_messages_with_memory). When provided, used directly.
                       When None, falls back to simple [system, user] construction.
        audit_logger:  Optional AuditLogger for recording each LLM call and tool
                       invocation. When None, no audit logging occurs (backward compat).
        tool_registry: Optional dict of tool name → ToolDefinition for this agent.
                       When provided, passed to LLM as function-calling tools.

    Returns:
        The LLM response content as a plain string.
        Returns a polite fallback message if the LLM pool is unreachable or
        returns a non-200 response.
        Returns a confirmation message if a tool with requires_confirmation=True
        was invoked — the caller should return this to the user and store the
        pending action.
    """
    if messages is None:
        # Fallback: simple two-message construction (backward compat)
        system_prompt = build_system_prompt(agent)
        user_text: str = msg.content.text or ""
        messages = build_messages(
            system_prompt=system_prompt,
            user_message=user_text,
        )

    # Build tool definitions for LiteLLM if a registry was provided
    tools_payload: list[dict] | None = None
    if tool_registry:
        from orchestrator.tools.registry import to_litellm_format
        tools_payload = to_litellm_format(tool_registry)

    # Mutable copy of messages for the tool loop
    loop_messages: list[dict[str, Any]] = list(messages)

    tenant_id: uuid.UUID | None = None
    if msg.tenant_id:
        try:
            tenant_id = uuid.UUID(str(msg.tenant_id))
        except ValueError:
            pass

    agent_uuid = agent.id if isinstance(agent.id, uuid.UUID) else uuid.UUID(str(agent.id))
    user_id: str = (
        msg.sender.user_id
        if msg.sender and msg.sender.user_id
        else (msg.thread_id or msg.id)
    )

    llm_pool_url = f"{settings.llm_pool_url}/complete"

    for iteration in range(_MAX_TOOL_ITERATIONS + 1):
        if iteration == _MAX_TOOL_ITERATIONS:
            logger.warning(
                "Agent %s reached max tool iterations (%d) for tenant=%s — stopping loop",
                agent.id,
                _MAX_TOOL_ITERATIONS,
                msg.tenant_id,
            )
            return _FALLBACK_RESPONSE

        # ------------------------------------------------------------------
        # Call LLM pool
        # ------------------------------------------------------------------
        payload: dict[str, Any] = {
            "model": agent.model_preference,
            "messages": loop_messages,
            "tenant_id": str(msg.tenant_id) if msg.tenant_id else "",
        }
        if tools_payload:
            payload["tools"] = tools_payload

        call_start = time.monotonic()
        async with httpx.AsyncClient(timeout=_LLM_TIMEOUT) as client:
            try:
                response = await client.post(llm_pool_url, json=payload)
            except httpx.RequestError:
                logger.exception(
                    "LLM pool unreachable for tenant=%s agent=%s url=%s",
                    msg.tenant_id,
                    agent.id,
                    llm_pool_url,
                )
                return _FALLBACK_RESPONSE

        if response.status_code != 200:
            logger.error(
                "LLM pool returned %d for tenant=%s agent=%s",
                response.status_code,
                msg.tenant_id,
                agent.id,
            )
            return _FALLBACK_RESPONSE

        call_latency_ms = int((time.monotonic() - call_start) * 1000)
        data = response.json()

        response_content: str = data.get("content", "") or ""
        response_tool_calls: list[dict] = data.get("tool_calls", []) or []

        # ------------------------------------------------------------------
        # Log LLM call to audit trail
        # ------------------------------------------------------------------
        if audit_logger and tenant_id:
            # Summarize input as last user message
            input_summary = _get_last_user_message(loop_messages)
            output_summary = response_content or f"[{len(response_tool_calls)} tool calls]"
            try:
                # Extract token usage from LLM pool response
                usage_data = data.get("usage", {}) or {}
                prompt_tokens = int(usage_data.get("prompt_tokens", 0))
                completion_tokens = int(usage_data.get("completion_tokens", 0))
                total_tokens = prompt_tokens + completion_tokens

                # Extract cost from response or estimate from usage
                cost_usd = float(data.get("cost_usd", 0.0))

                # Extract provider from model string (e.g. "anthropic/claude-sonnet-4" → "anthropic")
                model_str: str = data.get("model", agent.model_preference) or ""
                provider = model_str.split("/")[0] if "/" in model_str else model_str

                await audit_logger.log_llm_call(
                    tenant_id=tenant_id,
                    agent_id=agent_uuid,
                    user_id=user_id,
                    input_summary=input_summary,
                    output_summary=output_summary,
                    latency_ms=call_latency_ms,
                    metadata={
                        "model": model_str,
                        "provider": provider,
                        "iteration": iteration,
                        "tool_calls_count": len(response_tool_calls),
                        "prompt_tokens": prompt_tokens,
                        "completion_tokens": completion_tokens,
                        "total_tokens": total_tokens,
                        "cost_usd": cost_usd,
                    },
                )
            except Exception:
                logger.exception("Failed to log LLM call to audit trail")

        # ------------------------------------------------------------------
        # No tool calls — LLM returned plain text, we're done
        # ------------------------------------------------------------------
        if not response_tool_calls:
            return response_content or _FALLBACK_RESPONSE

        # ------------------------------------------------------------------
        # Process tool calls
        # ------------------------------------------------------------------
        if not tool_registry or not audit_logger or tenant_id is None:
            # No tool registry provided — cannot execute tools
            # Return content if available, or fallback
            return response_content or _FALLBACK_RESPONSE

        from orchestrator.tools.executor import execute_tool

        # Append assistant's tool-call message to the loop
        loop_messages.append({
            "role": "assistant",
            "content": response_content or None,
            "tool_calls": response_tool_calls,
        })

        for tool_call in response_tool_calls:
            tool_result = await execute_tool(
                tool_call=tool_call,
                registry=tool_registry,
                tenant_id=tenant_id,
                agent_id=agent_uuid,
                audit_logger=audit_logger,
            )

            # Check if this is a confirmation request (requires_confirmation=True)
            # The confirmation message template starts with "This action requires"
            if tool_result.startswith("This action requires your approval"):
                # Return confirmation message to user — stop the loop
                return tool_result

            # Append tool result as 'tool' role message for re-injection into LLM
            tool_call_id = tool_call.get("id", "call_0")
            loop_messages.append({
                "role": "tool",
                "tool_call_id": tool_call_id,
                "content": tool_result,
            })

        # Continue loop — re-call LLM with tool results appended

    # Should never reach here (loop guard above), but satisfy type checker
    return _FALLBACK_RESPONSE


async def run_agent_streaming(
    msg: KonstructMessage,
    agent: Agent,
    messages: list[dict] | None = None,
) -> AsyncGenerator[str, None]:
    """
    Stream the final LLM response for an agent, yielding token strings.

    This is the web-channel-only streaming path. It does NOT support tool calls —
    tool resolution must be completed BEFORE calling this function (use run_agent()
    for the tool-call loop, then only call this for streaming plain text responses).

    The caller (tasks._process_message) is responsible for:
      - Running the tool-call loop via run_agent() first if tools are registered
      - Only falling into this path when the response will be plain text
      - Publishing each yielded token to Redis pub-sub

    For simple conversations (no tools), this streams the response directly.

    Args:
        msg:      The inbound KonstructMessage.
        agent:    The ORM Agent instance.
        messages: Memory-enriched messages array (from build_messages_with_memory).
                  When None, falls back to simple [system, user] construction.

    Yields:
        Token strings as generated by the LLM.
        On error, yields the fallback response string as a single chunk.
    """
    if messages is None:
        from orchestrator.agents.builder import build_messages, build_system_prompt
        system_prompt = build_system_prompt(agent)
        user_text: str = msg.content.text or ""
        messages = build_messages(
            system_prompt=system_prompt,
            user_message=user_text,
        )

    llm_stream_url = f"{settings.llm_pool_url}/complete/stream"
    payload: dict[str, Any] = {
        "model": agent.model_preference,
        "messages": messages,
        "tenant_id": str(msg.tenant_id) if msg.tenant_id else "",
    }

    try:
        async with httpx.AsyncClient(timeout=_LLM_STREAM_TIMEOUT) as client:
            async with client.stream("POST", llm_stream_url, json=payload) as response:
                if response.status_code != 200:
                    logger.error(
                        "LLM pool stream returned %d for tenant=%s agent=%s",
                        response.status_code,
                        msg.tenant_id,
                        agent.id,
                    )
                    yield _FALLBACK_RESPONSE
                    return

                async for line in response.aiter_lines():
                    if not line.strip():
                        continue
                    try:
                        event = json.loads(line)
                    except json.JSONDecodeError:
                        continue

                    event_type = event.get("type")
                    if event_type == "chunk":
                        token = event.get("text", "")
                        if token:
                            yield token
                    elif event_type == "done":
                        return
                    elif event_type == "error":
                        logger.error(
                            "LLM pool stream error for tenant=%s agent=%s: %s",
                            msg.tenant_id,
                            agent.id,
                            event.get("message", "unknown"),
                        )
                        yield _FALLBACK_RESPONSE
                        return

    except httpx.RequestError:
        logger.exception(
            "LLM pool stream unreachable for tenant=%s agent=%s url=%s",
            msg.tenant_id,
            agent.id,
            llm_stream_url,
        )
        yield _FALLBACK_RESPONSE


def _get_last_user_message(messages: list[dict[str, Any]]) -> str:
    """Extract the content of the last user message for audit summary."""
    for msg in reversed(messages):
        if msg.get("role") == "user":
            content = msg.get("content", "")
            if isinstance(content, str):
                return content[:200]
            elif isinstance(content, list):
                # Multi-modal content
                for part in content:
                    if isinstance(part, dict) and part.get("type") == "text":
                        return str(part.get("text", ""))[:200]
    return "[no user message]"
