"""
Enforcement layer system for Verity audit engine.

Implements a 5-layer enforcement hierarchy:
- Layer 0: Audit Validity (Hard GATE)
- Layer 1: Data Integrity (Default GATE)
- Layer 2: Performance Fitness (Domain-aware GATE)
- Layer 3: Fairness & Governance (Conditional GATE)
- Layer 4: Risk Signals (WARN)
- Layer 5: Telemetry (REPORT)
"""
from typing import Dict, List, Optional, Any, Tuple
from enum import Enum
from audit.metric_contract import MetricSeverity


class EnforcementLayer(str, Enum):
    """Enforcement layers in the hierarchy."""
    VALIDITY = "validity"  # Layer 0 - Hard GATE
    DATA_INTEGRITY = "data_integrity"  # Layer 1 - Default GATE
    PERFORMANCE = "performance"  # Layer 2 - Domain-aware GATE
    FAIRNESS = "fairness"  # Layer 3 - Conditional GATE
    RISK_SIGNALS = "risk_signals"  # Layer 4 - WARN
    TELEMETRY = "telemetry"  # Layer 5 - REPORT


class AuditStatus(str, Enum):
    """Overall audit status."""
    PASSED = "passed"
    PASSED_WITH_WARNINGS = "passed_with_warnings"
    FAILED = "failed"
    INCOMPLETE = "incomplete"  # Missing required gates


class MetricEnforcement:
    """Represents a metric with its enforcement layer and requirements."""
    
    def __init__(
        self,
        metric_id: str,
        layer: EnforcementLayer,
        value: Any,
        threshold: Optional[float] = None,
        passed: Optional[bool] = None,
        condition: Optional[str] = None,  # Condition for conditional gates
        description: Optional[str] = None
    ):
        self.metric_id = metric_id
        self.layer = layer
        self.value = value
        self.threshold = threshold
        self.passed = passed
        self.condition = condition
        self.description = description
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "metric_id": self.metric_id,
            "layer": self.layer.value,
            "value": self.value,
            "threshold": self.threshold,
            "passed": self.passed,
            "condition": self.condition,
            "description": self.description
        }


