"""LLM Adapter

Thin abstraction over a concrete LLM backend (local or remote).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Dict, Optional


@dataclass
class LLMResponse:
    """Container for a single LLM interaction."""

    prompt: str
    response: str
    tokens_prompt: Optional[int] = None
    tokens_response: Optional[int] = None
    model_name: Optional[str] = None
    metadata: Dict[str, Any] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "prompt": self.prompt,
            "response": self.response,
            "tokens_prompt": self.tokens_prompt,
            "tokens_response": self.tokens_response,
            "model_name": self.model_name,
            "metadata": self.metadata or {},
        }


class LLMAdapter:
    """Adapter around a concrete LLM backend.

    backend: Callable that takes (prompt, **kwargs) and returns:
        - str
        - or dict with keys like:
            - response
            - tokens_prompt
            - tokens_response
            - model_name
    """

    def __init__(
        self,
        backend: Callable[..., Any],
        model_name: Optional[str] = None,
    ) -> None:
        self._backend = backend
        self._model_name = model_name or "unknown"

    def run(self, prompt: str, **kwargs: Any) -> LLMResponse:
        """Run the underlying LLM on a prompt and normalize the result."""
        raw = self._backend(prompt, **kwargs)

        if isinstance(raw, str):
            return LLMResponse(
                prompt=prompt,
                response=raw,
                model_name=self._model_name,
                metadata={},
            )

        if isinstance(raw, dict):
            return LLMResponse(
                prompt=prompt,
                response=str(raw.get("response", "")),
                tokens_prompt=raw.get("tokens_prompt"),
                tokens_response=raw.get("tokens_response"),
                model_name=raw.get("model_name", self._model_name),
                metadata={
                    k: v
                    for k, v in raw.items()
                    if k
                    not in {
                        "response",
                        "tokens_prompt",
                        "tokens_response",
                        "model_name",
                    }
                },
            )

        # Fallback: best-effort stringification
        return LLMResponse(
            prompt=prompt,
            response=str(raw),
            model_name=self._model_name,
            metadata={"raw_type": type(raw).__name__},
        )