"""Regulator-ready export functionality."""
import json
from pathlib import Path
from typing import Dict, Any, Optional
from datetime import datetime
from reportlab.lib import colors
from reportlab.lib.pagesizes import letter, A4
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
from reportlab.lib.units import inch
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, Table, TableStyle, PageBreak
from reportlab.lib.enums import TA_CENTER, TA_LEFT, TA_JUSTIFY


METRIC_DEFINITIONS = {
    "demographic_parity_diff": {
        "name": "Demographic Parity Difference",
        "definition": "The maximum difference in positive prediction rates across protected groups. Measures whether the model treats different groups equally in terms of approval rates.",
        "formula": "max(P(ŷ=1|A=a)) - min(P(ŷ=1|A=a))",
        "interpretation": "Lower values indicate better fairness. A value of 0 means perfect demographic parity.",
        "threshold_guidance": "Typically acceptable if ≤ 0.10 (10 percentage points)"
    },
    "equal_opportunity_diff": {
        "name": "Equal Opportunity Difference",
        "definition": "The maximum difference in true positive rates (TPR) across protected groups. Measures whether the model has equal accuracy for positive outcomes across groups.",
        "formula": "max(TPR|A=a) - min(TPR|A=a)",
        "interpretation": "Lower values indicate better fairness. A value of 0 means equal opportunity for all groups.",
        "threshold_guidance": "Typically acceptable if ≤ 0.10 (10 percentage points)"
    },
    "disparate_impact_ratio": {
        "name": "Disparate Impact Ratio",
        "definition": "The ratio of the minimum positive prediction rate to the maximum positive prediction rate across protected groups. Measures relative fairness.",
        "formula": "min(P(ŷ=1|A=a)) / max(P(ŷ=1|A=a))",
        "interpretation": "Values closer to 1.0 indicate better fairness. Values < 0.80 may indicate disparate impact.",
        "threshold_guidance": "Typically acceptable if ≥ 0.80 (80% rule)"
    }
}


METHODOLOGY = """
This audit follows established fairness metrics for algorithmic decision-making systems. 
The methodology is based on principles from:

1. **Equalized Odds**: The model should have equal true positive and false positive rates across protected groups.
2. **Demographic Parity**: The model should predict positive outcomes at similar rates across protected groups.
3. **Disparate Impact**: The model should not disproportionately affect any protected group.

**Audit Process:**
1. Dataset validation: Verify presence of sensitive attributes and label column
2. Model prediction: Generate predictions for all samples
3. Metric calculation: Compute fairness metrics for each sensitive attribute
4. Threshold evaluation: Compare metrics against defined thresholds
5. Evidence generation: Create immutable audit report with hashes and metadata

**Data Privacy:**
- No raw data is stored or transmitted
- Only aggregated statistics (counts) are included
- Dataset and model are identified by SHA-256 hashes only
- Sensitive attribute values are not included in reports
"""


def get_metric_definition(metric_name: str) -> Dict[str, str]:
    """Get definition for a metric by matching its base name."""
    # Extract base metric name (remove attribute prefix)
    base_name = None
    for key in METRIC_DEFINITIONS.keys():
        if key in metric_name or metric_name.endswith(f"_{key}"):
            base_name = key
            break
    
    if base_name:
        return METRIC_DEFINITIONS[base_name]
    
    # Default definition for unknown metrics
    return {
        "name": metric_name.replace("_", " ").title(),
        "definition": "Custom fairness metric",
        "formula": "N/A",
        "interpretation": "See metric value and threshold",
        "threshold_guidance": "See threshold configuration"
    }


