"""OpenAI and Azure OpenAI judges using structured outputs."""

from __future__ import annotations

import logging
import os
from typing import Any

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

logger = logging.getLogger(__name__)


class OpenAIJudge(BaseJudge):
    """Judge using ``chat.completions.parse(response_format=JudgeResponse)``."""

    def __init__(
        self,
        model: str = "gpt-4o-mini",
        api_key: str | None = None,
        temperature: float = 0.0,
        max_tokens: int = 1024,
        timeout: float = 60.0,
    ) -> None:
        try:
            import openai  # noqa: F401
        except ImportError as exc:
            raise ImportError("Install openai: pip install evalkit-llm[openai]") from exc

        self.model = model
        self.api_key = api_key
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.timeout = timeout

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

        logger.debug("OpenAIJudge calling model=%s, prompt_len=%d", self.model, len(prompt))
        client = (
            openai.OpenAI(api_key=self.api_key, timeout=self.timeout)
            if self.api_key
            else openai.OpenAI(timeout=self.timeout)
        )

        completion = client.chat.completions.parse(
            model=self.model,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            response_format=JudgeResponse,
            temperature=self.temperature,
            max_tokens=self.max_tokens,
        )

        parsed = completion.choices[0].message.parsed
        result = JudgeResponse(
            score=max(0.0, min(1.0, parsed.score)),
            reason=parsed.reason,
        )
        logger.debug("OpenAIJudge response: score=%.3f, reason=%s", result.score, result.reason[:100])
        return result


class AzureOpenAIJudge(BaseJudge):
    """Judge using Azure OpenAI with the same ``openai`` package."""

    def __init__(
        self,
        model: str = "gpt-4o-mini",
        azure_endpoint: str | None = None,
        api_key: str | None = None,
        api_version: str = "2024-10-21",
        temperature: float = 0.0,
        max_tokens: int = 1024,
        timeout: float = 60.0,
    ) -> None:
        try:
            import openai  # noqa: F401
        except ImportError as exc:
            raise ImportError("Install openai: pip install evalkit-llm[openai]") from exc

        self.model = model
        self.azure_endpoint = azure_endpoint or os.environ.get("AZURE_OPENAI_ENDPOINT")
        if not self.azure_endpoint:
            raise ValueError("azure_endpoint is required. Pass it directly or set AZURE_OPENAI_ENDPOINT env var.")
        self.api_key = api_key
        self.api_version = api_version
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.timeout = timeout

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

        logger.debug(
            "AzureOpenAIJudge calling model=%s, endpoint=%s, prompt_len=%d",
            self.model,
            self.azure_endpoint,
            len(prompt),
        )
        kwargs: dict[str, Any] = {
            "api_version": self.api_version,
            "azure_endpoint": self.azure_endpoint,
        }
        if self.api_key:
            kwargs["api_key"] = self.api_key
        kwargs["timeout"] = self.timeout

        client = openai.AzureOpenAI(**kwargs)

        completion = client.chat.completions.parse(
            model=self.model,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            response_format=JudgeResponse,
            temperature=self.temperature,
            max_tokens=self.max_tokens,
        )

        parsed = completion.choices[0].message.parsed
        result = JudgeResponse(
            score=max(0.0, min(1.0, parsed.score)),
            reason=parsed.reason,
        )
        logger.debug("AzureOpenAIJudge response: score=%.3f, reason=%s", result.score, result.reason[:100])
        return result
