"""Ollama judge — uses the ollama Python package with structured outputs."""

from __future__ import annotations

import json
import logging

from evalkit.judges.base import BaseJudge
from evalkit.models import JudgeResponse
from evalkit.prompts import SYSTEM_PROMPT

logger = logging.getLogger(__name__)


class OllamaJudge(BaseJudge):
    """Judge using a local Ollama server via the ollama Python package."""

    def __init__(
        self,
        model: str = "gpt-oss",
        base_url: str = "http://localhost:11434",
        temperature: float = 0.0,
        max_tokens: int = 1024,
        timeout: int = 300,
    ) -> None:
        try:
            import ollama  # noqa: F401
        except ImportError as exc:
            raise ImportError("Install ollama: pip install evalkit-llm[ollama]") from exc
        self.model = model
        self.base_url = base_url
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.timeout = timeout

    def evaluate(self, prompt: str) -> JudgeResponse:
        import ollama

        logger.debug("OllamaJudge calling model=%s at %s, prompt_len=%d", self.model, self.base_url, len(prompt))

        client = ollama.Client(host=self.base_url, timeout=self.timeout)

        try:
            response = client.chat(
                model=self.model,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": prompt},
                ],
                format=JudgeResponse.model_json_schema(),
                options={"temperature": self.temperature, "num_predict": self.max_tokens},
            )
        except ollama.ResponseError as exc:
            raise ConnectionError(f"Ollama API error for model '{self.model}': {exc}") from exc
        except Exception as exc:
            if "ConnectError" in type(exc).__name__ or "Timeout" in type(exc).__name__:
                raise ConnectionError(
                    f"Could not connect to Ollama server at {self.base_url}. "
                    "Make sure Ollama is running (e.g. `ollama serve`) and the "
                    f"model '{self.model}' is available (e.g. `ollama pull {self.model}`)."
                ) from exc
            raise

        content = response.message.content
        logger.debug("OllamaJudge raw content (%d chars): %.200s", len(content), content)

        raw = json.loads(content)
        result = JudgeResponse(
            score=max(0.0, min(1.0, float(raw["score"]))),
            reason=str(raw["reason"]),
        )
        logger.debug("OllamaJudge result: score=%.3f, reason=%s", result.score, result.reason[:100])
        return result
