"""
RAG Response Validation

Validates generated responses for grounding, citation accuracy,
and safety-critical requirements.
"""

import logging
import re
from dataclasses import dataclass, field
from enum import Enum

from sage.rag.context import RAGContext

logger = logging.getLogger(__name__)


class ValidationLevel(str, Enum):
    """Level of validation strictness."""

    STRICT = "strict"  # All claims must be grounded
    MODERATE = "moderate"  # Most claims should be grounded
    LENIENT = "lenient"  # Basic citation check only


@dataclass
class ValidationIssue:
    """A single validation issue found in a response."""

    issue_type: str
    message: str
    severity: str = "warning"  # error, warning, info
    location: str | None = None  # Where in response the issue is

    def __str__(self) -> str:
        prefix = f"[{self.severity.upper()}]"
        return f"{prefix} {self.issue_type}: {self.message}"


@dataclass
class ValidationResult:
    """Result of response validation."""

    is_valid: bool
    issues: list[ValidationIssue] = field(default_factory=list)

    # Grounding metrics
    grounded: bool = True
    grounding_score: float = 1.0

    # Citation metrics
    citations_found: list[str] = field(default_factory=list)
    citations_verified: list[str] = field(default_factory=list)
    citations_missing: list[str] = field(default_factory=list)

    # Safety flags
    has_safety_warnings: bool = False
    requires_verification: bool = False

    @property
    def error_count(self) -> int:
        return sum(1 for i in self.issues if i.severity == "error")

    @property
    def warning_count(self) -> int:
        return sum(1 for i in self.issues if i.severity == "warning")

    def add_issue(self, issue: ValidationIssue) -> None:
        """Add a validation issue."""
        self.issues.append(issue)
        if issue.severity == "error":
            self.is_valid = False


class CitationExtractor:
    """Extracts citations from generated text."""

    # Pattern for [DocID, Section] style citations
    CITATION_PATTERN = re.compile(
        r"\[([A-Za-z0-9\-/\.\s]+(?:,\s*[^,\]]+)?)\]"
    )

    # Pattern for (DocID) style citations
    PAREN_CITATION_PATTERN = re.compile(
        r"\(([A-Z][A-Za-z0-9\-/\.]+(?:-\d+)?(?:,\s*[^)]+)?)\)"
    )

    def extract(self, text: str) -> list[str]:
        """
        Extract all citations from text.

        Args:
            text: Generated response text

        Returns:
            List of citation strings found
        """
        citations = []

        # Find bracket-style citations
        for match in self.CITATION_PATTERN.finditer(text):
            citation = match.group(1).strip()
            if self._looks_like_citation(citation):
                citations.append(citation)

        # Find parenthetical citations
        for match in self.PAREN_CITATION_PATTERN.finditer(text):
            citation = match.group(1).strip()
            if self._looks_like_citation(citation):
                citations.append(citation)

        return list(set(citations))

    def _looks_like_citation(self, text: str) -> bool:
        """Check if text looks like a citation."""
        # Must contain some alphanumeric
        if not any(c.isalnum() for c in text):
            return False

        # Common citation patterns
        citation_indicators = [
            "NUREG", "ANSI", "ANS", "CFR", "DOE",
            "ORNL", "SAND", "LA-", "BNL", "INL",
            "Table", "Section", "Fig", "p.", "pg."
        ]

        for indicator in citation_indicators:
            if indicator.lower() in text.lower():
                return True

        # Pattern like "XXX-1234" or "10 CFR"
        if re.search(r"\d", text) and re.search(r"[A-Z]", text):
            return True

        return False

    def extract_doc_ids(self, citations: list[str]) -> list[str]:
        """
        Extract document IDs from citation strings.

        Args:
            citations: List of citation strings

        Returns:
            List of document IDs
        """
        doc_ids = []
        for citation in citations:
            # Take the first part before comma
            parts = citation.split(",")
            doc_id = parts[0].strip()
            doc_ids.append(doc_id)
        return list(set(doc_ids))


