"""GroundednessCheck: Verify claims are grounded in provided context."""

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

JUDGE_PROMPT = """You are an expert at evaluating whether claims in a response
are grounded in provided source material.

Given:
1. A response/output text
2. Source context chunks that may support claims

For each factual claim in the output:
- Identify which context chunk(s) support it
- If no chunk supports the claim, mark it as "ungrounded"

Score: Proportion of claims that are grounded in the context.
- 1.0 = all claims are grounded
- 0.5 = about half the claims are grounded
- 0.0 = no claims are grounded

Respond with a JSON object:
{
  "score": <float 0.0-1.0>,
  "evidence": "<list of ungrounded claims or grounding summary>"
}
"""


@register("groundedness")
class GroundednessCheck(BaseCheck):
    """Check that claims in output are grounded in provided context.

    For RAG systems, verifies every factual claim can be traced to a source.

    Args:
        threshold: Score threshold for passing (default 0.5).
    """

    requires_llm = True
    version = "1.0.0"
    description = "Verify claims are grounded in provided context"

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

        Args:
            ctx: The context to validate.

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

    def evaluate(self, ctx: Context) -> CheckResult:
        """Evaluate groundedness of claims in output.

        Args:
            ctx: Context with output and context to evaluate.

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

        provider = get_provider()

        # Build context string
        context_str = ""
        if isinstance(ctx.context, str):
            context_str = ctx.context
        elif isinstance(ctx.context, list):
            context_str = "\n\n".join(str(chunk) for chunk in ctx.context)
        elif isinstance(ctx.context, dict):
            context_str = "\n\n".join(
                f"{k}: {v}" for k, v in ctx.context.items()
            )

        prompt = (
            f"OUTPUT:\n{ctx.output}\n\n"
            f"CONTEXT:\n{context_str}\n\n"
            f"Evaluate how much of the output is grounded in the context."
        )

        resp = provider.judge(system=JUDGE_PROMPT, user=prompt)

        return CheckResult(
            name=self.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,
        )
