"""
SAGE Abstention Handler

Determines when SAGE should abstain from providing an answer.
This is a safety-critical module requiring 100% test coverage.

Abstention Triggers:
1. Outside validated domain (query doesn't match NCS scope)
2. Insufficient retrieval confidence (low RAG scores)
3. Conflicting sources (contradictory retrieved documents)
4. Calculation limitations (query requires tools beyond capability)
5. Dual-use concerns (potential misuse of information)
6. Ambiguous safety implications (unclear if response is conservative)
"""

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

logger = logging.getLogger(__name__)


class AbstentionTrigger(str, Enum):
    """Conditions that trigger abstention from answering."""

    OUTSIDE_VALIDATED_DOMAIN = "outside_validated_domain"
    INSUFFICIENT_CONFIDENCE = "insufficient_confidence"
    CONFLICTING_SOURCES = "conflicting_sources"
    CALCULATION_LIMITATIONS = "calculation_limitations"
    DUAL_USE_CONCERNS = "dual_use_concerns"
    AMBIGUOUS_SAFETY = "ambiguous_safety"


@dataclass
class AbstentionResult:
    """Result of abstention evaluation."""

    should_abstain: bool
    triggers: list[AbstentionTrigger] = field(default_factory=list)
    reason: str = ""
    partial_info: str | None = None
    next_steps: list[str] = field(default_factory=list)
    confidence: float = 1.0  # Confidence in the abstention decision

    def __bool__(self) -> bool:
        """Allow using result in boolean context."""
        return self.should_abstain


@dataclass
class AbstentionConfig:
    """Configuration for abstention thresholds."""

    # Minimum confidence score to not abstain
    min_confidence_threshold: float = 0.5

    # Maximum allowed conflict ratio between sources
    max_conflict_ratio: float = 0.3

    # Minimum number of sources needed
    min_sources_required: int = 1

    # Domain scope keywords for NCS
    domain_keywords: list[str] = field(default_factory=lambda: [
        "criticality", "nuclear", "fissile", "subcritical", "keff",
        "neutron", "enrichment", "uranium", "plutonium", "moderation",
        "reflection", "mass limit", "concentration", "fission",
        "u-235", "u-238", "pu-239", "pu-240", "ncs", "ansi/ans",
        "nureg", "10 cfr", "safety basis", "criticality safety",
    ])

    # Out of scope indicators
    out_of_scope_indicators: list[str] = field(default_factory=lambda: [
        "californium-251", "cf-251",  # Exotic isotopes not in standard tables
        "curium", "berkelium", "einsteinium",  # Transuranic elements with limited data
        "antimatter", "fusion",  # Not NCS domain
        "medical isotopes", "radiotherapy",  # Medical domain
    ])

    # Dual-use concern patterns
    dual_use_patterns: list[str] = field(default_factory=lambda: [
        r"weapon", r"bomb", r"explosive", r"detonate",
        r"improvised", r"homemade", r"DIY.*nuclear",
        r"how to make.*critical", r"construct.*device",
        r"assembly.*supercritical", r"achieve.*criticality.*outside",
        r"build.*nuclear.*device", r"build.*device",
        r"fastest.*criticality", r"quickest.*criticality",
        r"\bIND\b", r"improvised.*nuclear",  # IND = Improvised Nuclear Device
        r"nuclear.*device", r"critical.*assembly",
    ])

    # Calculation keywords requiring tool support
    calculation_keywords: list[str] = field(default_factory=lambda: [
        "calculate keff", "compute k-effective", "monte carlo simulation",
        "run scale", "run mcnp", "perform criticality calculation",
        "simulate", "model the geometry", "transport calculation",
    ])

    # Ambiguous safety patterns
    ambiguous_safety_patterns: list[str] = field(default_factory=lambda: [
        r"might be safe", r"probably okay", r"should be fine",
        r"likely acceptable", r"possibly within", r"could be conservative",
        r"uncertain whether", r"unclear if",
    ])


