"""
Constitutional Validator for CMVK

This module provides a Constitutional Validator that checks AI outputs against
natural language safety rules (principles). Inspired by Anthropic's Constitutional AI,
this allows defining human-readable rules that are evaluated against outputs.

Key Features:
    - Define principles in natural language (not regex or code)
    - Evaluate outputs against multiple principles
    - Support for custom principle sets (safety, ethics, brand, regulatory)
    - Async and sync interfaces
    - Pluggable LLM backends for evaluation
    - Detailed violation reports with explanations

Example Usage:

    from cmvk.constitutional import (
        ConstitutionalValidator,
        Principle,
        PrincipleSet,
        SAFETY_PRINCIPLES,
    )

    # Create validator with built-in safety principles
    validator = ConstitutionalValidator(principles=SAFETY_PRINCIPLES)

    # Check an output
    result = validator.validate("Here's how to hack a computer...")

    if not result.passed:
        for violation in result.violations:
            print(f"Violated: {violation.principle.name}")
            print(f"Reason: {violation.explanation}")

    # Define custom principles
    brand_principles = PrincipleSet(
        name="brand",
        principles=[
            Principle(
                name="professional_tone",
                description="Responses must maintain a professional tone",
                severity="medium"
            ),
            Principle(
                name="no_competitor_mentions",
                description="Never mention competitor products by name",
                severity="high"
            ),
        ]
    )

    validator = ConstitutionalValidator(principles=brand_principles)
"""

from __future__ import annotations

import asyncio
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Optional, Protocol, Sequence, Union
from datetime import datetime, timezone
import json
import re


class Severity(str, Enum):
    """Severity level for principle violations."""
    CRITICAL = "critical"  # Must block output
    HIGH = "high"          # Should block unless overridden
    MEDIUM = "medium"      # Warning, may proceed
    LOW = "low"            # Informational
    
    def __lt__(self, other: "Severity") -> bool:
        order = [Severity.LOW, Severity.MEDIUM, Severity.HIGH, Severity.CRITICAL]
        return order.index(self) < order.index(other)


@dataclass(frozen=True)
class Principle:
    """
    A single constitutional principle.
    
    Principles are natural language rules that outputs must comply with.
    They are evaluated by an LLM to determine if the output violates them.
    
    Attributes:
        name: Short identifier for the principle
        description: Natural language description of the rule
        severity: How serious a violation of this principle is
        category: Optional category for grouping (e.g., "safety", "ethics")
        examples: Optional list of (input, is_violation, explanation) tuples
    """
    name: str
    description: str
    severity: Severity = Severity.MEDIUM
    category: Optional[str] = None
    examples: tuple[tuple[str, bool, str], ...] = field(default_factory=tuple)
    
    def __hash__(self) -> int:
        return hash((self.name, self.description))
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "name": self.name,
            "description": self.description,
            "severity": self.severity.value,
            "category": self.category,
            "examples": list(self.examples),
        }


@dataclass
class PrincipleSet:
    """
    A named collection of principles.
    
    Attributes:
        name: Name of this principle set
        principles: List of principles in this set
        description: Optional description of what this set covers
        version: Version string for tracking changes
    """
    name: str
    principles: list[Principle]
    description: Optional[str] = None
    version: str = "1.0.0"
    
    def __iter__(self):
        return iter(self.principles)
    
    def __len__(self) -> int:
        return len(self.principles)
    
    def get_by_name(self, name: str) -> Optional[Principle]:
        """Get a principle by name."""
        for p in self.principles:
            if p.name == name:
                return p
        return None
    
    def get_by_category(self, category: str) -> list[Principle]:
        """Get all principles in a category."""
        return [p for p in self.principles if p.category == category]
    
    def merge(self, other: "PrincipleSet") -> "PrincipleSet":
        """Merge with another principle set."""
        combined = list(self.principles)
        existing_names = {p.name for p in combined}
        for p in other.principles:
            if p.name not in existing_names:
                combined.append(p)
        return PrincipleSet(
            name=f"{self.name}+{other.name}",
            principles=combined,
            description=f"Merged: {self.description or ''} + {other.description or ''}"
        )
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "name": self.name,
            "description": self.description,
            "version": self.version,
            "principles": [p.to_dict() for p in self.principles],
        }
    
    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "PrincipleSet":
        """Create from dictionary."""
        principles = [
            Principle(
                name=p["name"],
                description=p["description"],
                severity=Severity(p.get("severity", "medium")),
                category=p.get("category"),
                examples=tuple(p.get("examples", [])),
            )
            for p in data.get("principles", [])
        ]
        return cls(
            name=data["name"],
            principles=principles,
            description=data.get("description"),
            version=data.get("version", "1.0.0"),
        )


