"""Universal LLM provider via LiteLLM.

Supports 100+ providers through a single interface.
Model format: "provider/model" (e.g., "anthropic/claude-sonnet-4-6-20260415",
"openai/gpt-5.4", "ollama/llama3", "gemini/gemini-2.5-flash").

No provider-specific code. New providers work automatically
when LiteLLM updates.

Includes retry with exponential backoff for transient errors and
optional failover to a fallback model.
"""

from __future__ import annotations

import asyncio
import logging
import re
from collections.abc import AsyncIterator

import litellm

from agntspace.llm.base import LLMError, LLMProvider, LLMResponse, Message, StreamChunk, ToolCall

log = logging.getLogger(__name__)

# Suppress LiteLLM's verbose logging
litellm.suppress_debug_info = True
litellm.set_verbose = False

# ── error classification ─────────────────────────────────────────────────

# (pattern, reason, retryable)
_ERROR_PATTERNS: list[tuple[re.Pattern, str, bool]] = [
    (re.compile(r"credit balance is too low", re.I), "API credit balance is too low — top up at your provider's billing page.", False),
    (re.compile(r"invalid.*api.?key|incorrect.*api.?key|authentication", re.I), "Invalid API key — check your credentials in Settings.", False),
    (re.compile(r"rate.?limit|too many requests|429", re.I), "Rate limited by the provider — try again in a moment.", True),
    (re.compile(r"quota|exceeded.*limit", re.I), "API quota exceeded — check your plan limits.", False),
    (re.compile(r"model.*not found|does not exist|invalid model", re.I), "Model not found — check the model name in Settings.", False),
    (re.compile(r"context.?length|too.?long|max.*tokens", re.I), "Message too long for the model's context window.", False),
    (re.compile(r"timeout|timed out", re.I), "Request timed out — the provider may be slow or unreachable.", True),
    (re.compile(r"connection|connect|unreachable|refused", re.I), "Cannot reach the LLM provider — check your network or provider URL.", True),
    (re.compile(r"overloaded|capacity|503", re.I), "Provider is overloaded — try again shortly.", True),
    (re.compile(r"content.*policy|safety|blocked|moderation", re.I), "Message blocked by provider's content policy.", False),
    (re.compile(r"500|internal.?server.?error", re.I), "Provider internal error.", True),
    (re.compile(r"502|bad.?gateway", re.I), "Provider gateway error.", True),
]


# GPT-5 family (gpt-5, gpt-5-codex, gpt-5.1, gpt-5.4, …) rejects any
# temperature != 1. Rather than track OpenAI's per-version compat table,
# we drop temperature entirely for that family and let the provider use
# its default. This keeps every call site free of model-specific branches
# while still letting non-GPT-5 models honor the caller's temperature.
def _model_strips_temperature(model: str) -> bool:
    """True when this model family rejects a non-default temperature.

    Matches OpenAI's GPT-5 family regardless of LiteLLM routing prefix
    (``openai/gpt-5.4``, ``gpt-5-codex``, bare ``gpt-5``, etc).
    """
    if not model:
        return False
    name = model.split("/", 1)[-1].lower()  # strip "openai/" etc.
    return name.startswith("gpt-5")


def _strip_unsupported_params(model: str, kwargs: dict) -> dict:
    """Return a copy of ``kwargs`` with params the model family rejects removed.

    Currently only drops ``temperature`` for GPT-5 models. Returns the
    same dict when nothing needs stripping, so the happy path stays cheap.
    """
    if "temperature" in kwargs and _model_strips_temperature(model):
        out = dict(kwargs)
        out.pop("temperature", None)
        return out
    return kwargs


def _classify_error(exc: Exception, model: str) -> LLMError:
    """Turn a litellm/provider exception into a user-facing LLMError."""
    raw = str(exc)
    provider = model.split("/")[0] if "/" in model else "unknown"

    for pattern, reason, retryable in _ERROR_PATTERNS:
        if pattern.search(raw):
            return LLMError(reason, provider=provider, model=model, original=exc, retryable=retryable)

    # Fallback: include a trimmed version of the raw error
    short = raw[:200].strip()
    return LLMError(f"LLM request failed: {short}", provider=provider, model=model, original=exc, retryable=False)


