"""LLM provider abstraction — unified interface over LiteLLM and custom endpoints."""

from __future__ import annotations

import json
import logging
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
from typing import Any

import litellm

from agent_scout.events import Event, EventBus, EventType
from agent_scout.react_loop import Action, ActionType, Step

logger = logging.getLogger(__name__)

# Suppress litellm internal logging noise
litellm.suppress_debug_info = True


@dataclass
class LLMResponse:
    """Parsed response from an LLM call."""

    content: str = ""
    tool_calls: list[dict[str, Any]] = field(default_factory=list)
    input_tokens: int = 0
    output_tokens: int = 0
    model: str = ""
    finish_reason: str = ""


@dataclass
class CustomEndpoint:
    """User-defined OpenAI-compatible endpoint."""

    name: str
    base_url: str
    api_key: str
    model_name: str


@dataclass
class ProviderConfig:
    """Configuration for the LLM provider."""

    default_model: str = "anthropic/claude-sonnet-4-20250514"
    temperature: float = 0.3
    max_tokens: int = 4096
    custom_endpoints: list[CustomEndpoint] = field(default_factory=list)


class LLMClient:
    """Low-level LLM client wrapping LiteLLM.

    Handles API calls, streaming, custom endpoints, and token counting.
    All provider-specific logic is handled by LiteLLM.
    """

    def __init__(
        self,
        config: ProviderConfig | None = None,
        event_bus: EventBus | None = None,
    ) -> None:
        self._config = config or ProviderConfig()
        self._bus = event_bus
        self._custom_endpoints: dict[str, CustomEndpoint] = {}

        for ep in self._config.custom_endpoints:
            self._custom_endpoints[ep.name] = ep

    @property
    def config(self) -> ProviderConfig:
        return self._config

    async def complete(
        self,
        messages: list[dict[str, Any]],
        model: str | None = None,
        tools: list[dict[str, Any]] | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
    ) -> LLMResponse:
        """Send a completion request to the LLM.

        Uses LiteLLM to route to the correct provider based on model string.
        """
        effective_model = model or self._config.default_model
        effective_temp = temperature if temperature is not None else self._config.temperature
        effective_max = max_tokens or self._config.max_tokens

        kwargs: dict[str, Any] = {
            "model": effective_model,
            "messages": messages,
            "temperature": effective_temp,
            "max_tokens": effective_max,
        }

        # Check if this is a custom endpoint
        custom = self._resolve_custom_endpoint(effective_model)
        if custom:
            kwargs["model"] = custom.model_name
            kwargs["api_base"] = custom.base_url
            kwargs["api_key"] = custom.api_key

        if tools:
            kwargs["tools"] = tools

        try:
            response = await litellm.acompletion(**kwargs)

            result = self._parse_response(response, effective_model)

            if self._bus:
                await self._bus.emit(
                    Event(
                        event_type=EventType.COST_UPDATE,
                        data={
                            "model": result.model,
                            "input_tokens": result.input_tokens,
                            "output_tokens": result.output_tokens,
                        },
                    )
                )

            return result
        except Exception:
            logger.exception("LLM call failed for model '%s'", effective_model)
            raise

    async def stream(
        self,
        messages: list[dict[str, Any]],
        model: str | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
    ) -> AsyncIterator[str]:
        """Stream a completion response token by token."""
        effective_model = model or self._config.default_model
        effective_temp = temperature if temperature is not None else self._config.temperature
        effective_max = max_tokens or self._config.max_tokens

        kwargs: dict[str, Any] = {
            "model": effective_model,
            "messages": messages,
            "temperature": effective_temp,
            "max_tokens": effective_max,
            "stream": True,
        }

        custom = self._resolve_custom_endpoint(effective_model)
        if custom:
            kwargs["model"] = custom.model_name
            kwargs["api_base"] = custom.base_url
            kwargs["api_key"] = custom.api_key

        response = await litellm.acompletion(**kwargs)
        async for chunk in response:
            delta = chunk.choices[0].delta
            if delta and delta.content:
                yield delta.content

    def _resolve_custom_endpoint(self, model: str) -> CustomEndpoint | None:
        """Check if a model name maps to a custom endpoint."""
        # Custom endpoints are referenced by name: "custom/<name>"
        if model.startswith("custom/"):
            name = model[len("custom/") :]
            return self._custom_endpoints.get(name)
        return None

    def _parse_response(self, response: Any, model: str) -> LLMResponse:
        """Parse a LiteLLM response into our standardized format."""
        choice = response.choices[0]
        message = choice.message
        usage = response.usage or {}

        tool_calls: list[dict[str, Any]] = []
        if message.tool_calls:
            for tc in message.tool_calls:
                tool_calls.append(
                    {
                        "id": tc.id,
                        "name": tc.function.name,
                        "arguments": json.loads(tc.function.arguments),
                    }
                )

        return LLMResponse(
            content=message.content or "",
            tool_calls=tool_calls,
            input_tokens=getattr(usage, "prompt_tokens", 0),
            output_tokens=getattr(usage, "completion_tokens", 0),
            model=model,
            finish_reason=choice.finish_reason or "",
        )