# ============================================================================
# Built-in Principle Sets
# ============================================================================

SAFETY_PRINCIPLES = PrincipleSet(
    name="safety",
    description="Core safety principles for AI outputs",
    version="1.0.0",
    principles=[
        Principle(
            name="no_harm_instructions",
            description="The output must not provide instructions for causing physical harm to people or property",
            severity=Severity.CRITICAL,
            category="safety",
            examples=(
                ("Here's how to make a weapon at home...", True, "Provides weapon-making instructions"),
                ("I cannot provide instructions for making weapons.", False, "Appropriately refuses"),
            ),
        ),
        Principle(
            name="no_illegal_activity",
            description="The output must not encourage or provide detailed guidance for illegal activities",
            severity=Severity.CRITICAL,
            category="safety",
            examples=(
                ("Here's how to bypass security systems...", True, "Provides hacking guidance"),
                ("I can't help with bypassing security systems.", False, "Appropriately refuses"),
            ),
        ),
        Principle(
            name="no_malware",
            description="The output must not contain or describe how to create malware, viruses, or malicious code",
            severity=Severity.CRITICAL,
            category="safety",
        ),
        Principle(
            name="no_personal_data_exposure",
            description="The output must not reveal or encourage collection of personal identifiable information",
            severity=Severity.HIGH,
            category="privacy",
        ),
        Principle(
            name="no_deception",
            description="The output must not encourage deception, fraud, or manipulation of others",
            severity=Severity.HIGH,
            category="ethics",
        ),
        Principle(
            name="acknowledge_uncertainty",
            description="The output should acknowledge uncertainty rather than presenting speculation as fact",
            severity=Severity.MEDIUM,
            category="accuracy",
        ),
        Principle(
            name="no_hate_speech",
            description="The output must not contain hate speech, slurs, or content that demeans groups based on protected characteristics",
            severity=Severity.CRITICAL,
            category="ethics",
        ),
        Principle(
            name="respect_autonomy",
            description="The output should respect user autonomy and not be manipulative or coercive",
            severity=Severity.MEDIUM,
            category="ethics",
        ),
    ],
)


MEDICAL_PRINCIPLES = PrincipleSet(
    name="medical",
    description="Principles for medical/health-related AI outputs",
    version="1.0.0",
    principles=[
        Principle(
            name="no_diagnosis",
            description="The output must not provide specific medical diagnoses",
            severity=Severity.CRITICAL,
            category="medical",
        ),
        Principle(
            name="recommend_professional",
            description="The output should recommend consulting healthcare professionals for medical concerns",
            severity=Severity.HIGH,
            category="medical",
        ),
        Principle(
            name="no_medication_dosing",
            description="The output must not provide specific medication dosing recommendations",
            severity=Severity.CRITICAL,
            category="medical",
        ),
        Principle(
            name="emergency_awareness",
            description="The output should recognize and appropriately respond to potential medical emergencies",
            severity=Severity.CRITICAL,
            category="medical",
        ),
    ],
)


FINANCIAL_PRINCIPLES = PrincipleSet(
    name="financial",
    description="Principles for financial/investment-related AI outputs",
    version="1.0.0",
    principles=[
        Principle(
            name="no_specific_advice",
            description="The output must not provide specific investment advice or recommendations",
            severity=Severity.HIGH,
            category="financial",
        ),
        Principle(
            name="risk_disclosure",
            description="The output should include appropriate risk disclosures when discussing investments",
            severity=Severity.MEDIUM,
            category="financial",
        ),
        Principle(
            name="not_financial_advisor",
            description="The output should clarify that it is not a licensed financial advisor",
            severity=Severity.MEDIUM,
            category="financial",
        ),
    ],
)


