"""
HallucinationDetector: Graph-aligned verification for answer quality.

Implements the Composite Fidelity Index (CFI) for detecting hallucinations
by verifying entities and relations against the knowledge graph.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, NamedTuple
import re


class VerificationResult(NamedTuple):
    """Result of hallucination verification."""
    
    is_grounded: bool
    confidence_score: float  # 0.0 to 1.0
    entity_coverage: float  # Proportion of entities found in KG
    relation_coverage: float  # Proportion of relations found in KG
    temporal_consistency: float  # Score for temporal consistency
    ungrounded_claims: list[str]  # Claims not found in KG
    verification_details: dict  # Detailed verification info


@dataclass
class HallucinationDetector:
    """
    Detector for hallucinations using graph-aligned verification.
    
    Implements the Composite Fidelity Index (CFI):
    CFI = α * Entity_Coverage + β * Relation_Preservation + γ * Temporal_Consistency
    
    Where:
    - Entity_Coverage: % of mentioned entities found in KG
    - Relation_Preservation: % of stated relations verified in KG
    - Temporal_Consistency: Alignment with temporal facts
    
    Default weights: α=0.4, β=0.4, γ=0.2
    """
    
    # Graphiti instance for KG verification
    graphiti: Any = None
    
    # CFI weights
    alpha: float = 0.4  # Entity weight
    beta: float = 0.4   # Relation weight
    gamma: float = 0.2  # Temporal weight
    
    # Threshold for considering response grounded
    grounding_threshold: float = 0.7
    
    async def verify_response(
        self,
        response: str,
        query: str,
        context: dict | None = None,
    ) -> VerificationResult:
        """
        Verify a response against the knowledge graph.
        
        Args:
            response: The generated response to verify.
            query: The original query.
            context: Optional context used for generation.
            
        Returns:
            VerificationResult with grounding score and details.
        """
        # Extract entities from response
        entities = await self._extract_entities(response)
        
        # Extract relations from response
        relations = await self._extract_relations(response)
        
        # Verify entities against KG
        entity_score, grounded_entities, ungrounded_entities = await self._verify_entities(entities)
        
        # Verify relations against KG
        relation_score, grounded_relations, ungrounded_relations = await self._verify_relations(relations)
        
        # Check temporal consistency
        temporal_score = await self._check_temporal_consistency(response, context)
        
        # Calculate CFI
        cfi = (
            self.alpha * entity_score +
            self.beta * relation_score +
            self.gamma * temporal_score
        )
        
        # Collect ungrounded claims
        ungrounded_claims = []
        ungrounded_claims.extend([f"Entity: {e}" for e in ungrounded_entities])
        ungrounded_claims.extend([f"Relation: {r}" for r in ungrounded_relations])
        
        return VerificationResult(
            is_grounded=cfi >= self.grounding_threshold,
            confidence_score=cfi,
            entity_coverage=entity_score,
            relation_coverage=relation_score,
            temporal_consistency=temporal_score,
            ungrounded_claims=ungrounded_claims,
            verification_details={
                "total_entities": len(entities),
                "grounded_entities": len(grounded_entities),
                "total_relations": len(relations),
                "grounded_relations": len(grounded_relations),
                "alpha": self.alpha,
                "beta": self.beta,
                "gamma": self.gamma,
            },
        )
    
    async def _extract_entities(self, text: str) -> list[str]:
        """
        Extract entity mentions from text.
        
        Uses simple heuristics (capitalized words, quoted terms).
        For production, should use NER model.
        """
        entities = set()
        
        # Extract capitalized multi-word phrases (likely proper nouns)
        # Pattern: sequences of capitalized words
        pattern = r'\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\b'
        matches = re.findall(pattern, text)
        entities.update(matches)
        
        # Extract quoted terms
        quoted = re.findall(r'"([^"]+)"', text)
        entities.update(quoted)
        
        # Filter out common words that aren't entities
        stop_entities = {"The", "A", "An", "This", "That", "These", "Those", "It"}
        entities = [e for e in entities if e not in stop_entities]
        
        return list(entities)
    
    async def _extract_relations(self, text: str) -> list[tuple[str, str, str]]:
        """
        Extract relation triples from text.
        
        Returns list of (subject, predicate, object) tuples.
        Uses simple pattern matching. For production, should use relation extraction model.
        """
        relations = []
        
        # Common relation patterns
        patterns = [
            # "X is Y" pattern
            r'([A-Z][a-zA-Z\s]+)\s+is\s+(?:a|an|the)?\s*([a-zA-Z\s]+)',
            # "X was founded by Y" pattern
            r'([A-Z][a-zA-Z\s]+)\s+was\s+(\w+)\s+by\s+([A-Z][a-zA-Z\s]+)',
            # "X works at Y" pattern
            r'([A-Z][a-zA-Z\s]+)\s+works?\s+at\s+([A-Z][a-zA-Z\s]+)',
            # "X born in Y" pattern
            r'([A-Z][a-zA-Z\s]+)\s+(?:was\s+)?born\s+(?:in|on)\s+([\d\w\s,]+)',
        ]
        
        for pattern in patterns:
            matches = re.findall(pattern, text)
            for match in matches:
                if len(match) == 2:
                    relations.append((match[0].strip(), "related_to", match[1].strip()))
                elif len(match) == 3:
                    relations.append((match[0].strip(), match[1], match[2].strip()))
        
        return relations
    
    async def _verify_entities(
        self,
        entities: list[str],
    ) -> tuple[float, list[str], list[str]]:
        """
        Verify entities against the knowledge graph.
        
        Returns:
            Tuple of (coverage_score, grounded_entities, ungrounded_entities)
        """
        if not entities:
            return 1.0, [], []  # No entities to verify
        
        grounded = []
        ungrounded = []
        
        for entity in entities:
            # Search in Graphiti
            found = await self._entity_exists_in_kg(entity)
            if found:
                grounded.append(entity)
            else:
                ungrounded.append(entity)
        
        coverage = len(grounded) / len(entities) if entities else 1.0
        return coverage, grounded, ungrounded
    
    async def _verify_relations(
        self,
        relations: list[tuple[str, str, str]],
    ) -> tuple[float, list[tuple], list[tuple]]:
        """
        Verify relations against the knowledge graph.
        
        Returns:
            Tuple of (coverage_score, grounded_relations, ungrounded_relations)
        """
        if not relations:
            return 1.0, [], []  # No relations to verify
        
        grounded = []
        ungrounded = []
        
        for relation in relations:
            found = await self._relation_exists_in_kg(relation)
            if found:
                grounded.append(relation)
            else:
                ungrounded.append(relation)
        
        coverage = len(grounded) / len(relations) if relations else 1.0
        return coverage, grounded, ungrounded
    
    async def _entity_exists_in_kg(self, entity: str) -> bool:
        """Check if entity exists in knowledge graph."""
        if self.graphiti is None:
            return True  # Assume grounded if no KG available
        
        try:
            results = await self.graphiti.search(
                query=entity,
                num_results=1,
            )
            return len(results) > 0
        except Exception:
            return True  # Assume grounded on error
    
    async def _relation_exists_in_kg(self, relation: tuple[str, str, str]) -> bool:
        """Check if relation exists in knowledge graph."""
        if self.graphiti is None:
            return True  # Assume grounded if no KG available
        
        subject, predicate, obj = relation
        try:
            # Search for the relation fact
            query = f"{subject} {predicate} {obj}"
            results = await self.graphiti.search(
                query=query,
                num_results=3,
            )
            
            # Check if any result contains both subject and object
            for result in results:
                fact = result.fact if hasattr(result, "fact") else str(result)
                if subject.lower() in fact.lower() and obj.lower() in fact.lower():
                    return True
            
            return False
        except Exception:
            return True  # Assume grounded on error
    
    async def _check_temporal_consistency(
        self,
        response: str,
        context: dict | None,
    ) -> float:
        """
        Check temporal consistency of the response.
        
        Verifies that temporal statements align with KG facts.
        """
        if context is None:
            return 1.0  # No context to check against
        
        facts = context.get("facts", [])
        if not facts:
            return 1.0
        
        # Extract temporal statements from response
        temporal_patterns = [
            r'in\s+(\d{4})',  # "in 2024"
            r'on\s+([\w\s,]+\d{4})',  # "on January 15, 2024"
            r'since\s+(\d{4})',  # "since 2020"
        ]
        
        response_dates = []
        for pattern in temporal_patterns:
            matches = re.findall(pattern, response)
            response_dates.extend(matches)
        
        if not response_dates:
            return 1.0  # No temporal statements to verify
        
        # Check alignment with context facts
        # For now, simple overlap check
        context_text = " ".join([f.get("fact", "") for f in facts])
        
        matched = 0
        for date in response_dates:
            if date in context_text:
                matched += 1
        
        return matched / len(response_dates) if response_dates else 1.0
