"""
SAGE Output Classification

Classifies RAG responses into safety categories:
- GREEN: Safe to return, high confidence, well-cited
- YELLOW: Caution advised, review recommended
- RED: Do not return without expert review

This classification integrates with abstention and escalation systems.
"""

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

from sage.core.abstention import AbstentionResult
from sage.core.escalation import EscalationEvaluation, EscalationPriority

logger = logging.getLogger(__name__)


class OutputClassification(str, Enum):
    """Output safety classification levels."""

    GREEN = "green"  # Safe to return
    YELLOW = "yellow"  # Caution advised
    RED = "red"  # Do not return without review


@dataclass
class ClassificationResult:
    """Result of output classification."""

    classification: OutputClassification
    confidence: float = 1.0
    reasons: list[str] = field(default_factory=list)

    # Component results
    abstention: AbstentionResult | None = None
    escalation: EscalationEvaluation | None = None

    # Validation summary
    grounded: bool = True
    citations_valid: bool = True
    has_safety_warnings: bool = False

    def __str__(self) -> str:
        return f"[{self.classification.value.upper()}] {', '.join(self.reasons)}"


@dataclass
class ClassificationConfig:
    """Configuration for output classification."""

    # Confidence thresholds
    green_confidence_threshold: float = 0.8
    yellow_confidence_threshold: float = 0.5

    # Citation requirements
    min_citations_for_green: int = 1

    # Grounding requirements
    min_grounding_score_green: float = 0.8
    min_grounding_score_yellow: float = 0.5