class EnforcementAnalyzer:
    """Analyzes metrics and assigns them to enforcement layers."""
    
    def __init__(
        self,
        dataset_info: Optional[Dict[str, Any]] = None,
        has_protected_attributes: bool = False,
        task_type: Optional[str] = None,
        domain_profile: Optional[str] = None
    ):
        self.dataset_info = dataset_info or {}
        self.has_protected_attributes = has_protected_attributes
        self.task_type = task_type
        self.domain_profile = domain_profile
        
        # High-impact domains that require fairness gates
        self.high_impact_domains = ["lending", "housing", "hiring", "credit", "employment"]
        self.is_high_impact = (
            domain_profile is not None and 
            any(domain in domain_profile.lower() for domain in self.high_impact_domains)
        )
    
    def classify_metric(
        self,
        metric_id: str,
        value: Any,
        threshold: Optional[float] = None,
        passed: Optional[bool] = None
    ) -> MetricEnforcement:
        """
        Classify a metric into an enforcement layer.
        
        Returns:
            MetricEnforcement object with layer assignment
        """
        # Layer 0: Audit Validity (Hard GATE)
        if self._is_validity_metric(metric_id):
            return MetricEnforcement(
                metric_id=metric_id,
                layer=EnforcementLayer.VALIDITY,
                value=value,
                threshold=threshold,
                passed=passed,
                description=self._get_validity_description(metric_id)
            )
        
        # Layer 1: Data Integrity (Default GATE)
        if self._is_data_integrity_metric(metric_id):
            return MetricEnforcement(
                metric_id=metric_id,
                layer=EnforcementLayer.DATA_INTEGRITY,
                value=value,
                threshold=threshold,
                passed=passed,
                description=self._get_data_integrity_description(metric_id)
            )
        
        # Layer 2: Performance Fitness (Domain-aware GATE)
        if self._is_performance_metric(metric_id):
            return MetricEnforcement(
                metric_id=metric_id,
                layer=EnforcementLayer.PERFORMANCE,
                value=value,
                threshold=threshold,
                passed=passed,
                description=self._get_performance_description(metric_id)
            )
        
        # Layer 3: Fairness & Governance (Conditional GATE)
        if self._is_fairness_metric(metric_id):
            condition = None
            if not self.has_protected_attributes:
                condition = "protected_attributes_missing"
            
            return MetricEnforcement(
                metric_id=metric_id,
                layer=EnforcementLayer.FAIRNESS,
                value=value,
                threshold=threshold,
                passed=passed,
                condition=condition,
                description=self._get_fairness_description(metric_id)
            )
        
        # Layer 4: Risk Signals (WARN)
        if self._is_risk_signal_metric(metric_id):
            return MetricEnforcement(
                metric_id=metric_id,
                layer=EnforcementLayer.RISK_SIGNALS,
                value=value,
                threshold=threshold,
                passed=passed,
                description=self._get_risk_signal_description(metric_id)
            )
        
        # Layer 5: Telemetry (REPORT) - default
        return MetricEnforcement(
            metric_id=metric_id,
            layer=EnforcementLayer.TELEMETRY,
            value=value,
            threshold=threshold,
            passed=passed,
            description="Informational metric"
        )
    
    def _is_validity_metric(self, metric_id: str) -> bool:
        """Check if metric is a validity gate."""
        validity_patterns = [
            "dataset_exists",
            "model_loads",
            "label_exists",
            "feature_schema_matches",
            "dataset_size_minimum",
            "data_quality_schema_valid"
        ]
        return any(pattern in metric_id.lower() for pattern in validity_patterns)
    
    def _is_data_integrity_metric(self, metric_id: str) -> bool:
        """Check if metric is a data integrity gate."""
        integrity_patterns = [
            "missing_value_rate_max",
            "train_test_leakage",
            "label_distribution_degenerate",
            "schema_validation",
            "duplicate_row_rate_max"
        ]
        return any(pattern in metric_id.lower() for pattern in integrity_patterns)
    
    def _is_performance_metric(self, metric_id: str) -> bool:
        """Check if metric is a performance gate."""
        performance_patterns = [
            "performance_overall_nrmse_by_std",  # Primary gate - normalized by std
            "performance_overall_nmae",
            "performance_overall_r2",
            "performance_overall_accuracy",
            "performance_overall_f1",
            "performance_overall_precision",
            "performance_overall_recall"
        ]
        # Exclude raw MAE/RMSE and nrmse_by_mean - these are report-only
        if any(excluded in metric_id.lower() for excluded in ["nrmse_by_mean", "mae", "rmse"]):
            # But allow nmae (normalized MAE) which is a gate
            if "nmae" not in metric_id.lower():
                return False
        return any(pattern in metric_id.lower() for pattern in performance_patterns)
    
    def _is_fairness_metric(self, metric_id: str) -> bool:
        """Check if metric is a fairness gate."""
        fairness_patterns = [
            "demographic_parity",
            "equal_opportunity",
            "disparate_impact",
            "fairness",
            "bias",
            "parity_diff",
            "opportunity_diff",
            "impact_ratio"
        ]
        return any(pattern in metric_id.lower() for pattern in fairness_patterns)
    
    def _is_risk_signal_metric(self, metric_id: str) -> bool:
        """Check if metric is a risk signal."""
        risk_patterns = [
            "dataset_too_small",
            "protected_attributes_missing",
            "proxy_risk",
            "explainability_instability",
            "distribution_shift",
            "proxy_correlation",
            "sensitive_leakage"
        ]
        return any(pattern in metric_id.lower() for pattern in risk_patterns)
    
    def _get_validity_description(self, metric_id: str) -> str:
        """Get description for validity metric."""
        descriptions = {
            "dataset_exists": "Dataset file must exist",
            "model_loads": "Model must load successfully",
            "label_exists": "Label column must exist in dataset",
            "feature_schema_matches": "Feature schema must match model expectations",
            "dataset_size_minimum": "Dataset must have minimum required rows",
            "data_quality_schema_valid": "Dataset schema must be valid"
        }
        for key, desc in descriptions.items():
            if key in metric_id.lower():
                return desc
        return "Audit validity check"
    
    def _get_data_integrity_description(self, metric_id: str) -> str:
        """Get description for data integrity metric."""
        if "missing_value" in metric_id.lower():
            return "Missing value rate must be within acceptable limits"
        if "leakage" in metric_id.lower():
            return "Train/test leakage must not be present"
        if "label_distribution" in metric_id.lower():
            return "Label distribution must not be degenerate"
        return "Data integrity check"
    
    def _get_performance_description(self, metric_id: str) -> str:
        """Get description for performance metric."""
        if "nrmse" in metric_id.lower():
            return "Normalized RMSE must meet threshold"
        if "nmae" in metric_id.lower():
            return "Normalized MAE must meet threshold"
        if "r2" in metric_id.lower():
            return "R² (explained variance) must meet threshold"
        if "accuracy" in metric_id.lower():
            return "Accuracy must meet threshold"
        return "Performance fitness check"
    
    def _get_fairness_description(self, metric_id: str) -> str:
        """Get description for fairness metric."""
        if "demographic_parity" in metric_id.lower():
            return "Demographic parity must meet threshold"
        if "equal_opportunity" in metric_id.lower():
            return "Equal opportunity must meet threshold"
        if "disparate_impact" in metric_id.lower():
            return "Disparate impact ratio must meet threshold"
        return "Fairness check"
    
    def _get_risk_signal_description(self, metric_id: str) -> str:
        """Get description for risk signal."""
        if "dataset_too_small" in metric_id.lower():
            return "Dataset size may be insufficient for stable statistics"
        if "protected_attributes_missing" in metric_id.lower():
            return "Protected attributes missing - fairness metrics not computed"
        if "proxy" in metric_id.lower():
            return "Proxy feature detected - may leak protected attributes"
        return "Risk signal"
    
    def check_dataset_size(self) -> Tuple[List[MetricEnforcement], bool]:
        """
        Check dataset size with tiered minimums and return appropriate enforcement.
        
        Tiered rules:
        - Hard minimum to run: >= 10 rows (otherwise FAIL - invalid audit)
        - Minimum for performance gates: >= 30 rows (otherwise WARN)
        - Minimum for fairness: >= 200 overall and >= 30 per group (otherwise WARN + skip)
        
        Returns:
            (List[MetricEnforcement], is_failure: bool)
        """
        dataset_size = self.dataset_info.get('shape', {}).get('rows', 0)
        checks = []
        is_failure = False
        
        # Hard minimum to run anything: >= 10 rows
        if dataset_size < 10:
            checks.append(MetricEnforcement(
                metric_id="dataset_rows_min_run",
                layer=EnforcementLayer.VALIDITY,
                value=dataset_size,
                threshold=10,
                passed=False,
                description=f"Dataset has only {dataset_size} rows - below hard minimum of 10 required to run audit"
            ))
            is_failure = True
        else:
            checks.append(MetricEnforcement(
                metric_id="dataset_rows_min_run",
                layer=EnforcementLayer.VALIDITY,
                value=dataset_size,
                threshold=10,
                passed=True,
                description="Dataset meets minimum size requirement to run audit"
            ))
        
        # Minimum for reliable performance metrics: >= 30 rows
        if dataset_size < 30:
            checks.append(MetricEnforcement(
                metric_id="dataset_rows_min_reliable",
                layer=EnforcementLayer.RISK_SIGNALS,
                value=dataset_size,
                threshold=30,
                passed=False,
                description=f"Dataset has only {dataset_size} rows - below recommended minimum of 30 for reliable performance metrics. Results may be statistically unstable."
            ))
        else:
            checks.append(MetricEnforcement(
                metric_id="dataset_rows_min_reliable",
                layer=EnforcementLayer.TELEMETRY,
                value=dataset_size,
                threshold=30,
                passed=True,
                description=f"Dataset has {dataset_size} rows - adequate for reliable performance metrics"
            ))
        
        # Minimum for fairness: >= 200 overall
        if dataset_size < 200:
            checks.append(MetricEnforcement(
                metric_id="dataset_rows_min_fairness",
                layer=EnforcementLayer.RISK_SIGNALS,
                value=dataset_size,
                threshold=200,
                passed=False,
                description=f"Dataset has only {dataset_size} rows - below minimum of 200 required for reliable fairness metrics. Fairness checks will be skipped."
            ))
        else:
            checks.append(MetricEnforcement(
                metric_id="dataset_rows_min_fairness",
                layer=EnforcementLayer.TELEMETRY,
                value=dataset_size,
                threshold=200,
                passed=True,
                description=f"Dataset has {dataset_size} rows - adequate for fairness analysis"
            ))
        
        return checks, is_failure
    
    def check_protected_attributes(self, column_names: Optional[List[str]] = None) -> Optional[MetricEnforcement]:
        """
        Check if protected attributes are missing when they should be present.
        Provides actionable recommendations.
        
        Returns:
            MetricEnforcement if warning needed, None otherwise
        """
        if not self.has_protected_attributes:
            # Detect candidate columns that might be protected attributes
            candidate_columns = []
            if column_names:
                # Common protected attribute patterns
                protected_patterns = ['age', 'gender', 'sex', 'race', 'ethnicity', 'religion', 
                                    'disability', 'marital', 'nationality', 'country']
                # Common proxy patterns
                proxy_patterns = ['zip', 'postal', 'city', 'state', 'region', 'county', 'area_code']
                
                for col in column_names:
                    col_lower = col.lower()
                    if any(pattern in col_lower for pattern in protected_patterns):
                        candidate_columns.append(col)
                    elif any(pattern in col_lower for pattern in proxy_patterns):
                        candidate_columns.append(col)
            
            if self.is_high_impact:
                # High-impact domain without protected attributes → WARN with recommendations
                recommendation = ""
                if candidate_columns:
                    recommendation = f"\n\nDetected candidate columns: {', '.join(candidate_columns[:5])}{'...' if len(candidate_columns) > 5 else ''}"
                    recommendation += "\n\nTo enable fairness checks, add to verity.yml:"
                    recommendation += "\n  protected_columns: [age, race, gender]  # Add your actual protected columns"
                    if any('zip' in c.lower() or 'postal' in c.lower() for c in candidate_columns):
                        recommendation += "\n  proxy_columns: [zip_code]  # Optional: detect proxy risk"
                else:
                    recommendation = "\n\nTo enable fairness checks, add to verity.yml:"
                    recommendation += "\n  protected_columns: [age, race, gender]  # Add your actual protected columns"
                    recommendation += "\n  proxy_columns: [zip_code]  # Optional: detect proxy risk on geographic columns"
                
                recommendation += "\n\nFairness requirements:"
                recommendation += "\n  - Minimum 200 total rows"
                recommendation += "\n  - Minimum 30 rows per protected group"
                
                return MetricEnforcement(
                    metric_id="protected_attributes_missing_high_impact",
                    layer=EnforcementLayer.RISK_SIGNALS,
                    value=True,
                    threshold=False,
                    passed=False,
                    description=f"Protected attributes missing in high-impact domain ({self.domain_profile}) - fairness metrics not computed.{recommendation}"
                )
            else:
                # Non-high-impact domain → WARN but less critical
                return MetricEnforcement(
                    metric_id="protected_attributes_missing",
                    layer=EnforcementLayer.RISK_SIGNALS,
                    value=True,
                    threshold=False,
                    passed=False,
                    description="Protected attributes missing - fairness metrics not computed. Add 'protected_columns' to verity.yml to enable fairness checks."
                )
        return None
    
    def check_fairness_prerequisites(
        self,
        dataset_size: int,
        protected_columns: Optional[List[str]] = None,
        group_sizes: Optional[Dict[str, Dict[str, int]]] = None
    ) -> Tuple[bool, List[MetricEnforcement]]:
        """
        Check fairness prerequisites and return enforcement checks.
        
        Prerequisites:
        - protected_columns present
        - dataset_size >= 200
        - each group has >= 30 rows
        
        Returns:
            (fairness_enabled: bool, checks: List[MetricEnforcement])
        """
        checks = []
        fairness_enabled = True
        
        # Check if protected columns are present
        if not protected_columns or len(protected_columns) == 0:
            checks.append(MetricEnforcement(
                metric_id="fairness_prereq_protected_columns",
                layer=EnforcementLayer.RISK_SIGNALS,
                value=False,
                threshold=True,
                passed=False,
                description="Fairness checks require protected_columns to be specified in config"
            ))
            fairness_enabled = False
        
        # Check total dataset size for fairness
        if dataset_size < 200:
            checks.append(MetricEnforcement(
                metric_id="fairness_prereq_total_rows",
                layer=EnforcementLayer.RISK_SIGNALS,
                value=dataset_size,
                threshold=200,
                passed=False,
                description=f"Fairness requires minimum 200 total rows (found {dataset_size}). Fairness checks will be skipped."
            ))
            fairness_enabled = False
        else:
            checks.append(MetricEnforcement(
                metric_id="fairness_prereq_total_rows",
                layer=EnforcementLayer.TELEMETRY,
                value=dataset_size,
                threshold=200,
                passed=True,
                description=f"Dataset has {dataset_size} rows - meets minimum for fairness analysis"
            ))
        
        # Check minimum group sizes
        if group_sizes and protected_columns:
            min_group_size = 30
            for col in protected_columns:
                if col in group_sizes:
                    groups = group_sizes[col]
                    min_size = min(groups.values()) if groups else 0
                    if min_size < min_group_size:
                        checks.append(MetricEnforcement(
                            metric_id=f"fairness_prereq_min_group_size_{col}",
                            layer=EnforcementLayer.RISK_SIGNALS,
                            value=min_size,
                            threshold=min_group_size,
                            passed=False,
                            description=f"Protected column '{col}' has groups with < {min_group_size} rows (minimum: {min_size}). Fairness checks may be unreliable."
                        ))
                        fairness_enabled = False
                    else:
                        checks.append(MetricEnforcement(
                            metric_id=f"fairness_prereq_min_group_size_{col}",
                            layer=EnforcementLayer.TELEMETRY,
                            value=min_size,
                            threshold=min_group_size,
                            passed=True,
                            description=f"Protected column '{col}' meets minimum group size requirement ({min_size} >= {min_group_size})"
                        ))
        
        return fairness_enabled, checks
    
    def check_minimum_gates(
        self,
        gate_metrics: Dict[str, MetricEnforcement]
    ) -> Tuple[bool, List[str]]:
        """
        Check if minimum required gates are present.
        
        Requirements:
        - At least 1 validity gate
        - At least 1 data quality gate
        - At least 1 performance gate
        
        Returns:
            (meets_requirements: bool, missing_requirements: List[str])
        """
        missing = []
        
        # Check validity gates
        has_validity = any(
            m.layer == EnforcementLayer.VALIDITY 
            for m in gate_metrics.values()
        )
        if not has_validity:
            missing.append("validity gate")
        
        # Check data integrity gates
        has_data_integrity = any(
            m.layer == EnforcementLayer.DATA_INTEGRITY 
            for m in gate_metrics.values()
        )
        if not has_data_integrity:
            missing.append("data integrity gate")
        
        # Check performance gates
        has_performance = any(
            m.layer == EnforcementLayer.PERFORMANCE 
            for m in gate_metrics.values()
        )
        if not has_performance:
            missing.append("performance gate")
        
        return len(missing) == 0, missing
    
    def determine_overall_status(
        self,
        gate_metrics: Dict[str, MetricEnforcement],
        warn_metrics: Dict[str, MetricEnforcement]
    ) -> Tuple[AuditStatus, float, str]:
        """
        Determine overall audit status based on enforcement layers.
        
        Returns:
            (status: AuditStatus, confidence: float, reason: str)
        """
        # Check minimum gate requirements
        meets_requirements, missing = self.check_minimum_gates(gate_metrics)
        if not meets_requirements:
            return (
                AuditStatus.INCOMPLETE,
                0.0,
                f"Audit incomplete - missing required gates: {', '.join(missing)}"
            )
        
        # Check if any validity gates failed (Layer 0)
        validity_failed = any(
            not m.passed 
            for m in gate_metrics.values() 
            if m.layer == EnforcementLayer.VALIDITY and m.passed is not None
        )
        if validity_failed:
            return (
                AuditStatus.FAILED,
                0.0,
                "Audit validity checks failed - audit is invalid"
            )
        
        # Check if any data integrity gates failed (Layer 1)
        data_integrity_failed = any(
            not m.passed 
            for m in gate_metrics.values() 
            if m.layer == EnforcementLayer.DATA_INTEGRITY and m.passed is not None
        )
        if data_integrity_failed:
            return (
                AuditStatus.FAILED,
                0.0,
                "Data integrity checks failed"
            )
        
        # Check if any performance gates failed (Layer 2)
        performance_failed = any(
            not m.passed 
            for m in gate_metrics.values() 
            if m.layer == EnforcementLayer.PERFORMANCE and m.passed is not None
        )
        if performance_failed:
            return (
                AuditStatus.FAILED,
                0.0,
                "Performance fitness checks failed"
            )
        
        # Check if any fairness gates failed (Layer 3) - only if protected attributes exist
        fairness_failed = False
        if self.has_protected_attributes:
            fairness_failed = any(
                not m.passed 
                for m in gate_metrics.values() 
                if m.layer == EnforcementLayer.FAIRNESS and m.passed is not None
            )
            if fairness_failed:
                return (
                    AuditStatus.FAILED,
                    0.0,
                    "Fairness checks failed"
                )
        
        # Check for warnings
        has_warnings = len(warn_metrics) > 0 and any(
            not m.passed 
            for m in warn_metrics.values() 
            if m.passed is not None
        )
        
        # Calculate confidence based on warnings
        confidence = 1.0
        if has_warnings:
            # Reduce confidence based on number and severity of warnings
            warning_count = sum(
                1 for m in warn_metrics.values() 
                if m.passed is not None and not m.passed
            )
            confidence = max(0.5, 1.0 - (warning_count * 0.1))
        
        if has_warnings:
            return (
                AuditStatus.PASSED_WITH_WARNINGS,
                confidence,
                f"Audit passed with {warning_count} warning(s)"
            )
        
        return (
            AuditStatus.PASSED,
            1.0,
            "All gates passed"
        )