class AbstentionHandler:
    """
    Evaluates whether SAGE should abstain from answering a query.

    The handler checks multiple safety conditions and returns an
    AbstentionResult indicating whether to abstain and why.

    Example:
        handler = AbstentionHandler()

        # Evaluate a query with RAG results
        result = handler.evaluate(
            query="What is the critical mass of Cf-251?",
            rag_results=[...],  # Retrieved documents
            context={"domain": "ncs"}
        )

        if result.should_abstain:
            print(f"Abstaining: {result.reason}")
            print(f"Next steps: {result.next_steps}")
    """

    def __init__(self, config: AbstentionConfig | None = None):
        """Initialize abstention handler.

        Args:
            config: Optional configuration for thresholds
        """
        self.config = config or AbstentionConfig()

    def evaluate(
        self,
        query: str,
        rag_results: list[dict[str, Any]],
        context: dict[str, Any] | None = None,
    ) -> AbstentionResult:
        """
        Evaluate whether to abstain from answering.

        Args:
            query: User's question
            rag_results: Retrieved documents with confidence scores
            context: Optional context (domain, user role, etc.)

        Returns:
            AbstentionResult indicating whether to abstain
        """
        context = context or {}
        triggers: list[AbstentionTrigger] = []
        reasons: list[str] = []
        next_steps: list[str] = []
        partial_info: str | None = None

        # Check each abstention condition
        if self._check_outside_domain(query, context):
            triggers.append(AbstentionTrigger.OUTSIDE_VALIDATED_DOMAIN)
            reasons.append("Query is outside the validated NCS domain")
            next_steps.append("Consult domain-specific resources or experts")

        confidence_check = self._check_insufficient_confidence(rag_results)
        if confidence_check["triggered"]:
            triggers.append(AbstentionTrigger.INSUFFICIENT_CONFIDENCE)
            reasons.append(f"Retrieval confidence too low ({confidence_check['avg_confidence']:.2f})")
            next_steps.append("Try rephrasing the query or consult primary sources directly")
            if confidence_check.get("partial"):
                partial_info = confidence_check["partial"]

        conflict_check = self._check_conflicting_sources(rag_results)
        if conflict_check["triggered"]:
            triggers.append(AbstentionTrigger.CONFLICTING_SOURCES)
            reasons.append("Retrieved sources contain conflicting information")
            next_steps.append("Review source documents directly to resolve conflicts")
            if conflict_check.get("conflict_details"):
                partial_info = (partial_info or "") + "\n" + conflict_check["conflict_details"]

        if self._check_calculation_limitations(query):
            triggers.append(AbstentionTrigger.CALCULATION_LIMITATIONS)
            reasons.append("Query requires computational tools beyond current capabilities")
            next_steps.append("Use SCALE, MCNP, or other validated criticality codes")

        if self._check_dual_use_concerns(query):
            triggers.append(AbstentionTrigger.DUAL_USE_CONCERNS)
            reasons.append("Query raises potential dual-use or misuse concerns")
            next_steps.append("Consult with security officer or management")

        if self._check_ambiguous_safety(query, rag_results):
            triggers.append(AbstentionTrigger.AMBIGUOUS_SAFETY)
            reasons.append("Cannot determine if response would be sufficiently conservative")
            next_steps.append("Consult with qualified criticality safety engineer")

        # Build result
        should_abstain = len(triggers) > 0
        reason = "; ".join(reasons) if reasons else ""

        # Add default next step if abstaining
        if should_abstain and not next_steps:
            next_steps.append("Consult with a qualified criticality safety professional")

        return AbstentionResult(
            should_abstain=should_abstain,
            triggers=triggers,
            reason=reason,
            partial_info=partial_info.strip() if partial_info else None,
            next_steps=next_steps,
            confidence=self._calculate_abstention_confidence(triggers),
        )

    def _check_outside_domain(
        self, query: str, context: dict[str, Any]
    ) -> bool:
        """Check if query is outside validated NCS domain."""
        query_lower = query.lower()

        # Check for explicit out-of-scope indicators
        for indicator in self.config.out_of_scope_indicators:
            if indicator.lower() in query_lower:
                logger.info(f"Out of scope indicator found: {indicator}")
                return True

        # Check if query has any domain relevance
        domain_relevance = sum(
            1 for keyword in self.config.domain_keywords
            if keyword.lower() in query_lower
        )

        # If no domain keywords found and context doesn't indicate NCS
        if domain_relevance == 0:
            # Check context for domain hints
            if context.get("domain") == "ncs":
                # Even if no keywords, user explicitly in NCS context
                return False
            return True

        return False

    def _check_insufficient_confidence(
        self, rag_results: list[dict[str, Any]]
    ) -> dict[str, Any]:
        """Check if RAG confidence is too low."""
        if not rag_results:
            return {
                "triggered": True,
                "avg_confidence": 0.0,
                "partial": "No relevant documents found in knowledge base.",
            }

        # Calculate average confidence
        confidences = [
            r.get("confidence", r.get("score", 0.0))
            for r in rag_results
        ]
        avg_confidence = sum(confidences) / len(confidences) if confidences else 0.0

        # Check if below threshold
        if avg_confidence < self.config.min_confidence_threshold:
            # Gather any partial information
            partial_docs = [
                r.get("doc_title", r.get("doc_id", "Unknown"))
                for r in rag_results[:3]  # Top 3 results
            ]
            partial = f"Potentially relevant sources (low confidence): {', '.join(partial_docs)}"

            return {
                "triggered": True,
                "avg_confidence": avg_confidence,
                "partial": partial,
            }

        return {"triggered": False, "avg_confidence": avg_confidence}

    def _check_conflicting_sources(
        self, rag_results: list[dict[str, Any]]
    ) -> dict[str, Any]:
        """Check if sources contain conflicting information.

        Only flags conflicts when DIFFERENT documents report different
        values for the same measurement. Values within the same document
        (e.g., unreflected vs reflected limits) are not considered conflicts.
        """
        if len(rag_results) < 2:
            return {"triggered": False}

        # Extract numeric values from content, grouped by document
        # Structure: {unit: {doc_id: [values]}}
        doc_values: dict[str, dict[str, list[float]]] = {}

        for result in rag_results:
            content = result.get("content", "")
            doc_id = result.get("doc_id", "unknown")

            # Look for numeric limits (e.g., "20.1 kg", "18.5 kg")
            matches = re.findall(
                r"(\d+\.?\d*)\s*(kg|g|cm|mm|m|%|ppm)",
                content.lower()
            )

            for value_str, unit in matches:
                value = float(value_str)
                if unit not in doc_values:
                    doc_values[unit] = {}
                if doc_id not in doc_values[unit]:
                    doc_values[unit][doc_id] = []
                doc_values[unit][doc_id].append(value)

        # Check for conflicts between DIFFERENT documents
        conflicts = []
        for unit, docs in doc_values.items():
            if len(docs) < 2:
                continue

            # Get representative value from each document (use minimum for safety)
            doc_reps = [(doc_id, min(vals)) for doc_id, vals in docs.items()]

            # Compare values between different documents
            for i, (doc_a, val_a) in enumerate(doc_reps):
                for doc_b, val_b in doc_reps[i + 1:]:
                    if doc_a == doc_b:
                        continue  # Skip same document

                    # Check if values differ significantly (>10% difference)
                    max_val = max(val_a, val_b)
                    min_val = min(val_a, val_b)

                    if max_val > 0 and (max_val - min_val) / max_val > 0.1:
                        conflicts.append(
                            f"Conflict in {unit}: {val_a} ({doc_a}) vs {val_b} ({doc_b})"
                        )

        if conflicts:
            return {
                "triggered": True,
                "conflict_details": "Detected conflicts: " + "; ".join(conflicts),
            }

        return {"triggered": False}

    def _check_calculation_limitations(self, query: str) -> bool:
        """Check if query requires calculations beyond capability."""
        query_lower = query.lower()

        for keyword in self.config.calculation_keywords:
            if keyword.lower() in query_lower:
                logger.info(f"Calculation limitation keyword found: {keyword}")
                return True

        return False

    def _check_dual_use_concerns(self, query: str) -> bool:
        """Check for dual-use or misuse concerns."""
        query_lower = query.lower()

        for pattern in self.config.dual_use_patterns:
            if re.search(pattern, query_lower, re.IGNORECASE):
                logger.warning(f"Dual-use concern pattern matched: {pattern}")
                return True

        return False

    def _check_ambiguous_safety(
        self, query: str, rag_results: list[dict[str, Any]]
    ) -> bool:
        """Check if safety implications are ambiguous."""
        # Check query for ambiguous safety language
        query_lower = query.lower()

        for pattern in self.config.ambiguous_safety_patterns:
            if re.search(pattern, query_lower, re.IGNORECASE):
                return True

        # Check if retrieved content contains ambiguous safety language
        for result in rag_results:
            content = result.get("content", "").lower()
            for pattern in self.config.ambiguous_safety_patterns:
                if re.search(pattern, content, re.IGNORECASE):
                    return True

        return False

    def _calculate_abstention_confidence(
        self, triggers: list[AbstentionTrigger]
    ) -> float:
        """Calculate confidence in abstention decision."""
        if not triggers:
            return 1.0  # High confidence in NOT abstaining

        # More triggers = higher confidence in abstaining
        # Weight certain triggers more heavily
        weights = {
            AbstentionTrigger.DUAL_USE_CONCERNS: 1.0,
            AbstentionTrigger.OUTSIDE_VALIDATED_DOMAIN: 0.9,
            AbstentionTrigger.CONFLICTING_SOURCES: 0.8,
            AbstentionTrigger.INSUFFICIENT_CONFIDENCE: 0.7,
            AbstentionTrigger.CALCULATION_LIMITATIONS: 0.8,
            AbstentionTrigger.AMBIGUOUS_SAFETY: 0.75,
        }

        total_weight = sum(weights.get(t, 0.5) for t in triggers)
        max_possible = len(triggers) * 1.0

        return min(1.0, total_weight / max_possible) if max_possible > 0 else 0.5


def format_abstention_response(result: AbstentionResult) -> str:
    """
    Format an abstention result as a user-friendly response.

    Args:
        result: AbstentionResult from evaluation

    Returns:
        Formatted response string
    """
    if not result.should_abstain:
        return ""

    lines = [
        "I cannot provide a complete answer to this query.",
        "",
        f"**Reason:** {result.reason}",
    ]

    if result.partial_info:
        lines.extend([
            "",
            "**Available Information:**",
            result.partial_info,
        ])

    if result.next_steps:
        lines.extend([
            "",
            "**Recommended Next Steps:**",
        ])
        for i, step in enumerate(result.next_steps, 1):
            lines.append(f"{i}. {step}")

    lines.extend([
        "",
        "---",
        "*This abstention is a safety feature to ensure you receive accurate, "
        "well-supported information for criticality safety decisions.*"
    ])

    return "\n".join(lines)