class AgentLLMProvider:
    """Bridge between LLMClient and the LLMProvider protocol for the ReAct loop.

    Translates the ReAct loop's think/reflect calls into structured
    LLM API calls with proper prompt engineering.
    """

    def __init__(
        self,
        client: LLMClient,
        model: str | None = None,
    ) -> None:
        self._client = client
        self._model = model

    async def think(
        self,
        task: str,
        steps: list[Step],
        available_tools: list[dict[str, Any]],
    ) -> tuple[str, Action]:
        """Produce a thought and next action for the ReAct loop."""
        messages = self._build_think_messages(task, steps)
        tools = self._format_tools_for_api(available_tools)

        response = await self._client.complete(
            messages=messages,
            model=self._model,
            tools=tools if tools else None,
        )

        return self._parse_think_response(response)

    async def reflect(
        self,
        task: str,
        steps: list[Step],
        latest_observation: str,
    ) -> str:
        """Reflect on progress."""
        messages = [
            {
                "role": "system",
                "content": (
                    "You are reflecting on task progress. Given the task, "
                    "steps taken so far, and the latest observation, assess "
                    "whether we're making progress. Be brief (1-2 sentences)."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Task: {task}\n\n"
                    f"Steps taken: {len(steps)}\n"
                    f"Latest observation: {latest_observation[:500]}\n\n"
                    "Are we on track? Should we change approach?"
                ),
            },
        ]

        response = await self._client.complete(
            messages=messages,
            model=self._model,
            max_tokens=200,
        )
        return response.content

    def _build_think_messages(self, task: str, steps: list[Step]) -> list[dict[str, Any]]:
        """Build the message history for a think call."""
        messages: list[dict[str, Any]] = [
            {
                "role": "system",
                "content": (
                    "You are AgentScout, a helpful AI assistant that completes "
                    "tasks using available tools. Think step by step.\n\n"
                    "For each step:\n"
                    "1. Think about what you need to do next\n"
                    "2. Choose an action: use a tool, ask the user, or finish\n\n"
                    "When you have enough information to complete the task, "
                    "respond with your final answer."
                ),
            },
            {"role": "user", "content": f"Task: {task}"},
        ]

        for step in steps:
            # Add assistant's thought + action
            messages.append(
                {
                    "role": "assistant",
                    "content": f"Thought: {step.thought}",
                }
            )
            # Add observation as user message
            if step.observation:
                messages.append(
                    {
                        "role": "user",
                        "content": f"Observation: {step.observation}",
                    }
                )

        return messages

    def _format_tools_for_api(self, available_tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Format tools for LLM API tool calling."""
        formatted: list[dict[str, Any]] = []
        for tool in available_tools:
            formatted.append(
                {
                    "type": "function",
                    "function": {
                        "name": tool["name"],
                        "description": tool.get("description", ""),
                        "parameters": tool.get("parameters", {"type": "object", "properties": {}}),
                    },
                }
            )
        return formatted

    def _parse_think_response(self, response: LLMResponse) -> tuple[str, Action]:
        """Parse an LLM response into a thought and action."""
        thought = response.content

        # If the LLM made a tool call
        if response.tool_calls:
            tc = response.tool_calls[0]
            return thought, Action(
                action_type=ActionType.TOOL_CALL,
                tool_name=tc["name"],
                tool_input=tc.get("arguments", {}),
            )

        # If the response indicates a finish (no tool call, has content)
        if response.finish_reason in ("stop", "end_turn"):
            return thought, Action(
                action_type=ActionType.FINISH,
                message=thought,
            )

        # Default: finish with the content
        return thought, Action(
            action_type=ActionType.FINISH,
            message=thought,
        )
