"""Evidence collection and reporting."""
import json
import hashlib
import uuid
from pathlib import Path
from typing import Dict, Any, Optional
from datetime import datetime
import pandas as pd
import pickle

from audit.governance import check_governance_config
from audit.ethics import map_ethics_to_evidence


def _flatten_category_metrics(category_metrics: Dict[str, Any], flat_metrics: Dict[str, float], prefix: str = ""):
    """Flatten category-based metrics structure for threshold checking."""
    for key, value in category_metrics.items():
        full_key = f"{prefix}_{key}" if prefix else key
        
        if isinstance(value, dict):
            _flatten_category_metrics(value, flat_metrics, full_key)
        elif isinstance(value, (int, float)):
            flat_metrics[full_key] = float(value)
        elif isinstance(value, list):
            # Handle worst_case_slices and other lists
            if key == "worst_case_slices" or (isinstance(value, list) and value and isinstance(value[0], dict)):
                for i, item in enumerate(value):
                    if isinstance(item, dict) and "value" in item:
                        flat_metrics[f"{full_key}_{i}_value"] = float(item["value"])


def compute_file_hash(file_path: Path) -> str:
    """Compute SHA-256 hash of a file."""
    sha256_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        for byte_block in iter(lambda: f.read(4096), b""):
            sha256_hash.update(byte_block)
    return sha256_hash.hexdigest()


def get_model_metadata(model_path: Path) -> Dict[str, Any]:
    """Extract model metadata."""
    try:
        with open(model_path, 'rb') as f:
            model = pickle.load(f)
        
        model_type = type(model).__name__
        
        # Try to get version if available
        model_version = None
        if hasattr(model, '__version__'):
            model_version = str(model.__version__)
        elif hasattr(model, 'get_params'):
            params = model.get_params()
            if 'n_estimators' in params:
                model_version = f"n_estimators={params.get('n_estimators')}"
        
        return {
            "type": model_type,
            "version": model_version,
            "path": str(model_path)
        }
    except Exception as e:
        return {
            "type": "unknown",
            "version": None,
            "path": str(model_path),
            "error": str(e)
        }


def get_dataset_metadata(dataset_path: Path, label_col: str, sensitive_attrs: list) -> Dict[str, Any]:
    """Extract dataset metadata (counts only, no raw data)."""
    try:
        df = pd.read_csv(dataset_path)
        
        # Get sensitive attribute metadata (counts only, no actual values)
        sensitive_metadata = {}
        for attr in sensitive_attrs:
            if attr in df.columns:
                unique_count = df[attr].nunique()
                # Only store count, not actual values (for privacy)
                sensitive_metadata[attr] = {
                    "unique_value_count": int(unique_count),
                    "present": True
                }
        
        return {
            "source": str(dataset_path),
            "size": len(df),
            "features": len(df.columns),
            "sensitive_attributes_metadata": sensitive_metadata
        }
    except Exception as e:
        return {
            "source": str(dataset_path),
            "error": str(e)
        }


