"""LLM client — single `chat()` entry point that dispatches to the active provider.

Provider selection is via `settings.llm_provider`. Supported:
- "groq"      — hits Groq via the openai SDK (groq_base_url + groq_api_key).
- "anthropic" — native Anthropic SDK.
- "openai"    — openai SDK with the standard base_url (also handles vLLM /
                any OpenAI-compatible endpoint).

Every call opens a Langfuse `generation` span automatically, tagged with the
model name and populated with token usage from the provider's response.
Spans nest under whatever observation the caller has open, so the trace
hierarchy shows the wrapping step (e.g. `pick_for_question`) with the
`llm.chat` generation as a child carrying the usage.
"""

from __future__ import annotations

from dataclasses import dataclass
from functools import lru_cache
from typing import Callable

from anthropic import Anthropic
from openai import OpenAI

from api import langfuse_client as lf
from api.config import get_settings


# ── Provider clients (cached) ──

@lru_cache(maxsize=1)
def _anthropic() -> Anthropic:
    return Anthropic(api_key=get_settings().anthropic_api_key)


@lru_cache(maxsize=1)
def _groq() -> OpenAI:
    s = get_settings()
    return OpenAI(api_key=s.groq_api_key, base_url=s.groq_base_url)


@lru_cache(maxsize=1)
def _openai() -> OpenAI:
    s = get_settings()
    return OpenAI(api_key=s.openai_api_key, base_url=s.openai_base_url)


# ── Per-provider chat impls ──

@dataclass
class _Call:
    text: str
    usage: dict[str, int]   # {"input": N, "output": M, "total": N+M}


def _chat_anthropic(system: str, user: str, max_tokens: int) -> _Call:
    s = get_settings()
    msg = _anthropic().messages.create(
        model=s.anthropic_model,
        max_tokens=max_tokens,
        system=system,
        messages=[{"role": "user", "content": user}],
    )
    text = "".join(b.text for b in msg.content if getattr(b, "type", None) == "text").strip()
    return _Call(
        text=text,
        usage={
            "input": int(msg.usage.input_tokens),
            "output": int(msg.usage.output_tokens),
            "total": int(msg.usage.input_tokens + msg.usage.output_tokens),
        },
    )


def _chat_openai_compat(client: OpenAI, model: str, system: str, user: str, max_tokens: int) -> _Call:
    resp = client.chat.completions.create(
        model=model,
        max_tokens=max_tokens,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": user},
        ],
    )
    text = (resp.choices[0].message.content or "").strip()
    u = resp.usage
    usage = {
        "input": int(u.prompt_tokens) if u else 0,
        "output": int(u.completion_tokens) if u else 0,
        "total": int(u.total_tokens) if u else 0,
    }
    return _Call(text=text, usage=usage)


def _chat_groq(system: str, user: str, max_tokens: int) -> _Call:
    return _chat_openai_compat(_groq(), get_settings().groq_model, system, user, max_tokens)


def _chat_openai(system: str, user: str, max_tokens: int) -> _Call:
    return _chat_openai_compat(_openai(), get_settings().openai_model, system, user, max_tokens)


_PROVIDERS: dict[str, Callable[[str, str, int], _Call]] = {
    "groq": _chat_groq,
    "anthropic": _chat_anthropic,
    "openai": _chat_openai,
}


def _active_model() -> str:
    s = get_settings()
    return {
        "groq": s.groq_model,
        "anthropic": s.anthropic_model,
        "openai": s.openai_model,
    }.get(s.llm_provider.lower(), s.llm_provider)


# ── Public surface ──

def chat(*, system: str, user: str, max_tokens: int = 1024, span_name: str = "llm.chat") -> str:
    """Run a single chat completion. Opens a Langfuse generation span around
    the call with the model name + token usage. `span_name` is the label
    shown in Langfuse — pass something descriptive (e.g. "direct_answer.pick")
    so traces are scannable."""
    provider = get_settings().llm_provider.lower()
    impl = _PROVIDERS.get(provider)
    if impl is None:
        raise ValueError(
            f"unknown llm_provider {provider!r}; supported: {sorted(_PROVIDERS)}"
        )

    with lf.span(
        span_name,
        as_type="generation",
        model=_active_model(),
        input={"system": system, "user": user},
    ) as gen:
        call = impl(system, user, max_tokens)
        gen.update(output=call.text, usage_details=call.usage)
        return call.text