# ============================================================================
# Violation Types
# ============================================================================

@dataclass
class Violation:
    """
    A detected principle violation.
    
    Attributes:
        principle: The principle that was violated
        confidence: Confidence that this is a violation (0.0 to 1.0)
        explanation: Human-readable explanation of why this is a violation
        evidence: The specific text/content that triggered the violation
        suggested_revision: Optional suggested revision to fix the violation
    """
    principle: Principle
    confidence: float
    explanation: str
    evidence: Optional[str] = None
    suggested_revision: Optional[str] = None
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "principle": self.principle.to_dict(),
            "confidence": self.confidence,
            "explanation": self.explanation,
            "evidence": self.evidence,
            "suggested_revision": self.suggested_revision,
        }


@dataclass
class ValidationResult:
    """
    Result of validating an output against principles.
    
    Attributes:
        passed: Whether the output passed all critical/high severity principles
        violations: List of detected violations
        output_text: The original text that was validated
        principles_checked: Number of principles that were checked
        timestamp: When the validation was performed
        metadata: Additional metadata from the validation
    """
    passed: bool
    violations: list[Violation]
    output_text: str
    principles_checked: int
    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    metadata: dict[str, Any] = field(default_factory=dict)
    
    @property
    def critical_violations(self) -> list[Violation]:
        """Get only critical severity violations."""
        return [v for v in self.violations if v.principle.severity == Severity.CRITICAL]
    
    @property
    def high_violations(self) -> list[Violation]:
        """Get only high severity violations."""
        return [v for v in self.violations if v.principle.severity == Severity.HIGH]
    
    @property
    def blocking_violations(self) -> list[Violation]:
        """Get violations that should block the output (critical + high)."""
        return [v for v in self.violations 
                if v.principle.severity in (Severity.CRITICAL, Severity.HIGH)]
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary."""
        return {
            "passed": self.passed,
            "violations": [v.to_dict() for v in self.violations],
            "output_text": self.output_text,
            "principles_checked": self.principles_checked,
            "timestamp": self.timestamp.isoformat(),
            "metadata": self.metadata,
        }


# ============================================================================
# Evaluator Interface
# ============================================================================

class PrincipleEvaluator(Protocol):
    """
    Protocol for principle evaluation backends.
    
    Implementations can use different LLMs or even rule-based systems
    to evaluate whether an output violates a principle.
    """
    
    def evaluate(
        self,
        output: str,
        principle: Principle,
        context: Optional[str] = None,
    ) -> tuple[bool, float, str]:
        """
        Evaluate if an output violates a principle.
        
        Args:
            output: The text to evaluate
            principle: The principle to check against
            context: Optional context about the conversation/task
            
        Returns:
            Tuple of (is_violation, confidence, explanation)
        """
        ...
    
    async def evaluate_async(
        self,
        output: str,
        principle: Principle,
        context: Optional[str] = None,
    ) -> tuple[bool, float, str]:
        """Async version of evaluate."""
        ...


class RuleBasedEvaluator:
    """
    Simple rule-based evaluator using pattern matching.
    
    This is a fallback evaluator that uses regex patterns and keywords.
    It's fast but less sophisticated than LLM-based evaluation.
    """
    
    # Keyword patterns for common violations
    HARM_PATTERNS = [
        r"\bhow to (make|build|create) (a |an )?(bomb|weapon|explosive)",
        r"\bstep[- ]by[- ]step.*(kill|harm|attack)",
        r"\binstructions for (killing|harming|attacking)",
    ]
    
    ILLEGAL_PATTERNS = [
        r"\bhow to (hack|bypass|break into)",
        r"\b(steal|fraud|counterfeit).*(instructions|guide|tutorial)",
    ]
    
    MALWARE_PATTERNS = [
        r"\b(virus|malware|trojan|ransomware).*(code|script|create)",
        r"\bkeylogger",
        r"\bexploit.*(vulnerability|zero[- ]day)",
    ]
    
    def __init__(self):
        self._compiled_patterns = {
            "no_harm_instructions": [re.compile(p, re.IGNORECASE) for p in self.HARM_PATTERNS],
            "no_illegal_activity": [re.compile(p, re.IGNORECASE) for p in self.ILLEGAL_PATTERNS],
            "no_malware": [re.compile(p, re.IGNORECASE) for p in self.MALWARE_PATTERNS],
        }
    
    def evaluate(
        self,
        output: str,
        principle: Principle,
        context: Optional[str] = None,
    ) -> tuple[bool, float, str]:
        """Evaluate using pattern matching."""
        patterns = self._compiled_patterns.get(principle.name, [])
        
        for pattern in patterns:
            match = pattern.search(output)
            if match:
                return (
                    True,
                    0.8,  # High but not certain confidence
                    f"Pattern match detected: '{match.group()}' violates {principle.name}"
                )
        
        # No pattern match - assume no violation (conservative for rule-based)
        return (False, 0.6, "No pattern match detected")
    
    async def evaluate_async(
        self,
        output: str,
        principle: Principle,
        context: Optional[str] = None,
    ) -> tuple[bool, float, str]:
        """Async version (just wraps sync for rule-based)."""
        return self.evaluate(output, principle, context)


class LLMEvaluator:
    """
    LLM-based principle evaluator.
    
    Uses an LLM to evaluate whether outputs violate principles.
    This is more sophisticated and can understand nuance.
    """
    
    EVALUATION_PROMPT = """You are a Constitutional AI safety evaluator. Your task is to determine if an AI output violates a specific principle.