def write_report(
    evidence_dir: Path,
    category_metrics: Dict[str, Any],  # Changed from flat metrics to category-based
    thresholds: Dict[str, float],
    passed: bool,
    config_path: Optional[Path] = None,
    dataset_path: Optional[Path] = None,
    model_path: Optional[Path] = None,
    label_col: Optional[str] = None,
    sensitive_attrs: Optional[list] = None,
    dataset_info: Optional[Dict[str, Any]] = None,
    dataset_metadata: Optional[Dict[str, Any]] = None,
    has_evidence_storage: bool = False,
    has_api_access: bool = False,
    include_governance: bool = True,
    include_ethics: bool = True,
    task_type: Optional[str] = None
) -> Path:
    """Write audit report to JSON file."""
    evidence_dir.mkdir(parents=True, exist_ok=True)
    
    report_path = evidence_dir / "report.json"
    
    # Flatten category_metrics to flat metrics dict for threshold checking
    def _flatten_category_metrics(category_metrics: Dict[str, Any], flat_metrics: Dict[str, float], prefix: str = ""):
        """Flatten category-based metrics structure for threshold checking."""
        for key, value in category_metrics.items():
            full_key = f"{prefix}_{key}" if prefix else key
            
            if isinstance(value, dict):
                _flatten_category_metrics(value, flat_metrics, full_key)
            elif isinstance(value, (int, float)):
                flat_metrics[full_key] = float(value)
            elif isinstance(value, list):
                # Handle worst_case_slices and other lists
                if key == "worst_case_slices":
                    for i, slice_data in enumerate(value):
                        if isinstance(slice_data, dict) and "value" in slice_data:
                            flat_metrics[f"{full_key}_{i}_value"] = float(slice_data["value"])
    
    metrics = {}
    _flatten_category_metrics(category_metrics, metrics)
    
    # Check which metrics passed/failed
    metric_results = {}
    for metric_name, value in metrics.items():
        # Try to find threshold by exact name first
        threshold = thresholds.get(metric_name)
        threshold_key = None
        
        # If not found, try to match by base metric name (without attribute prefix)
        if threshold is None:
            # Try to find matching threshold by checking if threshold key appears in metric name
            for base_name in thresholds.keys():
                # Check if metric name ends with the threshold key (with or without underscore prefix)
                if metric_name.endswith(base_name) or metric_name.endswith(f"_{base_name}"):
                    threshold = thresholds.get(base_name)
                    threshold_key = base_name
                    break
                # Check if threshold key starts with a suffix of the metric name
                if '_' in metric_name:
                    # Try removing attribute prefix by finding where the threshold might start
                    for i in range(1, len(metric_name.split('_'))):
                        suffix = '_'.join(metric_name.split('_')[i:])
                        if base_name.startswith(suffix) or suffix in base_name:
                            threshold = thresholds.get(base_name)
                            threshold_key = base_name
                            break
                    if threshold is not None:
                        break
        
        if threshold is not None:
            # Determine if metric passed based on its type
            # Check threshold key name or metric name for type
            is_ratio = "ratio" in metric_name or (threshold_key and "ratio" in threshold_key)
            is_diff = "diff" in metric_name or (threshold_key and "diff" in threshold_key)
            
            if is_diff:
                # For difference metrics, pass if value <= threshold
                metric_passed = value <= threshold
            elif is_ratio:
                # For ratio metrics, pass if value >= threshold
                metric_passed = value >= threshold
            else:
                # Default: pass if value <= threshold
                metric_passed = value <= threshold
        else:
            metric_passed = None  # No threshold specified
        
        metric_results[metric_name] = {
            "value": float(value),  # Convert numpy types to Python types
            "threshold": float(threshold) if threshold is not None else None,
            "passed": bool(metric_passed) if metric_passed is not None else None
        }
    
    # Generate audit ID and timestamp
    audit_id = str(uuid.uuid4())
    timestamp = datetime.now().isoformat()
    
    # Compute file hashes
    file_hashes = {}
    if config_path and config_path.exists():
        file_hashes["config"] = {
            "path": str(config_path),
            "sha256": compute_file_hash(config_path)
        }
    if dataset_path and dataset_path.exists():
        file_hashes["dataset"] = {
            "path": str(dataset_path),
            "sha256": compute_file_hash(dataset_path)
        }
    if model_path and model_path.exists():
        file_hashes["model"] = {
            "path": str(model_path),
            "sha256": compute_file_hash(model_path)
        }
    
    # Get model metadata
    model_metadata = {}
    if model_path:
        model_metadata = get_model_metadata(model_path)
    
    # Get dataset metadata (combine config metadata with computed info)
    computed_metadata = {}
    if dataset_path and label_col and sensitive_attrs:
        computed_metadata = get_dataset_metadata(dataset_path, label_col, sensitive_attrs)
    
    # Merge config metadata with computed metadata
    # Config metadata takes precedence for description, population, jurisdiction, sample_size
    if dataset_metadata:
        final_dataset_metadata = {
            **computed_metadata,
            **dataset_metadata  # Override with config-provided metadata
        }
    else:
        final_dataset_metadata = computed_metadata
    
    # Add dataset info (shape and distributions - counts only, no raw data)
    if dataset_info:
        final_dataset_metadata["dataset_info"] = dataset_info
    
    # Phase 3: Governance checks (informational, non-blocking) - only for Pro/Enterprise
    governance_checks = None
    if include_governance:
        governance_checks = check_governance_config(
            config_path=config_path,
            thresholds=thresholds,
            sensitive_attrs=sensitive_attrs,
            has_evidence_storage=has_evidence_storage,
            has_api_access=has_api_access
        )
    
    # Phase 4: ACM Ethics mapping (traceable, no scoring) - only for Pro/Enterprise
    ethics_mapping = None
    if include_ethics:
        ethics_mapping = map_ethics_to_evidence(
            metrics=metric_results,
            thresholds=thresholds,
            governance_checks=governance_checks or {},
            has_evidence_storage=has_evidence_storage
        )
    
    # Build report according to Audit Report Schema v2 (expanded structure)
    report = {
        "audit_id": audit_id,
        "timestamp": timestamp,
        "overall_passed": passed,
        
        # Section 1: Fairness Metrics (existing - Phase 1 locked)
        "fairness_metrics": {
            "metrics": metric_results,
            "thresholds": thresholds,
            "summary": {
                "total_metrics": len(metrics),
                "metrics_with_thresholds": len([m for m in metric_results.values() if m["threshold"] is not None]),
                "metrics_passed": len([m for m in metric_results.values() if m["passed"] is True]),
                "metrics_failed": len([m for m in metric_results.values() if m["passed"] is False])
            }
        },
        
        # Section 2: Governance Checks (Phase 3)
        "governance_checks": governance_checks,
        
        # Section 3: Ethics Mapping (Phase 4)
        "ethics_mapping": ethics_mapping,
        
        # Section 4: Evidence Metadata (Phase 6)
        "evidence_metadata": {
            "audit_id": audit_id,
            "timestamp": timestamp,
            "file_hashes": file_hashes,
            "model_metadata": model_metadata,
            "dataset_metadata": final_dataset_metadata,
            "sensitive_attributes": sensitive_attrs if sensitive_attrs else [],
            "retention_status": "active" if has_evidence_storage else "disabled",
            "immutability": "enforced" if has_evidence_storage else "not_enforced"
        },
        
        # Legacy fields for backward compatibility
        "model_metadata": model_metadata,
        "dataset_metadata": final_dataset_metadata,
        "sensitive_attributes": sensitive_attrs if sensitive_attrs else [],
        "metrics": metric_results,
        "thresholds": thresholds,
        "file_hashes": file_hashes,
        "summary": {
            "total_metrics": len(metrics),
            "metrics_with_thresholds": len([m for m in metric_results.values() if m["threshold"] is not None]),
            "metrics_passed": len([m for m in metric_results.values() if m["passed"] is True]),
            "metrics_failed": len([m for m in metric_results.values() if m["passed"] is False])
        }
    }
    
    with open(report_path, 'w') as f:
        json.dump(report, f, indent=2)
    
    return report_path

