"""
Domain-aware enforcement policy for Verity audit engine.

Defines fairness requirements and strictness levels for different domains.
"""
from typing import Dict, Any, Optional
from enum import Enum


class FairnessRequirement(str, Enum):
    """Fairness requirement levels."""
    REQUIRED = "required"  # FAIL if missing protected attributes
    RECOMMENDED = "recommended"  # WARN if missing protected attributes
    OPTIONAL = "optional"  # No enforcement


class DomainProfile:
    """Domain profile with fairness requirements."""
    
    def __init__(
        self,
        domain_id: str,
        domain_name: str,
        fairness_requirement: FairnessRequirement,
        strict_fairness_required: bool = False,
        description: str = ""
    ):
        self.domain_id = domain_id
        self.domain_name = domain_name
        self.fairness_requirement = fairness_requirement
        self.strict_fairness_required = strict_fairness_required
        self.description = description


# Domain profile mappings
DOMAIN_PROFILES: Dict[str, DomainProfile] = {
    # Housing domain
    'housing_valuation': DomainProfile(
        domain_id='housing_valuation',
        domain_name='Housing Valuation',
        fairness_requirement=FairnessRequirement.RECOMMENDED,
        strict_fairness_required=False,
        description='Fairness is recommended for housing valuation models'
    ),
    'housing_lending': DomainProfile(
        domain_id='housing_lending',
        domain_name='Housing Lending',
        fairness_requirement=FairnessRequirement.REQUIRED,
        strict_fairness_required=True,
        description='Fairness is required for housing lending decisions (ECOA, FHA compliance)'
    ),
    
    # Lending domain
    'lending_approval': DomainProfile(
        domain_id='lending_approval',
        domain_name='Lending Approval',
        fairness_requirement=FairnessRequirement.REQUIRED,
        strict_fairness_required=True,
        description='Fairness is required for lending approval decisions (ECOA compliance)'
    ),
    'lending_pricing': DomainProfile(
        domain_id='lending_pricing',
        domain_name='Lending Pricing',
        fairness_requirement=FairnessRequirement.REQUIRED,
        strict_fairness_required=True,
        description='Fairness is required for lending pricing decisions (ECOA compliance)'
    ),
    
    # Hiring domain
    'hiring_selection': DomainProfile(
        domain_id='hiring_selection',
        domain_name='Hiring Selection',
        fairness_requirement=FairnessRequirement.REQUIRED,
        strict_fairness_required=True,
        description='Fairness is required for hiring decisions (Title VII compliance)'
    ),
    'hiring_ranking': DomainProfile(
        domain_id='hiring_ranking',
        domain_name='Hiring Ranking',
        fairness_requirement=FairnessRequirement.RECOMMENDED,
        strict_fairness_required=False,
        description='Fairness is recommended for hiring ranking systems'
    ),
    
    # Healthcare domain
    'healthcare_diagnosis': DomainProfile(
        domain_id='healthcare_diagnosis',
        domain_name='Healthcare Diagnosis',
        fairness_requirement=FairnessRequirement.RECOMMENDED,
        strict_fairness_required=False,
        description='Fairness is recommended for healthcare diagnosis models'
    ),
    'healthcare_treatment': DomainProfile(
        domain_id='healthcare_treatment',
        domain_name='Healthcare Treatment',
        fairness_requirement=FairnessRequirement.REQUIRED,
        strict_fairness_required=True,
        description='Fairness is required for healthcare treatment decisions'
    ),
    
    # General/default
    'general': DomainProfile(
        domain_id='general',
        domain_name='General',
        fairness_requirement=FairnessRequirement.OPTIONAL,
        strict_fairness_required=False,
        description='General purpose model - fairness is optional'
    )
}


def get_domain_profile(domain_profile_id: str) -> Optional[DomainProfile]:
    """Get domain profile by ID."""
    return DOMAIN_PROFILES.get(domain_profile_id, DOMAIN_PROFILES['general'])


def check_fairness_requirement(
    domain_profile_id: str,
    has_protected_columns: bool,
    audit_intent: str = 'performance_only'
) -> Dict[str, Any]:
    """
    Check fairness requirement for a domain and return enforcement decision.
    
    Args:
        domain_profile_id: Domain profile ID
        has_protected_columns: Whether protected columns are provided
        audit_intent: Audit intent (performance_only, fairness_enforced, telemetry_only)
    
    Returns:
        Dictionary with:
        - requirement_level: 'required', 'recommended', or 'optional'
        - should_fail: Whether to fail the audit
        - should_warn: Whether to warn
        - message: Human-readable message
    """
    profile = get_domain_profile(domain_profile_id)
    
    # If audit_intent is fairness_enforced, user explicitly wants fairness
    if audit_intent == 'fairness_enforced':
        if not has_protected_columns:
            return {
                'requirement_level': 'required',
                'should_fail': True,
                'should_warn': False,
                'message': 'audit_intent=fairness_enforced requires protected_columns'
            }
        else:
            return {
                'requirement_level': 'required',
                'should_fail': False,
                'should_warn': False,
                'message': 'Fairness enforcement enabled with protected columns'
            }
    
    # If audit_intent is performance_only or telemetry_only, fairness is optional
    if audit_intent in ['performance_only', 'telemetry_only']:
        return {
            'requirement_level': 'optional',
            'should_fail': False,
            'should_warn': False,
            'message': f'Audit intent is {audit_intent} - fairness is optional'
        }
    
    # Check domain-specific requirements
    if profile.fairness_requirement == FairnessRequirement.REQUIRED:
        if not has_protected_columns:
            if profile.strict_fairness_required:
                return {
                    'requirement_level': 'required',
                    'should_fail': True,
                    'should_warn': False,
                    'message': f'{profile.domain_name} requires protected columns for fairness analysis (strict mode)'
                }
            else:
                return {
                    'requirement_level': 'required',
                    'should_fail': False,
                    'should_warn': True,
                    'message': f'{profile.domain_name} requires protected columns for fairness analysis (recommended mode)'
                }
        else:
            return {
                'requirement_level': 'required',
                'should_fail': False,
                'should_warn': False,
                'message': f'{profile.domain_name} fairness requirements met'
            }
    
    elif profile.fairness_requirement == FairnessRequirement.RECOMMENDED:
        if not has_protected_columns:
            return {
                'requirement_level': 'recommended',
                'should_fail': False,
                'should_warn': True,
                'message': f'{profile.domain_name} recommends protected columns for fairness analysis'
            }
        else:
            return {
                'requirement_level': 'recommended',
                'should_fail': False,
                'should_warn': False,
                'message': f'{profile.domain_name} fairness recommendations met'
            }
    
    else:  # OPTIONAL
        return {
            'requirement_level': 'optional',
            'should_fail': False,
            'should_warn': False,
            'message': f'{profile.domain_name} - fairness is optional'
        }