class GroundingChecker:
    """
    Checks if response claims are grounded in context.

    Verifies that factual claims in the response have
    supporting evidence in the provided context.
    """

    def __init__(self, strict: bool = False):
        """
        Initialize grounding checker.

        Args:
            strict: If True, requires exact matches
        """
        self.strict = strict

    def check(
        self,
        response: str,
        context: RAGContext,
    ) -> tuple[bool, float, list[ValidationIssue]]:
        """
        Check if response is grounded in context.

        Args:
            response: Generated response text
            context: RAGContext used for generation

        Returns:
            Tuple of (is_grounded, grounding_score, issues)
        """
        issues = []

        # Extract key claims from response (simplified)
        claims = self._extract_claims(response)

        if not claims:
            return True, 1.0, []

        # Check each claim against context
        grounded_count = 0
        for claim in claims:
            is_grounded, issue = self._check_claim(claim, context)
            if is_grounded:
                grounded_count += 1
            elif issue:
                issues.append(issue)

        grounding_score = grounded_count / len(claims) if claims else 1.0
        is_grounded = grounding_score >= (0.9 if self.strict else 0.7)

        return is_grounded, grounding_score, issues

    def _extract_claims(self, response: str) -> list[str]:
        """Extract factual claims from response (simplified)."""
        claims = []

        # Split into sentences
        sentences = re.split(r"[.!?]+", response)

        for sentence in sentences:
            sentence = sentence.strip()
            if not sentence:
                continue

            # Look for numeric claims (likely factual)
            if re.search(r"\d+\.?\d*\s*(kg|g|cm|mm|m|%|ppm|atoms)", sentence.lower()):
                claims.append(sentence)

            # Look for definitive statements
            elif any(keyword in sentence.lower() for keyword in
                    ["must", "shall", "requires", "limit", "maximum", "minimum"]):
                claims.append(sentence)

        return claims

    def _check_claim(
        self, claim: str, context: RAGContext
    ) -> tuple[bool, ValidationIssue | None]:
        """Check if a claim is grounded in context."""
        # Simple check: look for key terms in context
        claim_lower = claim.lower()

        # Extract numbers from claim
        numbers = re.findall(r"\d+\.?\d*", claim)

        # Check if any context chunk contains the claim elements
        for chunk in context.chunks:
            content_lower = chunk.content.lower()

            # Check for number presence
            numbers_found = sum(1 for n in numbers if n in content_lower)

            # Check for key term overlap
            # Filter out common words
            stop_words = {"the", "a", "an", "is", "are", "was", "were", "be", "been",
                         "being", "have", "has", "had", "do", "does", "did", "will",
                         "would", "could", "should", "may", "might", "must", "shall",
                         "to", "of", "in", "for", "on", "with", "at", "by", "from",
                         "as", "this", "that", "these", "those", "it", "its", "and",
                         "or", "but", "if", "then", "when", "where", "which", "who"}

            claim_words = set(claim_lower.split()) - stop_words
            content_words = set(content_lower.split()) - stop_words
            overlap = len(claim_words & content_words)

            # If we find numbers and some overlap, consider grounded
            if numbers_found > 0 and overlap >= 2:
                return True, None

            # Check for substantial text overlap (30% of meaningful words)
            if claim_words and overlap >= max(2, len(claim_words) * 0.3):
                return True, None

            # Check for key NCS terms matching
            ncs_terms = {"limit", "critical", "subcritical", "mass", "concentration",
                        "moderation", "reflection", "u-235", "pu-239", "enrichment"}
            claim_ncs = claim_words & ncs_terms
            content_ncs = content_words & ncs_terms
            if claim_ncs and (claim_ncs & content_ncs):
                return True, None

        # Not grounded - create issue
        issue = ValidationIssue(
            issue_type="ungrounded_claim",
            message=f"Claim may not be fully supported by context: {claim[:100]}...",
            severity="warning",
        )
        return False, issue


class CitationVerifier:
    """Verifies that citations reference actual sources in context."""

    def __init__(self):
        self.extractor = CitationExtractor()

    def verify(
        self,
        response: str,
        context: RAGContext,
    ) -> tuple[list[str], list[str], list[ValidationIssue]]:
        """
        Verify citations in response against context.

        Args:
            response: Generated response text
            context: RAGContext with source chunks

        Returns:
            Tuple of (verified_citations, missing_citations, issues)
        """
        issues = []

        # Extract citations from response
        citations = self.extractor.extract(response)
        cited_doc_ids = self.extractor.extract_doc_ids(citations)

        # Get doc IDs from context
        context_doc_ids = set(context.source_doc_ids)

        # Check which citations are verified
        verified = []
        missing = []

        for doc_id in cited_doc_ids:
            # Check for exact or partial match
            if self._doc_id_in_context(doc_id, context_doc_ids):
                verified.append(doc_id)
            else:
                missing.append(doc_id)
                issues.append(ValidationIssue(
                    issue_type="fabricated_citation",
                    message=f"Citation '{doc_id}' not found in provided context",
                    severity="error",
                ))

        return verified, missing, issues

    def _doc_id_in_context(self, doc_id: str, context_ids: set[str]) -> bool:
        """Check if doc_id matches any context document."""
        doc_id_lower = doc_id.lower()

        for context_id in context_ids:
            # Exact match
            if doc_id_lower == context_id.lower():
                return True

            # Partial match (e.g., "NUREG-1520" matches "NUREG-1520-Rev2")
            if doc_id_lower in context_id.lower() or context_id.lower() in doc_id_lower:
                return True

        return False