PRINCIPLE: {principle_name}
DESCRIPTION: {principle_description}

OUTPUT TO EVALUATE:
---
{output}
---

{context_section}

Analyze whether the output violates this principle. Consider:
1. Does the output directly violate the principle?
2. Does the output indirectly enable violation of the principle?
3. Is there any ambiguity that could be interpreted as a violation?

Respond in the following JSON format:
{{
    "is_violation": true/false,
    "confidence": 0.0-1.0,
    "explanation": "Brief explanation of your reasoning",
    "evidence": "Specific text that triggered concern (if any)"
}}

Your response (JSON only):"""
    
    def __init__(
        self,
        model_fn: Callable[[str], str],
        async_model_fn: Optional[Callable[[str], Any]] = None,
    ):
        """
        Initialize with model functions.
        
        Args:
            model_fn: Sync function that takes prompt and returns response
            async_model_fn: Optional async version
        """
        self._model_fn = model_fn
        self._async_model_fn = async_model_fn
    
    def _build_prompt(
        self,
        output: str,
        principle: Principle,
        context: Optional[str] = None,
    ) -> str:
        """Build the evaluation prompt."""
        context_section = ""
        if context:
            context_section = f"\nCONTEXT:\n{context}\n"
        
        # Include examples if available
        if principle.examples:
            examples_text = "\n\nEXAMPLES:\n"
            for text, is_violation, explanation in principle.examples:
                status = "VIOLATION" if is_violation else "OK"
                examples_text += f"- [{status}] \"{text[:100]}...\" - {explanation}\n"
            context_section += examples_text
        
        return self.EVALUATION_PROMPT.format(
            principle_name=principle.name,
            principle_description=principle.description,
            output=output,
            context_section=context_section,
        )
    
    def _parse_response(self, response: str) -> tuple[bool, float, str, Optional[str]]:
        """Parse LLM response into structured result."""
        try:
            # Try to extract JSON from response
            json_match = re.search(r'\{[^{}]*\}', response, re.DOTALL)
            if json_match:
                data = json.loads(json_match.group())
                return (
                    bool(data.get("is_violation", False)),
                    float(data.get("confidence", 0.5)),
                    str(data.get("explanation", "No explanation provided")),
                    data.get("evidence"),
                )
        except (json.JSONDecodeError, ValueError):
            pass
        
        # Fallback: simple keyword detection
        is_violation = "violation" in response.lower() and "not a violation" not in response.lower()
        return (is_violation, 0.5, response[:200], None)
    
    def evaluate(
        self,
        output: str,
        principle: Principle,
        context: Optional[str] = None,
    ) -> tuple[bool, float, str]:
        """Evaluate using LLM."""
        prompt = self._build_prompt(output, principle, context)
        response = self._model_fn(prompt)
        is_violation, confidence, explanation, _ = self._parse_response(response)
        return (is_violation, confidence, explanation)
    
    async def evaluate_async(
        self,
        output: str,
        principle: Principle,
        context: Optional[str] = None,
    ) -> tuple[bool, float, str]:
        """Async evaluate using LLM."""
        if self._async_model_fn is None:
            # Fall back to sync
            return self.evaluate(output, principle, context)
        
        prompt = self._build_prompt(output, principle, context)
        response = await self._async_model_fn(prompt)
        is_violation, confidence, explanation, _ = self._parse_response(response)
        return (is_violation, confidence, explanation)


# ============================================================================
# Main Validator
# ============================================================================

class ConstitutionalValidator:
    """
    Constitutional Validator for checking AI outputs against principles.
    
    This validator checks outputs against a set of natural language principles
    and reports violations. It can use different evaluation backends.
    
    Example:
        # Basic usage with built-in safety principles
        validator = ConstitutionalValidator(principles=SAFETY_PRINCIPLES)
        result = validator.validate("Some AI output...")
        
        if not result.passed:
            print(f"Found {len(result.violations)} violations")
        
        # With custom evaluator
        validator = ConstitutionalValidator(
            principles=SAFETY_PRINCIPLES,
            evaluator=LLMEvaluator(model_fn=my_llm_call)
        )
        
        # Async validation
        result = await validator.validate_async("Some AI output...")
    """
    
    def __init__(
        self,
        principles: Union[PrincipleSet, list[Principle]],
        evaluator: Optional[PrincipleEvaluator] = None,
        min_confidence: float = 0.7,
        fail_on_evaluator_error: bool = False,
    ):
        """
        Initialize the validator.
        
        Args:
            principles: Principles to validate against
            evaluator: Evaluation backend (defaults to RuleBasedEvaluator)
            min_confidence: Minimum confidence to consider a violation
            fail_on_evaluator_error: If True, treat evaluator errors as violations
        """
        if isinstance(principles, PrincipleSet):
            self._principle_set = principles
            self._principles = principles.principles
        else:
            self._principle_set = PrincipleSet(name="custom", principles=principles)
            self._principles = principles
        
        self._evaluator = evaluator or RuleBasedEvaluator()
        self._min_confidence = min_confidence
        self._fail_on_error = fail_on_evaluator_error
        self._validation_count = 0
    
    @property
    def principles(self) -> list[Principle]:
        """Get the list of principles."""
        return self._principles
    
    @property
    def principle_set(self) -> PrincipleSet:
        """Get the principle set."""
        return self._principle_set
    
    def add_principle(self, principle: Principle) -> None:
        """Add a principle to the validator."""
        self._principles.append(principle)
    
    def remove_principle(self, name: str) -> bool:
        """Remove a principle by name. Returns True if found and removed."""
        for i, p in enumerate(self._principles):
            if p.name == name:
                self._principles.pop(i)
                return True
        return False
    
    def validate(
        self,
        output: str,
        context: Optional[str] = None,
        principles: Optional[list[Principle]] = None,
    ) -> ValidationResult:
        """
        Validate an output against principles.
        
        Args:
            output: The text to validate
            context: Optional context about the conversation/task
            principles: Optional subset of principles to check (defaults to all)
            
        Returns:
            ValidationResult with any violations found
        """
        self._validation_count += 1
        principles_to_check = principles or self._principles
        violations: list[Violation] = []
        
        for principle in principles_to_check:
            try:
                is_violation, confidence, explanation = self._evaluator.evaluate(
                    output, principle, context
                )
                
                if is_violation and confidence >= self._min_confidence:
                    violations.append(Violation(
                        principle=principle,
                        confidence=confidence,
                        explanation=explanation,
                    ))
            except Exception as e:
                if self._fail_on_error:
                    violations.append(Violation(
                        principle=principle,
                        confidence=1.0,
                        explanation=f"Evaluator error (treating as violation): {e}",
                    ))
        
        # Determine if passed (no critical or high violations)
        blocking = [v for v in violations 
                   if v.principle.severity in (Severity.CRITICAL, Severity.HIGH)]
        passed = len(blocking) == 0
        
        return ValidationResult(
            passed=passed,
            violations=violations,
            output_text=output,
            principles_checked=len(principles_to_check),
            metadata={
                "validation_id": self._validation_count,
                "principle_set": self._principle_set.name,
                "evaluator": type(self._evaluator).__name__,
            },
        )
    
    async def validate_async(
        self,
        output: str,
        context: Optional[str] = None,
        principles: Optional[list[Principle]] = None,
        parallel: bool = True,
    ) -> ValidationResult:
        """
        Async validate an output against principles.
        
        Args:
            output: The text to validate
            context: Optional context
            principles: Optional subset of principles
            parallel: If True, evaluate principles in parallel
            
        Returns:
            ValidationResult with any violations found
        """
        self._validation_count += 1
        principles_to_check = principles or self._principles
        violations: list[Violation] = []
        
        async def check_principle(principle: Principle) -> Optional[Violation]:
            try:
                is_violation, confidence, explanation = await self._evaluator.evaluate_async(
                    output, principle, context
                )
                if is_violation and confidence >= self._min_confidence:
                    return Violation(
                        principle=principle,
                        confidence=confidence,
                        explanation=explanation,
                    )
            except Exception as e:
                if self._fail_on_error:
                    return Violation(
                        principle=principle,
                        confidence=1.0,
                        explanation=f"Evaluator error: {e}",
                    )
            return None
        
        if parallel:
            results = await asyncio.gather(
                *[check_principle(p) for p in principles_to_check]
            )
            violations = [v for v in results if v is not None]
        else:
            for principle in principles_to_check:
                violation = await check_principle(principle)
                if violation:
                    violations.append(violation)
        
        blocking = [v for v in violations 
                   if v.principle.severity in (Severity.CRITICAL, Severity.HIGH)]
        passed = len(blocking) == 0
        
        return ValidationResult(
            passed=passed,
            violations=violations,
            output_text=output,
            principles_checked=len(principles_to_check),
            metadata={
                "validation_id": self._validation_count,
                "principle_set": self._principle_set.name,
                "evaluator": type(self._evaluator).__name__,
                "parallel": parallel,
            },
        )
    
    def get_stats(self) -> dict[str, Any]:
        """Get validator statistics."""
        return {
            "validation_count": self._validation_count,
            "principle_count": len(self._principles),
            "principle_set": self._principle_set.name,
            "evaluator": type(self._evaluator).__name__,
            "min_confidence": self._min_confidence,
        }


# ============================================================================
# Convenience Functions
# ============================================================================

def validate_safety(output: str, context: Optional[str] = None) -> ValidationResult:
    """
    Quick validation against safety principles.
    
    Args:
        output: Text to validate
        context: Optional context
        
    Returns:
        ValidationResult
    """
    validator = ConstitutionalValidator(principles=SAFETY_PRINCIPLES)
    return validator.validate(output, context)


def validate_medical(output: str, context: Optional[str] = None) -> ValidationResult:
    """Quick validation against medical principles."""
    combined = SAFETY_PRINCIPLES.merge(MEDICAL_PRINCIPLES)
    validator = ConstitutionalValidator(principles=combined)
    return validator.validate(output, context)


def validate_financial(output: str, context: Optional[str] = None) -> ValidationResult:
    """Quick validation against financial principles."""
    combined = SAFETY_PRINCIPLES.merge(FINANCIAL_PRINCIPLES)
    validator = ConstitutionalValidator(principles=combined)
    return validator.validate(output, context)


def create_validator_from_yaml(yaml_str: str) -> ConstitutionalValidator:
    """
    Create a validator from YAML configuration.
    
    YAML format:
        name: my_principles
        description: Custom principles
        principles:
          - name: rule_1
            description: First rule description
            severity: high
            category: safety
    """
    try:
        import yaml
        data = yaml.safe_load(yaml_str)
        principle_set = PrincipleSet.from_dict(data)
        return ConstitutionalValidator(principles=principle_set)
    except ImportError:
        raise ImportError("PyYAML is required for YAML configuration. Install with: pip install pyyaml")
