"""
LLM-based assertion helpers using LiteLLM.

Provides the core judge functionality for semantic evaluation of responses.
"""

import json
import logging
from dataclasses import dataclass
from typing import Optional

logger = logging.getLogger(__name__)

try:
    import litellm
    litellm.drop_params = True  # Auto-drop unsupported params (e.g., temperature for GPT-5)
    LITELLM_AVAILABLE = True
except ImportError:
    LITELLM_AVAILABLE = False


@dataclass
class JudgmentResult:
    """Result of an LLM judgment."""

    passed: bool
    reasoning: str
    intent: str
    actual: str


class LLMJudge:
    """
    LLM-based judge for evaluating response content.

    Uses LiteLLM to support multiple providers (OpenAI, Anthropic, Azure, etc.).

    Example:
        judge = LLMJudge(model="gpt-4o-mini", api_key="sk-...")
        await expect(response).to_satisfy("offers a friendly greeting", llm=judge)
    """

    JUDGE_SYSTEM_PROMPT = """You are an evaluator for voice agent responses. Your task is to determine if an agent's response matches the given intent.

You will receive:
1. The agent's actual response
2. The expected intent/behavior

Evaluate whether the response satisfies the intent. Be reasonably flexible - the response doesn't need to match word-for-word, but should convey the same meaning or accomplish the same goal.

Respond with a JSON object:
{
    "reasoning": "Brief explanation of why it passed or failed",
    "passed": true/false
}"""

    def __init__(
        self,
        model: str,
        api_key: Optional[str] = None,
        api_base: Optional[str] = None,
        **kwargs
    ):
        """
        Initialize the LLM judge.

        Args:
            model: LiteLLM model string. Examples:
                - "gpt-4o-mini" (OpenAI)
                - "azure/<deployment-name>" (Azure OpenAI)
                - "claude-3-haiku-20240307" (Anthropic)
                - "groq/llama3-8b-8192" (Groq)
            api_key: API key. If not provided, uses provider's env var (e.g., OPENAI_API_KEY)
            api_base: API base URL (for Azure or custom endpoints)
            **kwargs: Additional arguments passed to litellm.completion
        """
        if not LITELLM_AVAILABLE:
            raise ImportError(
                "litellm is required for LLM assertions. "
                "Install it with: pip install litellm"
            )

        self.model = model
        self.api_key = api_key
        self.api_base = api_base
        self.extra_kwargs = kwargs

        # Disable LiteLLM's verbose logging
        litellm.set_verbose = False

    async def judge(
        self,
        actual: str,
        intent: str,
    ) -> JudgmentResult:
        """
        Judge whether the actual response matches the expected intent.

        Args:
            actual: The actual response text
            intent: The expected intent/behavior description

        Returns:
            JudgmentResult with passed status and reasoning
        """
        user_prompt = f"""Actual response: "{actual}"

Expected intent: "{intent}"

Does the response satisfy the intent?"""

        logger.debug(f"Judging response with model={self.model}")
        logger.debug(f"Intent: {intent}")
        logger.debug(f"Actual: {actual}")

        kwargs = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": self.JUDGE_SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt}
            ],
            "response_format": {"type": "json_object"},
            **self.extra_kwargs
        }

        if self.api_key:
            kwargs["api_key"] = self.api_key
        if self.api_base:
            kwargs["api_base"] = self.api_base

        response = await litellm.acompletion(**kwargs)

        content = response.choices[0].message.content
        result = json.loads(content)

        passed = result.get("passed", False)
        reasoning = result.get("reasoning", "")

        if passed:
            logger.info(f"Judgment PASSED: {reasoning}")
        else:
            logger.warning(f"Judgment FAILED: {reasoning}")

        return JudgmentResult(
            passed=passed,
            reasoning=reasoning,
            intent=intent,
            actual=actual,
        )

    def judge_sync(
        self,
        actual: str,
        intent: str,
    ) -> JudgmentResult:
        """
        Synchronous version of judge().

        Args:
            actual: The actual response text
            intent: The expected intent/behavior description

        Returns:
            JudgmentResult with passed status and reasoning
        """
        user_prompt = f"""Actual response: "{actual}"

Expected intent: "{intent}"

Does the response satisfy the intent?"""

        logger.debug(f"Judging response with model={self.model}")
        logger.debug(f"Intent: {intent}")
        logger.debug(f"Actual: {actual}")

        kwargs = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": self.JUDGE_SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt}
            ],
            "temperature": 0,
            "response_format": {"type": "json_object"},
            **self.extra_kwargs
        }

        if self.api_key:
            kwargs["api_key"] = self.api_key
        if self.api_base:
            kwargs["api_base"] = self.api_base

        response = litellm.completion(**kwargs)

        content = response.choices[0].message.content
        result = json.loads(content)

        passed = result.get("passed", False)
        reasoning = result.get("reasoning", "")

        if passed:
            logger.info(f"Judgment PASSED: {reasoning}")
        else:
            logger.warning(f"Judgment FAILED: {reasoning}")

        return JudgmentResult(
            passed=passed,
            reasoning=reasoning,
            intent=intent,
            actual=actual,
        )
