"""CustomCheck: User-defined validation criteria with natural language prompts."""

from __future__ import annotations

from typing import TYPE_CHECKING

from ai_qa.checks.registry import register
from ai_qa.core.check import BaseCheck
from ai_qa.core.result import CheckResult
from ai_qa.exceptions import ContextValidationError
from ai_qa.providers import get_provider

if TYPE_CHECKING:
    from ai_qa.core.context import Context


@register("custom")
class CustomCheck(BaseCheck):
    """Allow users to define custom validation criteria with natural language.

    Users provide a description of what to check, and the LLM evaluates
    the output against that criteria.

    Args:
        criteria: Natural language description of the validation criteria.
        custom_name: Name for this check instance (default "custom").
        threshold: Score threshold for passing (default 0.5).
    """

    requires_llm = True
    version = "1.0.0"
    description = "User-defined validation criteria"

    def __init__(
        self,
        criteria: str,
        custom_name: str = "custom",
        threshold: float = 0.5,
    ) -> None:
        """Initialize CustomCheck.

        Args:
            criteria: Validation criteria description.
            custom_name: Name for this check instance.
            threshold: Score threshold for passing.
        """
        super().__init__(threshold=threshold)
        self.criteria = criteria
        self.custom_name = custom_name
        self.name = custom_name

    def validate_context(self, ctx: Context) -> None:
        """Validate that output is available.

        Args:
            ctx: The context to validate.

        Raises:
            ContextValidationError: If output is missing.
        """
        if ctx.output is None:
            msg = "output is required for CustomCheck"
            raise ContextValidationError(msg)

    def evaluate(self, ctx: Context) -> CheckResult:
        """Evaluate output against custom criteria.

        Args:
            ctx: Context with output to evaluate.

        Returns:
            CheckResult with custom evaluation score.
        """
        self.validate_context(ctx)

        judge_prompt = (
            f"You are an AI output evaluator. Evaluate the OUTPUT against "
            f"the following criteria:\n\n"
            f"CRITERIA: {self.criteria}\n\n"
            f"Score from 0.0 (completely fails criteria) to 1.0 "
            f"(perfectly meets criteria).\n\n"
            f"Respond with ONLY a JSON object:\n"
            f'{{\n'
            f'  "score": <float>,\n'
            f'  "evidence": "<explanation>"\n'
            f'}}'
        )

        input_str = (
            ""
            if not ctx.input
            else (
                ctx.input
                if isinstance(ctx.input, str)
                else str(ctx.input)
            )
        )

        output_str = ctx.output
        if not isinstance(output_str, str):
            output_str = str(output_str)

        prompt = f"INPUT: {input_str}\n\nOUTPUT:\n{output_str}"

        provider = get_provider()
        resp = provider.judge(system=judge_prompt, user=prompt)

        return CheckResult(
            name=self.custom_name,
            score=resp.get("score", 0.5),
            passed=resp.get("score", 0.5) >= self.threshold,
            threshold=self.threshold,
            evidence=resp.get("evidence", ""),
            cost_usd=resp.get("cost_usd", 0.0),
            tokens_used=resp.get("tokens", {}),
            check_version=self.version,
        )
