"""Thin wrapper for LLM API calls used by judges and condition runners."""

from __future__ import annotations

import asyncio
import hashlib
import json
import os
from pathlib import Path
from typing import Optional

from openai import AsyncOpenAI, RateLimitError


class JudgeLLM:
    """Stateless LLM client for judge evaluations (and condition generation)."""

    def __init__(
        self,
        model: Optional[str] = None,
        temperature: Optional[float] = None,
        cache_dir: Optional[Path] = None,
    ):
        self.client = AsyncOpenAI(api_key=os.environ["EVAL_JUDGE_API_KEY"])
        self.model = model or os.environ.get("EVAL_JUDGE_MODEL", "gpt-4o")
        self.temperature = (
            temperature
            if temperature is not None
            else float(os.environ.get("EVAL_JUDGE_TEMPERATURE", "0.0"))
        )
        self.cache_dir = cache_dir or Path(
            os.environ.get("EVAL_CACHE_DIR", "evaluation/results")
        )

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def judge(self, system_prompt: str, user_prompt: str) -> str:
        """Call judge LLM with caching.  Returns raw text response."""
        cache_key = self._cache_key(system_prompt, user_prompt)
        cached = self._load_cache(cache_key)
        if cached is not None:
            return cached

        text = await self._call_with_retry(system_prompt, user_prompt)
        self._save_cache(cache_key, text)
        return text

    async def generate(self, prompt: str) -> str:
        """Simple single-turn generation (used by condition runners)."""
        return await self.judge("You are a helpful analytical assistant.", prompt)

    # ------------------------------------------------------------------
    # Retry logic for rate limits
    # ------------------------------------------------------------------

    async def _call_with_retry(
        self, system_prompt: str, user_prompt: str, max_retries: int = 8
    ) -> str:
        """Call OpenAI API with exponential backoff on rate limit errors."""
        for attempt in range(max_retries + 1):
            try:
                response = await self.client.chat.completions.create(
                    model=self.model,
                    temperature=self.temperature,
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_prompt},
                    ],
                )
                return response.choices[0].message.content or ""
            except RateLimitError:
                if attempt == max_retries:
                    raise
                wait = min(2 ** attempt * 2, 60)  # 2s, 4s, 8s, … 60s cap
                await asyncio.sleep(wait)

    # ------------------------------------------------------------------
    # Cache helpers
    # ------------------------------------------------------------------

    def _cache_key(self, system: str, user: str) -> str:
        content = f"{self.model}|{self.temperature}|{system}|{user}"
        return hashlib.sha256(content.encode()).hexdigest()[:16]

    def _load_cache(self, key: str) -> Optional[str]:
        path = self.cache_dir / "judge_cache" / f"{key}.json"
        if path.exists():
            data = json.loads(path.read_text(encoding="utf-8"))
            return data.get("response")
        return None

    def _save_cache(self, key: str, response: str) -> None:
        path = self.cache_dir / "judge_cache" / f"{key}.json"
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(
            json.dumps({"response": response}, ensure_ascii=False),
            encoding="utf-8",
        )
