"""
Ollama LLM Client — wraps the Ollama REST API running in Docker.

Provides a unified interface for all experiment scripts to call local models.
"""

import json
import time
from typing import Optional
import ollama

# ── Default Configuration ─────────────────────────────────────────────
DEFAULT_BASE_URL = "http://localhost:11434"
DEFAULT_MODEL = "qwen2.5:32b"


class LLMClient:
    """A lightweight wrapper around the Ollama API."""

    def __init__(
        self,
        model: str = DEFAULT_MODEL,
        base_url: str = DEFAULT_BASE_URL,
        temperature: float = 0.7,
        top_p: float = 0.9,
        max_tokens: int = 2048,
        num_ctx: int = 8192,
        seed: Optional[int] = None,
    ):
        self.model = model
        self.client = ollama.Client(host=base_url)
        self.options = dict(
            temperature=temperature,
            top_p=top_p,
            num_predict=max_tokens,
            num_ctx=num_ctx,
        )
        if seed is not None:
            self.options["seed"] = seed

    # ── Public API ────────────────────────────────────────────────────

    def generate(self, prompt: str, system: Optional[str] = None) -> str:
        """Simple text generation (no chat history)."""
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})
        return self._chat(messages)

    def chat(self, messages: list) -> str:
        """Multi-turn chat. messages = [{"role": "user"/"assistant", "content": ...}]."""
        return self._chat(messages)

    def generate_with_metadata(self, prompt: str, system: Optional[str] = None) -> dict:
        """Like generate() but returns full response metadata."""
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})
        return self._chat_full(messages)

    def count_tokens(self, text: str) -> int:
        """Estimate token count (rough: ~4 chars per token for Chinese+English)."""
        return len(text) // 3 + 1

    @property
    def model_name(self) -> str:
        return self.model

    # ── Internal ──────────────────────────────────────────────────────

    def _chat(self, messages: list) -> str:
        resp = self._chat_full(messages)
        return resp["message"]["content"]

    def _chat_full(self, messages: list) -> dict:
        resp = self.client.chat(
            model=self.model,
            messages=messages,
            options=self.options,
        )
        return resp

    def __repr__(self) -> str:
        return f"LLMClient(model={self.model})"


# ── Convenience Factory ───────────────────────────────────────────────

def make_client(
    model: str = DEFAULT_MODEL,
    temperature: float = 0.0,
    max_tokens: int = 2048,
    seed: int = 42,
) -> LLMClient:
    """Create a deterministic client (temperature=0) suitable for evaluation."""
    return LLMClient(
        model=model,
        temperature=temperature,
        max_tokens=max_tokens,
        seed=seed,
    )


# ── Quick Test ────────────────────────────────────────────────────────

if __name__ == "__main__":
    # Test the client
    client = make_client()
    print(f"Model: {client.model_name}")
    resp = client.generate("Say hello in exactly one word.")
    print(f"Response: {resp}")