class OutputClassifier:
    """
    Classifies RAG outputs into GREEN/YELLOW/RED categories.

    Combines results from:
    - Abstention handler (should we answer at all?)
    - Escalation evaluation (does this need review?)
    - Validation results (is the response well-grounded?)

    Example:
        classifier = OutputClassifier()

        result = classifier.classify(
            confidence=0.9,
            grounded=True,
            citations_valid=True,
            abstention_result=abstention_result,
            escalation_result=escalation_result,
        )

        if result.classification == OutputClassification.GREEN:
            # Safe to return response
            return response
        elif result.classification == OutputClassification.YELLOW:
            # Return with caution notice
            return response + "\\n\\n*Review recommended*"
        else:  # RED
            # Hold response for review
            create_escalation_request()
    """

    def __init__(self, config: ClassificationConfig | None = None):
        """Initialize output classifier.

        Args:
            config: Optional configuration
        """
        self.config = config or ClassificationConfig()

    def classify(
        self,
        confidence: float,
        grounded: bool = True,
        grounding_score: float = 1.0,
        citations_valid: bool = True,
        citation_count: int = 0,
        has_safety_warnings: bool = False,
        abstention_result: AbstentionResult | None = None,
        escalation_result: EscalationEvaluation | None = None,
    ) -> ClassificationResult:
        """
        Classify an output based on multiple factors.

        Args:
            confidence: Retrieval confidence score (0-1)
            grounded: Whether response is grounded in context
            grounding_score: Grounding score (0-1)
            citations_valid: Whether all citations are valid
            citation_count: Number of citations
            has_safety_warnings: Whether safety warnings were raised
            abstention_result: Optional abstention evaluation result
            escalation_result: Optional escalation evaluation result

        Returns:
            ClassificationResult with classification and reasons
        """
        reasons: list[str] = []

        # Check for RED conditions (immediate blockers)
        if abstention_result and abstention_result.should_abstain:
            reasons.append(f"Abstention triggered: {abstention_result.reason}")
            return ClassificationResult(
                classification=OutputClassification.RED,
                confidence=abstention_result.confidence,
                reasons=reasons,
                abstention=abstention_result,
                escalation=escalation_result,
                grounded=grounded,
                citations_valid=citations_valid,
                has_safety_warnings=has_safety_warnings,
            )

        if escalation_result and escalation_result.should_escalate:
            # Critical priority = RED
            if escalation_result.priority == EscalationPriority.CRITICAL:
                triggers = [t.value for t in escalation_result.triggers]
                reasons.append(f"Critical escalation: {', '.join(triggers)}")
                return ClassificationResult(
                    classification=OutputClassification.RED,
                    confidence=confidence,
                    reasons=reasons,
                    abstention=abstention_result,
                    escalation=escalation_result,
                    grounded=grounded,
                    citations_valid=citations_valid,
                    has_safety_warnings=has_safety_warnings,
                )

        # Check for fabricated citations (always RED)
        if not citations_valid:
            reasons.append("Response contains unverified citations")
            return ClassificationResult(
                classification=OutputClassification.RED,
                confidence=confidence,
                reasons=reasons,
                abstention=abstention_result,
                escalation=escalation_result,
                grounded=grounded,
                citations_valid=citations_valid,
                has_safety_warnings=has_safety_warnings,
            )

        # Check for YELLOW conditions
        yellow_reasons: list[str] = []

        if escalation_result and escalation_result.should_escalate:
            # Non-critical escalation = YELLOW
            triggers = [t.value for t in escalation_result.triggers]
            yellow_reasons.append(f"Escalation recommended: {', '.join(triggers)}")

        if confidence < self.config.green_confidence_threshold:
            if confidence < self.config.yellow_confidence_threshold:
                # Very low confidence = RED
                reasons.append(f"Very low confidence ({confidence:.2f})")
                return ClassificationResult(
                    classification=OutputClassification.RED,
                    confidence=confidence,
                    reasons=reasons,
                    abstention=abstention_result,
                    escalation=escalation_result,
                    grounded=grounded,
                    citations_valid=citations_valid,
                    has_safety_warnings=has_safety_warnings,
                )
            yellow_reasons.append(f"Moderate confidence ({confidence:.2f})")

        if not grounded or grounding_score < self.config.min_grounding_score_green:
            if grounding_score < self.config.min_grounding_score_yellow:
                # Very low grounding = RED
                reasons.append(f"Response poorly grounded ({grounding_score:.2f})")
                return ClassificationResult(
                    classification=OutputClassification.RED,
                    confidence=confidence,
                    reasons=reasons,
                    abstention=abstention_result,
                    escalation=escalation_result,
                    grounded=grounded,
                    citations_valid=citations_valid,
                    has_safety_warnings=has_safety_warnings,
                )
            yellow_reasons.append(f"Moderate grounding ({grounding_score:.2f})")

        if has_safety_warnings:
            yellow_reasons.append("Safety warnings present")

        if citation_count < self.config.min_citations_for_green:
            yellow_reasons.append("Insufficient citations")

        # If any YELLOW reasons, return YELLOW
        if yellow_reasons:
            return ClassificationResult(
                classification=OutputClassification.YELLOW,
                confidence=confidence,
                reasons=yellow_reasons,
                abstention=abstention_result,
                escalation=escalation_result,
                grounded=grounded,
                citations_valid=citations_valid,
                has_safety_warnings=has_safety_warnings,
            )

        # All checks passed = GREEN
        reasons.append("High confidence, well-grounded, valid citations")
        return ClassificationResult(
            classification=OutputClassification.GREEN,
            confidence=confidence,
            reasons=reasons,
            abstention=abstention_result,
            escalation=escalation_result,
            grounded=grounded,
            citations_valid=citations_valid,
            has_safety_warnings=has_safety_warnings,
        )

    def classify_from_validation(
        self,
        validation_result: Any,  # ValidationResult from sage.rag.validation
        confidence: float,
        abstention_result: AbstentionResult | None = None,
        escalation_result: EscalationEvaluation | None = None,
    ) -> ClassificationResult:
        """
        Classify based on a ValidationResult.

        Convenience method that extracts validation fields.

        Args:
            validation_result: ValidationResult from ResponseValidator
            confidence: Retrieval confidence
            abstention_result: Optional abstention result
            escalation_result: Optional escalation result

        Returns:
            ClassificationResult
        """
        return self.classify(
            confidence=confidence,
            grounded=validation_result.grounded,
            grounding_score=validation_result.grounding_score,
            citations_valid=len(validation_result.citations_missing) == 0,
            citation_count=len(validation_result.citations_verified),
            has_safety_warnings=validation_result.has_safety_warnings,
            abstention_result=abstention_result,
            escalation_result=escalation_result,
        )


def get_classification_badge(classification: OutputClassification) -> str:
    """Get a display badge for classification.

    Args:
        classification: The output classification

    Returns:
        Badge string for display
    """
    badges = {
        OutputClassification.GREEN: "[GREEN ✓]",
        OutputClassification.YELLOW: "[YELLOW ⚠]",
        OutputClassification.RED: "[RED ✗]",
    }
    return badges.get(classification, "[UNKNOWN]")


def format_classification_notice(result: ClassificationResult) -> str:
    """
    Format a classification result as a notice.

    Args:
        result: ClassificationResult

    Returns:
        Formatted notice string
    """
    badge = get_classification_badge(result.classification)

    lines = [f"**Classification:** {badge}"]

    if result.reasons:
        lines.append("")
        for reason in result.reasons:
            lines.append(f"- {reason}")

    if result.classification == OutputClassification.YELLOW:
        lines.extend([
            "",
            "*Caution: This response should be reviewed before use in safety decisions.*"
        ])
    elif result.classification == OutputClassification.RED:
        lines.extend([
            "",
            "*This response is held for expert review and cannot be provided directly.*"
        ])

    return "\n".join(lines)
