"""Lightweight wrappers for invoking large language models."""

from __future__ import annotations

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


CompletionFn = Callable[[str], str]


@dataclass
class LLMClient:
    """Minimal client that delegates completions to a callable."""

    completion_fn: CompletionFn

    def complete(self, prompt: str) -> str:
        return self.completion_fn(prompt)

    @classmethod
    def from_hf_pipeline(
        cls,
        text_generator: Callable[..., list[dict[str, Any]]],
        tokenizer: Any | None = None,
        *,
        generation_kwargs: Optional[Dict[str, Any]] = None,
        strip_prompt: bool = True,
    ) -> "LLMClient":
        """Builds a client around a HuggingFace text-generation pipeline."""

        if generation_kwargs is None:
            generation_kwargs = {}
        else:
            generation_kwargs = dict(generation_kwargs)

        if (
            tokenizer is not None
            and getattr(tokenizer, "eos_token_id", None) is not None
            and "eos_token_id" not in generation_kwargs
        ):
            generation_kwargs["eos_token_id"] = tokenizer.eos_token_id

        def _completion(prompt: str) -> str:
            outputs = text_generator(prompt, **generation_kwargs)
            if not outputs:
                return ""
            generated = outputs[0].get("generated_text", "")
            if strip_prompt and generated.startswith(prompt):
                generated = generated[len(prompt) :]
            return generated.strip()

        return cls(_completion)


class MockLLMClient(LLMClient):
    """Deterministic mock useful for tests and notebooks."""

    def __init__(self, responses: Optional[dict[str, str]] = None):
        self.responses = responses or {}
        super().__init__(self._default_completion)

    def _default_completion(self, prompt: str) -> str:
        for key, value in self.responses.items():
            if key in prompt:
                return value
        # fallback: echo empty JSON structure to avoid crashes
        return self.responses.get("__default__", "[]")

