"""Configuration loading and validation."""
from pathlib import Path
from typing import Any, Dict, List, Optional
import yaml
import os
import requests
from urllib.parse import urlparse


class Config:
    """Audit configuration."""
    
    def __init__(self, config_path: str):
        """Load configuration from YAML file."""
        self.config_path = Path(config_path)
        if not self.config_path.exists():
            raise FileNotFoundError(f"Config file not found: {config_path}")
        
        with open(self.config_path, 'r') as f:
            self.data = yaml.safe_load(f)
        
        # Get mode (new field, defaults to dataset_model for backward compatibility)
        self.mode = self.data.get('mode', 'dataset_model')
        self.domain_profile = self.data.get('domain_profile', 'general')
        self.jurisdiction = self.data.get('jurisdiction', 'United States')
        
        # Audit intent: performance_only | fairness_enforced | telemetry_only
        # Default to performance_only for safety
        self.audit_intent = self.data.get('audit_intent', 'performance_only')
        if self.audit_intent not in ['performance_only', 'fairness_enforced', 'telemetry_only']:
            raise ValueError(f"Invalid audit_intent: {self.audit_intent}. Must be one of: performance_only, fairness_enforced, telemetry_only")
        
        # Fairness configuration block
        fairness_config = self.data.get('fairness', {})
        self.fairness_enabled = fairness_config.get('enabled', False)
        self.fairness_protected_columns = fairness_config.get('protected_columns', [])
        if isinstance(self.fairness_protected_columns, str):
            self.fairness_protected_columns = [self.fairness_protected_columns]
        self.fairness_min_group_size = fairness_config.get('min_group_size', 30)
        self.fairness_min_groups = fairness_config.get('min_groups', 2)
        
        # UI hints (optional, for generator)
        self.ui_hints = self.data.get('ui_hints', {})
        self.ui_hints_missing_fields = self.ui_hints.get('missing_fields', [])
        self.ui_hints_suggested_fields = self.ui_hints.get('suggested_fields', {})
        
        # Data source options for dataset
        # Priority: dataset_upload_id > dataset_url > dataset_env > dataset_path
        self.dataset_upload_id = self.data.get('dataset_upload_id')
        self.dataset_url = self.data.get('dataset_url')
        self.dataset_env = self.data.get('dataset_env')
        
        # Mode-specific field handling
        if self.mode == 'dataset_model' or self.mode == 'hybrid':
            # Handle dataset: can be path, URL, or env var reference
            resolved_dataset = self._resolve_dataset()
            if resolved_dataset is None:
                raise ValueError(
                    f"Dataset is required for {self.mode} mode. "
                    f"Provide one of: 'dataset', 'dataset_url', 'dataset_env', or 'dataset_sample'"
                )
            self.dataset = resolved_dataset
            
            # Resolve model path (required for dataset_model and hybrid)
            model_path = self.data.get('model')
            if not model_path:
                raise ValueError("'model' field is required for dataset_model and hybrid modes")
            
            model_path = Path(model_path)
            if model_path.is_absolute():
                self.model = model_path
            else:
                self.model = Path.cwd() / model_path
            
            self.label = self.data.get('label', 'target')
            # Support both 'sensitive' (legacy) and 'protected_columns' (new format)
            # Priority: fairness.protected_columns > protected_columns > sensitive
            self.sensitive = (
                self.fairness_protected_columns or 
                self.data.get('protected_columns') or 
                self.data.get('sensitive', [])
            )
            # Ensure it's a list
            if isinstance(self.sensitive, str):
                self.sensitive = [self.sensitive]
            elif not isinstance(self.sensitive, list):
                self.sensitive = []
            # Extract thresholds - support both old format (top-level) and new format (nested in metrics)
            top_level_thresholds = self.data.get('thresholds', {})
            metrics_thresholds = {}
            
            # Extract thresholds from nested metrics structure (new format)
            metrics = self.data.get('metrics', {})
            if isinstance(metrics, dict):
                for metric_key, metric_config in metrics.items():
                    if isinstance(metric_config, dict):
                        threshold_config = metric_config.get('threshold')
                        if threshold_config:
                            # Support both simple value and dict with value field
                            if isinstance(threshold_config, dict):
                                threshold_value = threshold_config.get('value')
                                metric_id = threshold_config.get('metric_id', metric_key)
                                # Store both the metric_id and the threshold value
                                # This allows us to map config metric IDs to actual computed metric names
                                metrics_thresholds[metric_id] = {
                                    'value': threshold_value,
                                    'unit': threshold_config.get('unit'),
                                    'scope': threshold_config.get('scope'),
                                    'config_key': metric_key  # Keep reference to original config key
                                }
                            else:
                                threshold_value = threshold_config
                                if threshold_value is not None:
                                    # Store as dict format for consistency
                                    metrics_thresholds[metric_key] = {'value': threshold_value}
            
            # Merge: top-level thresholds take precedence, then metrics thresholds
            # For top-level thresholds, convert to dict format if needed
            merged_thresholds = {}
            for key, value in metrics_thresholds.items():
                if isinstance(value, dict):
                    merged_thresholds[key] = value
                else:
                    merged_thresholds[key] = {'value': value}
            
            # Add top-level thresholds (convert to dict format)
            for key, value in top_level_thresholds.items():
                if key not in merged_thresholds:
                    merged_thresholds[key] = {'value': value} if not isinstance(value, dict) else value
            
            self.thresholds = merged_thresholds
            self.task_type = self.data.get('task_type', None)  # binary_classification, multi_class_classification, regression
            self.positive_class = self.data.get('positive_class', None)  # For binary classification
        else:
            # API model mode - dataset and model are optional
            self.dataset = None
            self.model = None
            self.label = None
            self.sensitive = []
            self.thresholds = {}
            
            # Try to resolve dataset if provided (for evaluation)
            # Only if it's explicitly provided and not "NONE"
            dataset_value = self.data.get('dataset')
            if (self.data.get('dataset_url') or self.data.get('dataset_env') or self.data.get('dataset_sample') or 
                (dataset_value and dataset_value != 'NONE' and dataset_value != '')):
                try:
                    resolved = self._resolve_dataset()
                    if resolved is not None:
                        self.dataset = resolved
                except Exception as e:
                    # Dataset optional for API mode - only log if it was explicitly provided
                    if dataset_value and dataset_value != 'NONE' and dataset_value != '':
                        # User provided a dataset but it failed - this is fine for API mode
                        pass
        
        # API-specific fields (for api_model and hybrid modes)
        if self.mode == 'api_model' or self.mode == 'hybrid':
            self.api_provider = self.data.get('api_provider', 'unknown')
            self.api_model = self.data.get('api_model', 'unknown')
            self.api_key_env = self.data.get('api_key_env', 'API_KEY')
            self.test_suite_source = self.data.get('test_suite_source', 'tests/prompts.jsonl')
            self.protected_demographics = self.data.get('protected_demographics', [])
            self.governance_thresholds = self.data.get('governance_thresholds', {})
        else:
            self.api_provider = None
            self.api_model = None
            self.api_key_env = None
            self.test_suite_source = None
            self.protected_demographics = []
            self.governance_thresholds = {}
        
        # Dataset metadata (required for production audits)
        self.dataset_metadata = self.data.get('dataset_metadata', {})
        self.dataset_description = self.dataset_metadata.get('description', '')
        self.dataset_population = self.dataset_metadata.get('population', '')
        self.dataset_jurisdiction = self.dataset_metadata.get('jurisdiction', '')
        self.dataset_sample_size = self.dataset_metadata.get('sample_size', None)
        
        # Evidence dir is relative to current working directory
        evidence_dir_path = self.data.get('evidence_dir', 'audit_evidence')
        if Path(evidence_dir_path).is_absolute():
            self.evidence_dir = Path(evidence_dir_path)
        else:
            self.evidence_dir = Path.cwd() / evidence_dir_path
        
        # Validate paths (mode-specific)
        if self.mode == 'dataset_model' or self.mode == 'hybrid':
            self._validate_dataset()
            if self.model and not self.model.exists():
                raise FileNotFoundError(f"Model not found: {self.model}")
        elif self.mode == 'api_model':
            # For API mode, validate API configuration
            if not self.api_provider or self.api_provider == 'unknown':
                raise ValueError("'api_provider' field is required for api_model mode")
            if not self.api_key_env:
                raise ValueError("'api_key_env' field is required for api_model mode")
        
        # Validate audit intent requirements
        self._validate_audit_intent()
    
    def _resolve_dataset(self) -> Optional[Path]:
        """Resolve dataset path from config, supporting URLs, env vars, and local paths."""
        # Priority: dataset_url > dataset_env > dataset_sample > dataset
        dataset_url = self.data.get('dataset_url')
        dataset_env = self.data.get('dataset_env')
        dataset_sample = self.data.get('dataset_sample')
        dataset_path = self.data.get('dataset')
        
        # For API models, dataset is optional - if it's "NONE" or empty, return None
        if self.mode == 'api_model':
            if not dataset_url and not dataset_env and not dataset_sample:
                if not dataset_path or dataset_path == 'NONE' or dataset_path == '':
                    return None  # API models don't require dataset
        
        # Option 1: URL-based dataset
        if dataset_url:
            return self._download_dataset_from_url(dataset_url)
        
        # Option 2: Environment variable reference
        if dataset_env:
            env_value = os.getenv(dataset_env)
            if not env_value:
                raise ValueError(
                    f"Environment variable '{dataset_env}' not set. "
                    f"Set it to a dataset path or URL."
                )
            # Check if it's a URL
            if self._is_url(env_value):
                return self._download_dataset_from_url(env_value)
            # Otherwise treat as path
            dataset_path = env_value
        
        # Option 3: Sample dataset (for CI)
        if dataset_sample:
            sample_path = Path(dataset_sample)
            if sample_path.is_absolute():
                return sample_path
            return Path.cwd() / sample_path
        
        # Option 4: Regular path (original behavior)
        # For API models, if dataset is "NONE" or empty, it's optional
        if not dataset_path or (self.mode == 'api_model' and (dataset_path == 'NONE' or dataset_path == '')):
            if self.mode == 'api_model':
                return None  # API models don't require dataset
            raise ValueError("One of 'dataset', 'dataset_url', 'dataset_env', or 'dataset_sample' is required")
        
        dataset_path = Path(dataset_path)
        if dataset_path.is_absolute():
            return dataset_path
        return Path.cwd() / dataset_path
    
    def _is_url(self, value: str) -> bool:
        """Check if a string is a URL."""
        try:
            result = urlparse(value)
            return all([result.scheme, result.netloc])
        except:
            return False
    
    def _download_dataset_from_url(self, url: str) -> Path:
        """Download dataset from URL and return local path."""
        from rich.console import Console
        console = Console()
        
        # Determine local filename from URL or use default
        parsed = urlparse(url)
        filename = os.path.basename(parsed.path) or "dataset.csv"
        if not filename.endswith(('.csv', '.pkl', '.json')):
            filename = "dataset.csv"
        
        local_path = Path.cwd() / ".verity_cache" / filename
        local_path.parent.mkdir(parents=True, exist_ok=True)
        
        # Check if already downloaded (simple cache)
        if local_path.exists():
            console.print(f"[dim]Using cached dataset: {local_path}[/dim]")
            return local_path
        
        console.print(f"[cyan]Downloading dataset from {url}...[/cyan]")
        
        try:
            # Handle different URL schemes
            if url.startswith(('http://', 'https://')):
                response = requests.get(url, timeout=300, stream=True)
                response.raise_for_status()
                with open(local_path, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        f.write(chunk)
            elif url.startswith('s3://'):
                # AWS S3 - requires boto3
                try:
                    import boto3
                    s3 = boto3.client('s3')
                    bucket, key = url[5:].split('/', 1)
                    s3.download_file(bucket, key, str(local_path))
                except ImportError:
                    raise ValueError(
                        "S3 URLs require 'boto3'. Install with: pip install boto3"
                    )
            elif url.startswith('gs://'):
                # Google Cloud Storage - requires google-cloud-storage
                try:
                    from google.cloud import storage
                    client = storage.Client()
                    bucket_name, blob_name = url[5:].split('/', 1)
                    bucket = client.bucket(bucket_name)
                    blob = bucket.blob(blob_name)
                    blob.download_to_filename(str(local_path))
                except ImportError:
                    raise ValueError(
                        "GCS URLs require 'google-cloud-storage'. Install with: pip install google-cloud-storage"
                    )
            elif url.startswith('azure://'):
                # Azure Blob Storage - requires azure-storage-blob
                try:
                    from azure.storage.blob import BlobClient
                    # Format: azure://account/container/blob
                    parts = url[8:].split('/', 2)
                    if len(parts) < 3:
                        raise ValueError("Azure URL format: azure://account/container/blob")
                    account, container, blob = parts
                    # Get connection string from env or use account name
                    conn_str = os.getenv('AZURE_STORAGE_CONNECTION_STRING')
                    blob_client = BlobClient.from_connection_string(conn_str, container, blob)
                    with open(local_path, 'wb') as f:
                        blob_client.download_blob().readinto(f)
                except ImportError:
                    raise ValueError(
                        "Azure URLs require 'azure-storage-blob'. Install with: pip install azure-storage-blob"
                    )
            else:
                raise ValueError(f"Unsupported URL scheme: {url}")
            
            console.print(f"[green]Dataset downloaded to: {local_path}[/green]")
            return local_path
            
        except Exception as e:
            if local_path.exists():
                local_path.unlink()
            raise ValueError(f"Failed to download dataset from {url}: {e}")
    
    def _validate_dataset(self):
        """Validate dataset exists, with helpful error messages."""
        # Skip validation if dataset is None (valid for API models)
        if self.dataset is None:
            return
        if not self.dataset.exists():
            # Check if we're in CI environment
            is_ci = os.getenv('CI') == 'true' or os.getenv('GITHUB_ACTIONS') == 'true'
            
            # Check if file might be gitignored
            gitignore_hint = ""
            if self.config_path.parent.exists():
                gitignore_path = self.config_path.parent / '.gitignore'
                if gitignore_path.exists():
                    with open(gitignore_path, 'r') as f:
                        gitignore_content = f.read()
                        dataset_name = self.dataset.name
                        if dataset_name in gitignore_content or '*.csv' in gitignore_content:
                            gitignore_hint = (
                                f"\n\n💡 This file appears to be in .gitignore. "
                                f"In CI/CD, gitignored files aren't available.\n"
                                f"Solutions:\n"
                                f"  1. Use 'dataset_url' in verity.yml with a secure URL (S3, GCS, etc.)\n"
                                f"  2. Use 'dataset_env' to reference a GitHub Secret containing the URL\n"
                                f"  3. Use 'dataset_sample' for a small sample dataset committed to the repo\n"
                                f"  4. Download the dataset in your workflow before running verity\n"
                            )
            
            error_msg = f"Dataset not found: {self.dataset}"
            if is_ci:
                error_msg += (
                    "\n\n⚠️  Running in CI/CD environment. "
                    "Files in .gitignore are not available in GitHub Actions."
                )
            error_msg += gitignore_hint
            
            error_msg += (
                "\n\nFor more information, see: "
                "https://docs.verity.com/datasets-in-ci"
            )
            
            raise FileNotFoundError(error_msg)
    
    def get_threshold(self, metric_name: str):
        """Get threshold for a metric, or None if not specified.
        
        Returns the threshold value (float) or a dict with 'value', 'unit', 'scope' if available.
        """
        threshold = self.thresholds.get(metric_name)
        if threshold is None:
            return None
        if isinstance(threshold, dict):
            return threshold.get('value') if 'value' in threshold else threshold
        return threshold
    
    def find_threshold_for_metric(self, computed_metric_name: str):
        """Find threshold for a computed metric by mapping config metric IDs to computed metric patterns.
        
        This handles the mismatch between config metric IDs (e.g., 'mean_prediction_error_by_group')
        and actual computed metric names (e.g., 'fairness_group_metrics_gender_gender_mae_diff').
        
        Returns tuple (threshold_value, threshold_info_dict) or (None, None) if not found.
        """
        # EXCLUDE overall performance metrics - these are informational only, never have thresholds
        if computed_metric_name.startswith('performance_overall_'):
            return None, None
        
        # First, try direct match
        threshold = self.thresholds.get(computed_metric_name)
        if threshold is not None:
            if isinstance(threshold, dict):
                return threshold.get('value'), threshold
            return threshold, None
        
        # Mapping from config metric IDs to patterns/keywords that match computed metric names
        # This handles the mismatch between config metric IDs and actual computed metric names
        metric_mappings = {
            # Regression fairness metrics
            'mean_prediction_error_by_group': [
                'group_metrics',  # From fairness.group_metrics.{attr}
                'error_bias',  # Matches {attr}_error_bias from group_metrics
                'performance_by_group',  # From performance.by_group.{attr}
                # Must NOT match 'performance_overall' - exclude overall metrics
            ],
            'mae_rmse_by_protected_group': [
                'mae_diff',  # Matches {attr}_mae_diff from fairness.group_metrics
                'rmse_diff',  # Matches {attr}_rmse_diff from fairness.group_metrics
                'performance_by_group',  # Matches performance.by_group.{attr}
                # Must NOT match 'performance_overall' - exclude overall metrics
            ],
            'over_under_estimation_bias': [
                'error_bias',  # Matches error bias metrics
            ],
            'residual_distribution_parity': [
                'error_90th_percentile_diff',  # Matches 90th percentile error differences
                'residual_parity',  # Matches residual parity metrics
            ],
            'worst_group_error_analysis': [
                'worst_group_gap',  # Matches worst group gap metrics
                'worst_case_slices',  # Matches worst case slice metrics
            ],
            'performance_by_group': [
                'performance_by_group',  # Direct match for per-group performance
                # But NOT 'performance_overall' - that's a different metric
            ],
            # Normalized regression metrics (for thresholding)
            # Separate mappings for nrmse_by_std and nrmse_by_mean
            'nrmse_by_std': [
                'nrmse_by_std',  # Matches normalized RMSE by std
                'performance_overall_nrmse_by_std',  # Matches flattened performance_overall.nrmse_by_std
            ],
            'nrmse_by_mean': [
                'nrmse_by_mean',  # Matches normalized RMSE by mean
                'performance_overall_nrmse_by_mean',  # Matches flattened performance_overall.nrmse_by_mean
            ],
            # Legacy support for nrmse_max (maps to both)
            'nrmse_max': [
                'nrmse_by_std',  # Matches normalized RMSE by std
                'nrmse_by_mean',  # Matches normalized RMSE by mean
                'performance_overall_nrmse_by_std',  # Matches flattened performance_overall.nrmse_by_std
                'performance_overall_nrmse_by_mean',  # Matches flattened performance_overall.nrmse_by_mean
                'performance_overall_nrmse',  # Matches any performance_overall_nrmse variant
            ],
            'nmae_max': [
                'nmae',  # Matches normalized MAE
                'performance_overall_nmae',  # Matches flattened performance_overall.nmae
            ],
            'mape_max': [
                'mape',  # Matches mean absolute percentage error
                'performance_overall_mape',  # Matches flattened performance_overall.mape
            ],
            'smape_max': [
                'smape',  # Matches symmetric MAPE
                'performance_overall_smape',  # Matches flattened performance_overall.smape
            ],
            # Classification fairness metrics
            'demographic_parity': [
                'demographic_parity_diff',  # Matches demographic parity difference
            ],
            'equal_opportunity': [
                'equal_opportunity_diff',  # Matches equal opportunity difference
            ],
            'disparate_impact': [
                'disparate_impact_ratio',  # Matches disparate impact ratio
            ],
            'fpr_fnr_gaps_by_group': [
                'fpr_diff',  # Matches false positive rate difference
                'fnr_diff',  # Matches false negative rate difference
            ],
            'calibration_gaps': [
                'calibration_gap',  # Matches calibration gap metrics
            ],
            # Performance metrics - NOTE: overall_performance should NOT have thresholds
            # (It's informational only, thresholds are for group-based metrics)
            # 'overall_performance': [
            #     'performance_overall',  # Intentionally excluded - no thresholds for overall metrics
            # ],
            # Data quality metrics
            'missing_value_rates': [
                'missingness',  # Matches missing value metrics
            ],
            'data_quality_missing_value_rate_max': [
                'data_quality_missing_value_rate_max',  # Direct match
                'missing_value_rate_max',  # Alternative name
            ],
            'data_quality_schema_valid': [
                'data_quality_schema_valid',  # Direct match
                'schema_valid',  # Alternative name
            ],
            'train_test_leakage_detected': [
                'train_test_leakage_detected',  # Direct match
                'train_test_leakage',  # Alternative name
            ],
            'distribution_shifts': [
                'drift',  # Matches drift detection metrics
            ],
            # Explainability metrics
            'proxy_detection': [
                'proxy_risk',  # Matches proxy risk metrics
            ],
            'sensitive_leakage_detection': [
                'sensitive_leakage',  # Matches sensitive leakage metrics
            ],
        }
        
        # Try to match computed metric name against patterns
        import re
        for config_metric_id, patterns in metric_mappings.items():
            threshold_info = self.thresholds.get(config_metric_id)
            if threshold_info is None:
                continue
            
            # Extract threshold value
            if isinstance(threshold_info, dict):
                threshold_value = threshold_info.get('value')
            else:
                threshold_value = threshold_info
            
            if threshold_value is None:
                continue
            
            # Check if computed metric name contains any pattern (more flexible than exact match)
            for pattern in patterns:
                if pattern in computed_metric_name:
                    return threshold_value, threshold_info if isinstance(threshold_info, dict) else None
        
        return None, None
    
    def _validate_audit_intent(self):
        """Validate audit intent requirements (consent gating)."""
        if self.audit_intent == 'fairness_enforced':
            # For fairness_enforced, require protected_columns OR allow "pending suggestions" state
            if not self.sensitive and not self.fairness_protected_columns:
                # Check if we're in a "pending suggestions" state (indicated by ui_hints)
                if not self.ui_hints_missing_fields or 'protected_columns' not in self.ui_hints_missing_fields:
                    # This is a hard requirement - raise error
                    raise ValueError(
                        "audit_intent='fairness_enforced' requires 'protected_columns' or 'fairness.protected_columns'. "
                        "Either provide protected columns or set audit_intent='performance_only'."
                    )