def export_json(
    report_path: Path,
    output_path: Path,
    config_path: Optional[Path] = None
) -> Path:
    """
    Export JSON report.
    
    NOTE: Local exports are for validation only. They do NOT provide:
    - Signed evidence for compliance
    - Regulator-ready summaries
    - Audit history or trend analysis
    
    For authoritative audit records, upload to Verity's platform.
    """
    # Load original report
    with open(report_path, 'r') as f:
        report = json.load(f)
    
    # Check if this is a local audit (no evidence storage metadata)
    is_local = not report.get("evidence_metadata", {}).get("retention_status") == "active"
    
    # Create enhanced export with methodology
    export_data = {
        "export_metadata": {
            "exported_at": datetime.now().isoformat(),
            "export_format": "json",
            "export_version": "1.0",
            "original_report": str(report_path),
            "audit_type": "local_validation" if is_local else "platform_recorded",
            "notice": "This is a local validation export. For governance enforcement, compliance evidence, and regulatory reporting, audits must be uploaded to Verity's platform." if is_local else None
        },
        "audit_summary": {
            "audit_id": report.get("audit_id"),
            "timestamp": report.get("timestamp"),
            "overall_passed": report.get("overall_passed"),
            "total_metrics": report.get("summary", {}).get("total_metrics", 0),
            "metrics_passed": report.get("summary", {}).get("metrics_passed", 0),
            "metrics_failed": report.get("summary", {}).get("metrics_failed", 0)
        },
        "methodology": METHODOLOGY.strip(),
        "metric_definitions": {},
        "detailed_results": report
    }
    
    # Add metric definitions for all metrics in report
    for metric_name in report.get("metrics", {}).keys():
        export_data["metric_definitions"][metric_name] = get_metric_definition(metric_name)
    
    # Write export
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, 'w') as f:
        json.dump(export_data, f, indent=2)
    
    return output_path