class SafetyChecker:
    """
    Checks response for safety-critical issues.

    Specifically designed for nuclear safety domain.
    """

    # Phrases that might indicate non-conservative advice
    NON_CONSERVATIVE_PHRASES = [
        "can be safely ignored",
        "not necessary",
        "overkill",
        "excessive margin",
        "too conservative",
        "doesn't matter",
        "negligible risk",
    ]

    # Phrases requiring verification
    VERIFICATION_TRIGGERS = [
        "calculate",
        "computation",
        "keff",
        "k-effective",
        "monte carlo",
        "scale",
        "mcnp",
        "critical mass",
        "critical concentration",
    ]

    def check(self, response: str) -> tuple[bool, bool, list[ValidationIssue]]:
        """
        Check response for safety issues.

        Args:
            response: Generated response text

        Returns:
            Tuple of (has_warnings, requires_verification, issues)
        """
        issues = []
        has_warnings = False
        requires_verification = False

        response_lower = response.lower()

        # Check for non-conservative language
        for phrase in self.NON_CONSERVATIVE_PHRASES:
            if phrase in response_lower:
                has_warnings = True
                issues.append(ValidationIssue(
                    issue_type="non_conservative",
                    message=f"Potentially non-conservative statement detected: '{phrase}'",
                    severity="warning",
                ))

        # Check if verification is needed
        for trigger in self.VERIFICATION_TRIGGERS:
            if trigger in response_lower:
                requires_verification = True
                break

        if requires_verification:
            issues.append(ValidationIssue(
                issue_type="verification_needed",
                message="Response contains calculations that should be independently verified",
                severity="info",
            ))

        return has_warnings, requires_verification, issues


class ResponseValidator:
    """
    Main validator that combines all validation checks.

    Example:
        validator = ResponseValidator()
        result = validator.validate(response_text, context)

        if not result.is_valid:
            for issue in result.issues:
                print(issue)
    """

    def __init__(
        self,
        level: ValidationLevel = ValidationLevel.MODERATE,
        check_grounding: bool = True,
        check_citations: bool = True,
        check_safety: bool = True,
    ):
        """
        Initialize response validator.

        Args:
            level: Validation strictness level
            check_grounding: Whether to check claim grounding
            check_citations: Whether to verify citations
            check_safety: Whether to check for safety issues
        """
        self.level = level
        self.check_grounding = check_grounding
        self.check_citations = check_citations
        self.check_safety = check_safety

        # Initialize sub-validators
        self.grounding_checker = GroundingChecker(strict=level == ValidationLevel.STRICT)
        self.citation_verifier = CitationVerifier()
        self.safety_checker = SafetyChecker()

    def validate(
        self,
        response: str,
        context: RAGContext,
    ) -> ValidationResult:
        """
        Validate a generated response.

        Args:
            response: Generated response text
            context: RAGContext used for generation

        Returns:
            ValidationResult with all issues found
        """
        result = ValidationResult(is_valid=True)

        # Check if response is empty
        if not response or not response.strip():
            result.add_issue(ValidationIssue(
                issue_type="empty_response",
                message="Response is empty",
                severity="error",
            ))
            return result

        # Check grounding
        if self.check_grounding:
            grounded, score, issues = self.grounding_checker.check(response, context)
            result.grounded = grounded
            result.grounding_score = score
            for issue in issues:
                result.add_issue(issue)

        # Verify citations
        if self.check_citations:
            verified, missing, issues = self.citation_verifier.verify(response, context)
            result.citations_found = verified + missing
            result.citations_verified = verified
            result.citations_missing = missing
            for issue in issues:
                result.add_issue(issue)

        # Check safety
        if self.check_safety:
            has_warnings, needs_verify, issues = self.safety_checker.check(response)
            result.has_safety_warnings = has_warnings
            result.requires_verification = needs_verify
            for issue in issues:
                result.add_issue(issue)

        # Final validity check based on level
        if self.level == ValidationLevel.STRICT:
            # Strict: no errors or warnings allowed
            result.is_valid = result.error_count == 0 and result.warning_count == 0
        elif self.level == ValidationLevel.MODERATE:
            # Moderate: no errors, limited warnings
            result.is_valid = result.error_count == 0 and result.warning_count <= 2
        else:
            # Lenient: only errors matter
            result.is_valid = result.error_count == 0

        return result


def validate_response(
    response: str,
    context: RAGContext,
    strict: bool = False,
) -> ValidationResult:
    """
    Convenience function to validate a response.

    Args:
        response: Generated response text
        context: RAGContext used for generation
        strict: Whether to use strict validation

    Returns:
        ValidationResult
    """
    level = ValidationLevel.STRICT if strict else ValidationLevel.MODERATE
    validator = ResponseValidator(level=level)
    return validator.validate(response, context)