# ── cache token extraction ───────────────────────────────────────────────


def _extract_cache_tokens(resp) -> tuple[int, int]:
    """Extract cache read/write token counts from a LiteLLM response.

    LiteLLM passes through provider-specific usage fields:
    - Anthropic: usage.cache_read_input_tokens, usage.cache_creation_input_tokens
    - OpenAI: usage.prompt_tokens_details.cached_tokens
    Returns (cache_read, cache_write).
    """
    cache_read = 0
    cache_write = 0
    usage = getattr(resp, "usage", None)
    if not usage:
        return 0, 0

    # Anthropic style (LiteLLM passes these through on Claude models)
    cache_read = getattr(usage, "cache_read_input_tokens", 0) or 0
    cache_write = getattr(usage, "cache_creation_input_tokens", 0) or 0

    # OpenAI style (prompt_tokens_details.cached_tokens)
    if not cache_read:
        details = getattr(usage, "prompt_tokens_details", None)
        if details:
            cache_read = getattr(details, "cached_tokens", 0) or 0

    return cache_read, cache_write


def _extract_provider_cost(resp) -> float | None:
    """Use LiteLLM's built-in cost calculator for exact pricing.

    LiteLLM maintains an auto-updated pricing database for 100+ providers.
    This is more accurate than any manually-maintained YAML because it tracks
    price changes, tiered pricing (>200K tokens), and provider-specific
    discounts.  Returns None on failure so the caller falls back to YAML.
    """
    try:
        cost = litellm.completion_cost(completion_response=resp)
        return float(cost) if cost and cost > 0 else None
    except Exception:
        return None


# ── retry config ─────────────────────────────────────────────────────────

MAX_RETRIES = 3
RETRY_BASE_DELAY = 1.0  # seconds, doubles each retry (1s, 2s, 4s)
RETRY_MAX_DELAY = 8.0

# HTTP-level timeout injected into every litellm.acompletion() call.
# LiteLLM forwards this to httpx.Timeout, which interrupts at the socket
# read — unlike asyncio.wait_for which only sends CancelledError and
# depends on the coroutine honoring it. Deliberately slightly shorter
# than the WorkflowSupervisor step timeout (240s) so an LLM hang fails
# at this layer first (clean retryable error) rather than as a step-
# level timeout (more opaque). Callers can override by passing
# ``timeout=`` explicitly to ``complete()`` / ``complete_with_tools()``.
HTTP_TIMEOUT_SECONDS = 180.0