def export_pdf(
    report_path: Path,
    output_path: Path,
    config_path: Optional[Path] = None
) -> Path:
    """
    Export PDF report.
    
    NOTE: Local exports are for validation only. They do NOT provide:
    - Signed evidence for compliance
    - Regulator-ready summaries
    - Audit history or trend analysis
    
    For authoritative audit records, upload to Verity's platform.
    """
    # Load original report
    with open(report_path, 'r') as f:
        report = json.load(f)
    
    # Check if this is a local audit
    is_local = not report.get("evidence_metadata", {}).get("retention_status") == "active"
    
    # Create PDF document
    doc = SimpleDocTemplate(str(output_path), pagesize=letter)
    story = []
    
    # Define styles
    styles = getSampleStyleSheet()
    title_style = ParagraphStyle(
        'CustomTitle',
        parent=styles['Heading1'],
        fontSize=24,
        textColor=colors.HexColor('#1a1a1a'),
        spaceAfter=30,
        alignment=TA_CENTER
    )
    
    heading_style = ParagraphStyle(
        'CustomHeading',
        parent=styles['Heading2'],
        fontSize=16,
        textColor=colors.HexColor('#2c3e50'),
        spaceAfter=12,
        spaceBefore=20
    )
    
    body_style = ParagraphStyle(
        'CustomBody',
        parent=styles['BodyText'],
        fontSize=11,
        leading=14,
        alignment=TA_JUSTIFY
    )
    
    # Title
    title_text = "AI Model Fairness Audit Report"
    if is_local:
        title_text += " (Local Validation Only)"
    story.append(Paragraph(title_text, title_style))
    
    # Add notice for local audits
    if is_local:
        notice_style = ParagraphStyle(
            'Notice',
            parent=styles['BodyText'],
            fontSize=10,
            textColor=colors.HexColor('#e74c3c'),
            leading=12,
            alignment=TA_CENTER,
            backColor=colors.HexColor('#fff3cd'),
            borderPadding=8
        )
        story.append(Spacer(1, 0.2*inch))
        story.append(Paragraph(
            "<b>NOTICE:</b> This is a local validation export. It does not provide signed evidence, "
            "regulator-ready summaries, or audit history. For governance enforcement, compliance evidence, "
            "and regulatory reporting, audits must be uploaded to Verity's platform.",
            notice_style
        ))
    
    story.append(Spacer(1, 0.3*inch))
    
    # Audit Summary
    story.append(Paragraph("Audit Summary", heading_style))
    
    summary_data = [
        ["Audit ID", report.get("audit_id", "N/A")],
        ["Timestamp", report.get("timestamp", "N/A")],
        ["Overall Status", "PASSED" if report.get("overall_passed") else "FAILED"],
        ["Total Metrics", str(report.get("summary", {}).get("total_metrics", 0))],
        ["Metrics Passed", str(report.get("summary", {}).get("metrics_passed", 0))],
        ["Metrics Failed", str(report.get("summary", {}).get("metrics_failed", 0))]
    ]
    
    summary_table = Table(summary_data, colWidths=[2*inch, 4*inch])
    summary_table.setStyle(TableStyle([
        ('BACKGROUND', (0, 0), (0, -1), colors.HexColor('#ecf0f1')),
        ('TEXTCOLOR', (0, 0), (-1, -1), colors.black),
        ('ALIGN', (0, 0), (-1, -1), 'LEFT'),
        ('FONTNAME', (0, 0), (0, -1), 'Helvetica-Bold'),
        ('FONTSIZE', (0, 0), (-1, -1), 10),
        ('BOTTOMPADDING', (0, 0), (-1, -1), 8),
        ('TOPPADDING', (0, 0), (-1, -1), 8),
        ('GRID', (0, 0), (-1, -1), 1, colors.grey)
    ]))
    story.append(summary_table)
    story.append(Spacer(1, 0.2*inch))
    
    # Methodology
    story.append(Paragraph("Methodology", heading_style))
    for line in METHODOLOGY.strip().split('\n'):
        if line.strip():
            if line.strip().startswith('**') and line.strip().endswith('**'):
                # Bold heading
                text = line.strip().strip('*')
                story.append(Paragraph(f"<b>{text}</b>", body_style))
            elif line.strip().startswith(tuple('123456789')):
                # Numbered list
                story.append(Paragraph(line.strip(), body_style))
            else:
                story.append(Paragraph(line.strip(), body_style))
            story.append(Spacer(1, 0.1*inch))
    
    story.append(PageBreak())
    
    # Metric Definitions
    story.append(Paragraph("Metric Definitions", heading_style))
    
    metrics = report.get("metrics", {})
    for metric_name, metric_data in metrics.items():
        definition = get_metric_definition(metric_name)
        
        story.append(Paragraph(f"<b>{definition['name']}</b>", heading_style))
        story.append(Paragraph(f"<b>Definition:</b> {definition['definition']}", body_style))
        story.append(Spacer(1, 0.1*inch))
        story.append(Paragraph(f"<b>Formula:</b> {definition['formula']}", body_style))
        story.append(Spacer(1, 0.1*inch))
        story.append(Paragraph(f"<b>Interpretation:</b> {definition['interpretation']}", body_style))
        story.append(Spacer(1, 0.1*inch))
        story.append(Paragraph(f"<b>Threshold Guidance:</b> {definition['threshold_guidance']}", body_style))
        story.append(Spacer(1, 0.2*inch))
    
    story.append(PageBreak())
    
    # Detailed Results
    story.append(Paragraph("Detailed Results", heading_style))
    
    # Model Metadata
    model_meta = report.get("model_metadata", {})
    if model_meta:
        story.append(Paragraph("Model Information", heading_style))
        model_data = [
            ["Type", model_meta.get("type", "N/A")],
            ["Version", model_meta.get("version", "N/A")],
            ["Path", model_meta.get("path", "N/A")]
        ]
        model_table = Table(model_data, colWidths=[2*inch, 4*inch])
        model_table.setStyle(TableStyle([
            ('BACKGROUND', (0, 0), (0, -1), colors.HexColor('#ecf0f1')),
            ('TEXTCOLOR', (0, 0), (-1, -1), colors.black),
            ('ALIGN', (0, 0), (-1, -1), 'LEFT'),
            ('FONTNAME', (0, 0), (0, -1), 'Helvetica-Bold'),
            ('FONTSIZE', (0, 0), (-1, -1), 10),
            ('BOTTOMPADDING', (0, 0), (-1, -1), 8),
            ('TOPPADDING', (0, 0), (-1, -1), 8),
            ('GRID', (0, 0), (-1, -1), 1, colors.grey)
        ]))
        story.append(model_table)
        story.append(Spacer(1, 0.2*inch))
    
    # Dataset Metadata
    dataset_meta = report.get("dataset_metadata", {})
    if dataset_meta:
        story.append(Paragraph("Dataset Information", heading_style))
        dataset_data = [
            ["Source", dataset_meta.get("source", "N/A")],
            ["Size", str(dataset_meta.get("size", "N/A"))],
            ["Features", str(dataset_meta.get("features", "N/A"))]
        ]
        if dataset_meta.get("description"):
            dataset_data.append(["Description", dataset_meta.get("description")])
        if dataset_meta.get("population"):
            dataset_data.append(["Population", dataset_meta.get("population")])
        if dataset_meta.get("jurisdiction"):
            dataset_data.append(["Jurisdiction", dataset_meta.get("jurisdiction")])
        
        dataset_table = Table(dataset_data, colWidths=[2*inch, 4*inch])
        dataset_table.setStyle(TableStyle([
            ('BACKGROUND', (0, 0), (0, -1), colors.HexColor('#ecf0f1')),
            ('TEXTCOLOR', (0, 0), (-1, -1), colors.black),
            ('ALIGN', (0, 0), (-1, -1), 'LEFT'),
            ('FONTNAME', (0, 0), (0, -1), 'Helvetica-Bold'),
            ('FONTSIZE', (0, 0), (-1, -1), 10),
            ('BOTTOMPADDING', (0, 0), (-1, -1), 8),
            ('TOPPADDING', (0, 0), (-1, -1), 8),
            ('GRID', (0, 0), (-1, -1), 1, colors.grey)
        ]))
        story.append(dataset_table)
        story.append(Spacer(1, 0.2*inch))
    
    # Metrics Table
    story.append(Paragraph("Metric Results", heading_style))
    
    metrics_data = [["Metric", "Value", "Threshold", "Status"]]
    for metric_name, metric_data in metrics.items():
        value = f"{metric_data.get('value', 0):.4f}"
        threshold = f"{metric_data.get('threshold', 'N/A')}" if metric_data.get('threshold') else "N/A"
        status = "PASS" if metric_data.get('passed') else "FAIL" if metric_data.get('passed') is False else "N/A"
        metrics_data.append([metric_name, value, threshold, status])
    
    metrics_table = Table(metrics_data, colWidths=[2.5*inch, 1*inch, 1*inch, 1*inch])
    metrics_table.setStyle(TableStyle([
        ('BACKGROUND', (0, 0), (-1, 0), colors.HexColor('#34495e')),
        ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
        ('ALIGN', (0, 0), (-1, -1), 'CENTER'),
        ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
        ('FONTSIZE', (0, 0), (-1, 0), 11),
        ('BOTTOMPADDING', (0, 0), (-1, 0), 12),
        ('BACKGROUND', (0, 1), (-1, -1), colors.beige),
        ('TEXTCOLOR', (0, 1), (-1, -1), colors.black),
        ('FONTSIZE', (0, 1), (-1, -1), 9),
        ('GRID', (0, 0), (-1, -1), 1, colors.grey),
        ('ROWBACKGROUNDS', (0, 1), (-1, -1), [colors.white, colors.HexColor('#f8f9fa')])
    ]))
    story.append(metrics_table)
    
    # File Hashes
    story.append(Spacer(1, 0.3*inch))
    story.append(Paragraph("File Integrity Hashes", heading_style))
    file_hashes = report.get("file_hashes", {})
    if file_hashes:
        hash_data = [["File Type", "SHA-256 Hash"]]
        for file_type, hash_info in file_hashes.items():
            hash_data.append([file_type.title(), hash_info.get("sha256", "N/A")])
        
        hash_table = Table(hash_data, colWidths=[2*inch, 4*inch])
        hash_table.setStyle(TableStyle([
            ('BACKGROUND', (0, 0), (0, 0), colors.HexColor('#34495e')),
            ('TEXTCOLOR', (0, 0), (-1, 0), colors.whitesmoke),
            ('ALIGN', (0, 0), (-1, -1), 'LEFT'),
            ('FONTNAME', (0, 0), (-1, 0), 'Helvetica-Bold'),
            ('FONTSIZE', (0, 0), (-1, 0), 11),
            ('FONTSIZE', (0, 1), (-1, -1), 8),
            ('FONTNAME', (0, 1), (-1, -1), 'Courier'),
            ('BOTTOMPADDING', (0, 0), (-1, -1), 8),
            ('TOPPADDING', (0, 0), (-1, -1), 8),
            ('GRID', (0, 0), (-1, -1), 1, colors.grey)
        ]))
        story.append(hash_table)
    
    # Digital Signature
    story.append(Spacer(1, 0.3*inch))
    story.append(Paragraph("Digital Signature", heading_style))
    if is_local:
        story.append(Paragraph(
            "<i><b>This is a local validation export.</b> Digital signatures, signed evidence, "
            "and regulator-ready summaries are only available for audits uploaded to Verity's platform.</i>",
            body_style
        ))
    else:
        story.append(Paragraph(
            "<i>Digital signature functionality will be available in a future release.</i>",
            body_style
        ))
    story.append(Paragraph(f"<i>Report generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S UTC')}</i>", body_style))
    
    # Build PDF
    doc.build(story)
    
    return output_path