class UniversalProvider(LLMProvider):
    """LLM provider using LiteLLM — works with any provider.

    Supports:
    - Automatic retry with exponential backoff for transient errors
    - Optional failover to a fallback model on persistent failure
    """

    def __init__(
        self,
        model: str,
        api_key: str | None = None,
        api_base: str | None = None,
        fallback_model: str | None = None,
        fallback_api_key: str | None = None,
    ) -> None:
        self._model = model
        self._api_key = api_key
        self._api_base = api_base
        self._fallback_model = fallback_model
        self._fallback_api_key = fallback_api_key
        self._healthy = True
        self._last_error: str = ""

    def _effective_api_base(self, effective_model: str) -> str | None:
        """Return the api_base LiteLLM should use for this specific call.

        A provider built for cloud (anthropic/openai) carries no _api_base.
        But the ModelRouter's hybrid routing may hand back an `ollama/...`
        model on a per-call basis (background tier routed local). In that
        case we MUST set api_base to the local Ollama URL — otherwise
        LiteLLM's default ollama provider points at the wrong host.

        For any non-ollama model, return the configured `_api_base`
        (preserves existing behavior — None for cloud, custom for
        self-hosted).
        """
        if effective_model.startswith("ollama/"):
            return self._api_base or "http://localhost:11434"
        return self._api_base

    @property
    def model_name(self) -> str:
        """Currently active model identifier."""
        return self._model

    @property
    def healthy(self) -> bool:
        """Whether the last LLM call succeeded."""
        return self._healthy

    @property
    def last_error(self) -> str:
        """Last error message, empty if healthy."""
        return self._last_error

    async def check_health(self) -> tuple[bool, str]:
        """Quick connectivity check. Returns (healthy, error_or_empty)."""
        try:
            from agntspace.llm.base import Message
            await self.complete(
                messages=[Message(role="user", content="ping")],
                system="Reply with exactly: pong",
                max_tokens=10,
                temperature=0,
            )
            self._healthy = True
            self._last_error = ""
            return True, ""
        except Exception as exc:
            self._healthy = False
            self._last_error = str(exc)[:200]
            return False, self._last_error

    # ── internal retry helper ────────────────────────────────────────

    async def _call_with_retry(
        self,
        method_name: str,
        model: str,
        api_key: str | None,
        api_base: str | None,
        **kwargs,
    ):
        """Call litellm.acompletion with retry + optional failover.

        Returns the raw litellm response object.
        """
        # GPT-5 family rejects any temperature != 1 (OpenAI API constraint).
        # Rather than track per-version compatibility (gpt-5, gpt-5.1,
        # gpt-5.4, gpt-5-codex all behave slightly differently), drop the
        # temperature kwarg entirely for that family and let the provider
        # use its own default. This lets callers keep their ``temperature=``
        # arguments for non-GPT-5 models without every call site having
        # to branch on the model string.
        kwargs = _strip_unsupported_params(model, kwargs)

        # Enforce an HTTP-level timeout on every LiteLLM call. Without
        # this, a stalled OpenAI/Anthropic connection hangs the whole
        # request forever — the outer ``asyncio.wait_for`` we added in
        # WorkflowSupervisor sends CancelledError but LiteLLM's httpx
        # socket doesn't always honor it. Passing ``timeout=`` to
        # ``litellm.acompletion`` maps to ``httpx.Timeout`` at the network
        # layer, which IS interruptible. Callers can override by passing
        # ``timeout=`` explicitly; otherwise we inject the default.
        kwargs.setdefault("timeout", HTTP_TIMEOUT_SECONDS)

        last_err: LLMError | None = None

        # Try primary model with retries
        for attempt in range(MAX_RETRIES):
            try:
                result = await litellm.acompletion(
                    model=model,
                    api_key=api_key,
                    api_base=api_base,
                    **kwargs,
                )
                self._healthy = True
                self._last_error = ""
                return result
            except LLMError:
                raise
            except Exception as exc:
                last_err = _classify_error(exc, model)
                if not last_err.retryable or attempt == MAX_RETRIES - 1:
                    break
                delay = min(RETRY_BASE_DELAY * (2 ** attempt), RETRY_MAX_DELAY)
                log.warning(
                    "LLM %s attempt %d/%d failed (retryable): %s — retrying in %.1fs",
                    method_name, attempt + 1, MAX_RETRIES, last_err.reason, delay,
                )
                await asyncio.sleep(delay)

        # Log LLM failure to activity logger
        if last_err and hasattr(self, "_activity") and self._activity:
            try:
                self._activity.log("llm", "error", f"{model} failed: {last_err.reason}",
                                   {"model": model, "error": last_err.reason,
                                    "retryable": last_err.retryable, "method": method_name},
                                   importance="high")
            except Exception:
                pass

        # Try fallback model if available
        if last_err and self._fallback_model and model != self._fallback_model:
            log.warning(
                "Primary model %s failed, falling back to %s: %s",
                model, self._fallback_model, last_err.reason,
            )
            if hasattr(self, "_activity") and self._activity:
                try:
                    self._activity.log("llm", "fallback", f"Switching to {self._fallback_model}",
                                       {"primary": model, "fallback": self._fallback_model},
                                       importance="high")
                except Exception:
                    pass
            try:
                fb_kwargs = _strip_unsupported_params(self._fallback_model, kwargs)
                result = await litellm.acompletion(
                    model=self._fallback_model,
                    api_key=self._fallback_api_key or api_key,
                    api_base=api_base if self._fallback_model.startswith("ollama/") else None,
                    **fb_kwargs,
                )
                self._healthy = True
                self._last_error = ""
                return result
            except LLMError:
                raise
            except Exception as exc:
                fallback_err = _classify_error(exc, self._fallback_model)
                log.error("Fallback model %s also failed: %s", self._fallback_model, fallback_err.reason)
                # Raise the original error — it's more informative
                raise last_err from exc

        if last_err:
            self._healthy = False
            self._last_error = last_err.reason
            log.info("LLM %s failed: %s", method_name, last_err.reason)
            raise last_err

    # ── public API ───────────────────────────────────────────────────

    async def complete(
        self,
        messages: list[Message],
        system: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
        model: str | None = None,
    ) -> LLMResponse:
        effective_model = model or self._model
        msgs = self._build_messages(messages, system, model=effective_model)

        # Ollama models default to 4096 context — extend for longer prompts.
        # Fires for any ollama/ model regardless of provider construction
        # (cloud-built provider routing a local model still needs num_ctx
        # bumped — _effective_api_base handles the URL, this handles ctx).
        extra_kwargs: dict = {}
        if effective_model.startswith("ollama/"):
            # Estimate tokens needed: ~4 chars/token for English
            input_chars = sum(len(m.get("content", "")) for m in msgs if isinstance(m.get("content"), str))
            estimated_tokens = input_chars // 3 + max_tokens + 512  # input + output + overhead
            if estimated_tokens > 4096:
                extra_kwargs["num_ctx"] = min(estimated_tokens, 32768)

        resp = await self._call_with_retry(
            "complete",
            model=effective_model,
            api_key=self._api_key,
            api_base=self._effective_api_base(effective_model),
            messages=msgs,
            max_tokens=max_tokens,
            temperature=temperature,
            **extra_kwargs,
        )

        choice = resp.choices[0]
        cache_read, cache_write = _extract_cache_tokens(resp)
        return LLMResponse(
            content=choice.message.content or "",
            model=resp.model or self._model,
            input_tokens=resp.usage.prompt_tokens if resp.usage else 0,
            output_tokens=resp.usage.completion_tokens if resp.usage else 0,
            cache_read_tokens=cache_read,
            cache_write_tokens=cache_write,
            provider_cost=_extract_provider_cost(resp),
            stop_reason=choice.finish_reason,
        )

    async def complete_with_tools(
        self,
        messages: list[Message],
        tools: list[dict],
        system: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
        model: str | None = None,
    ) -> LLMResponse:
        import json

        effective_model = model or self._model
        msgs = self._build_messages(messages, system, model=effective_model)

        # Convert tool format to OpenAI function calling format
        # (LiteLLM uses OpenAI format universally)
        oai_tools = []
        for t in tools:
            oai_tools.append({
                "type": "function",
                "function": {
                    "name": t["name"],
                    "description": t.get("description", ""),
                    "parameters": t.get("input_schema", {}),
                },
            })

        resp = await self._call_with_retry(
            "complete_with_tools",
            model=effective_model,
            api_key=self._api_key,
            api_base=self._effective_api_base(effective_model),
            messages=msgs,
            max_tokens=max_tokens,
            temperature=temperature,
            tools=oai_tools if oai_tools else None,
        )

        choice = resp.choices[0]
        text = choice.message.content or ""
        tool_calls = []

        if choice.message.tool_calls:
            for tc in choice.message.tool_calls:
                args = tc.function.arguments
                parsed = json.loads(args) if isinstance(args, str) else args
                tool_calls.append(ToolCall(id=tc.id, name=tc.function.name, input=parsed))

        cache_read, cache_write = _extract_cache_tokens(resp)
        return LLMResponse(
            content=text,
            model=resp.model or self._model,
            input_tokens=resp.usage.prompt_tokens if resp.usage else 0,
            output_tokens=resp.usage.completion_tokens if resp.usage else 0,
            cache_read_tokens=cache_read,
            cache_write_tokens=cache_write,
            provider_cost=_extract_provider_cost(resp),
            stop_reason=choice.finish_reason,
            tool_calls=tool_calls,
            raw_content=choice.message,
        )

    async def stream(
        self,
        messages: list[Message],
        system: str | None = None,
        max_tokens: int = 4096,
        temperature: float = 0.7,
        model: str | None = None,
    ) -> AsyncIterator[StreamChunk]:
        effective_model = model or self._model
        msgs = self._build_messages(messages, system, model=effective_model)

        resp = await self._call_with_retry(
            "stream",
            model=effective_model,
            api_key=self._api_key,
            api_base=self._effective_api_base(effective_model),
            messages=msgs,
            max_tokens=max_tokens,
            temperature=temperature,
            stream=True,
            # Opt in to the trailing usage chunk. OpenAI/LiteLLM omit usage from
            # streaming responses unless this is set; without it every chat_stream
            # row is recorded with 0 tokens and $0.
            stream_options={"include_usage": True},
        )

        full_content = ""
        stop_reason: str | None = None
        final_chunk = None
        async for chunk in resp:
            if chunk.choices:
                delta = chunk.choices[0].delta
                if delta.content:
                    full_content += delta.content
                    yield StreamChunk(delta=delta.content)
                if chunk.choices[0].finish_reason and stop_reason is None:
                    stop_reason = chunk.choices[0].finish_reason
            # Usage typically arrives in a trailing chunk with empty choices.
            # Keep the most recent chunk that carries a populated usage block.
            if getattr(chunk, "usage", None):
                final_chunk = chunk

        usage = getattr(final_chunk, "usage", None) if final_chunk else None
        input_tokens = getattr(usage, "prompt_tokens", 0) or 0 if usage else 0
        output_tokens = getattr(usage, "completion_tokens", 0) or 0 if usage else 0
        cache_read, cache_write = _extract_cache_tokens(final_chunk) if final_chunk else (0, 0)
        provider_cost = _extract_provider_cost(final_chunk) if final_chunk else None

        yield StreamChunk(
            delta="",
            done=True,
            response=LLMResponse(
                content=full_content,
                model=self._model,
                stop_reason=stop_reason or "stop",
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                cache_read_tokens=cache_read,
                cache_write_tokens=cache_write,
                provider_cost=provider_cost,
            ),
        )

    def _build_messages(self, messages: list[Message], system: str | None, model: str | None = None) -> list[dict]:
        msgs: list[dict] = []
        if system:
            # Anthropic prompt caching: mark the system prompt for caching.
            # The system prompt (identity, skills, memory) is largely static
            # across calls, so caching gives ~90% discount on cached input tokens.
            # OpenAI caches automatically for prompts >1024 tokens — no markers needed.
            effective = model or self._model
            if effective.startswith("anthropic/") or "claude" in effective.lower():
                msgs.append({
                    "role": "system",
                    "content": [
                        {
                            "type": "text",
                            "text": system,
                            "cache_control": {"type": "ephemeral"},
                        }
                    ],
                })
            else:
                msgs.append({"role": "system", "content": system})
        for m in messages:
            if m.role == "tool" and m.tool_call_id:
                # Tool result: OpenAI format
                msgs.append({
                    "role": "tool",
                    "tool_call_id": m.tool_call_id,
                    "content": m.content if isinstance(m.content, str) else str(m.content),
                })
            elif m.role == "assistant" and not isinstance(m.content, str):
                # Raw LLM message object (has tool_calls) — pass as dict
                try:
                    msgs.append(m.content.model_dump())
                except (AttributeError, TypeError):
                    msgs.append({"role": "assistant", "content": str(m.content)})
            else:
                msgs.append({"role": m.role, "content": m.content})
        return msgs
