"""Command-line interface for audit tool."""
import sys
from pathlib import Path
import argparse
from rich.console import Console
from rich.table import Table
from rich import box

from audit.config import Config
from audit.metrics import calculate_metrics
from audit.evidence import write_report
from audit.export import export_pdf, export_json
from typing import Optional, Dict, Any
import os
import json
import requests
import pandas as pd


console = Console()


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"])


def _flatten_api_metrics(api_results: Dict[str, Any], flat_metrics: Dict[str, float]):
    """Flatten API audit metrics for threshold checking."""
    for category, data in api_results.get("api_categories", {}).items():
        metrics = data.get("metrics", {})
        for metric_name, value in metrics.items():
            if isinstance(value, (int, float)):
                flat_metrics[f"{category}_{metric_name}"] = float(value)


def run_audit(
    config_path: str,
    fail_on_threshold: bool = True,
    output_dir: Optional[str] = None,
    upload: bool = False
) -> bool:
    """
    Run audit and return True if passed, False if failed.
    
    NOTE: Local audits are for developer validation only. They do NOT provide:
    - Audit history or trend analysis
    - Comparison against past runs
    - Regulator-ready summaries
    - Signed evidence for compliance
    - CI/CD enforcement capabilities
    
    For governance enforcement, compliance evidence, and regulatory reporting,
    audits must be uploaded to Verity's platform.
    """
    # Show license notice for local audits (when no API key)
    api_key = os.getenv("VERITY_API_KEY")
    if not api_key and not upload:
        console.print("\n[yellow]⚠️  LOCAL AUDIT MODE[/yellow]")
        console.print("[dim]Local audits are for developer validation only.[/dim]")
        console.print("[dim]For governance enforcement and compliance evidence, upload to Verity's platform.[/dim]\n")
    
    # Load configuration
    try:
        config = Config(config_path)
    except Exception as e:
        console.print(f"[red]Error loading config: {e}[/red]")
        return False
    
    # Handle API model mode differently
    if config.mode == 'api_model':
        console.print("[cyan]API Model Audit Mode[/cyan]")
        console.print("[yellow]⚠️  API model audits require platform upload for full functionality.[/yellow]")
        console.print("[yellow]   Local API audits are not yet supported in CLI mode.[/yellow]")
        console.print("[yellow]   Please use the Verity platform for API model audits.[/yellow]")
        console.print("[yellow]   For CI/CD, ensure your workflow uploads results to Verity.[/yellow]")
        return True  # Don't fail, just inform user
    
    # Warn if dataset metadata is missing (required for production)
    if not config.dataset_metadata:
        console.print("[yellow]⚠️  Warning: dataset_metadata not provided in config[/yellow]")
        console.print("[yellow]   For production audits, include: description, population, jurisdiction, sample_size[/yellow]")
    
    # For hybrid mode, we still need dataset for the dataset portion
    if config.mode == 'hybrid' and not config.dataset:
        console.print("[red]Error: Hybrid mode requires a dataset for the dataset_model portion[/red]")
        return False
    
    # Handle hybrid mode
    if config.mode == 'hybrid':
        console.print("[cyan]Hybrid Audit Mode (Dataset + API + Integration Tests)[/cyan]")
        
        try:
            from audit.hybrid_audit import run_hybrid_audit
            
            # Get test suite path
            test_suite_path = Path(config.test_suite_source) if config.test_suite_source else None
            if not test_suite_path or not test_suite_path.exists():
                console.print("[yellow]⚠️  Test suite not found for API component. Creating starter suite...[/yellow]")
                from audit.api_test_suite import create_starter_test_suite
                starter_suite = create_starter_test_suite()
                test_suite_path = Path.cwd() / "tests" / "verity_starter_suite.yaml"
                starter_suite.save_to_file(test_suite_path)
                console.print(f"[green]Created starter test suite at: {test_suite_path}[/green]")
            
            console.print("[cyan]Running hybrid audit...[/cyan]")
            
            hybrid_results = run_hybrid_audit(
                dataset_path=str(config.dataset),
                model_path=str(config.model),
                label_col=config.label,
                sensitive_attrs=config.sensitive,
                api_provider=config.api_provider,
                api_model=config.api_model,
                api_key_env=config.api_key_env,
                test_suite_path=test_suite_path,
                task_type=config.task_type,
                positive_class=config.positive_class,
                domain_profile=config.domain_profile
            )
            
            # Combine dataset and API metrics
            category_metrics = {
                **hybrid_results["dataset_component"],
                "api_categories": hybrid_results.get("api_component", {}),
                "end_to_end_integration": hybrid_results.get("end_to_end_integration", {})
            }
            
            # Flatten for threshold checking
            metrics = {}
            _flatten_category_metrics(hybrid_results["dataset_component"], metrics)
            if hybrid_results.get("api_component"):
                _flatten_api_metrics({"api_categories": hybrid_results["api_component"]}, metrics)
            
            # Log summary
            console.print(f"[dim]Dataset audit: completed[/dim]")
            console.print(f"[dim]API audit: {'completed' if hybrid_results.get('api_component') else 'skipped'}[/dim]")
            console.print(f"[dim]Integration tests: completed[/dim]")
            
            # Integration test summary
            integration = hybrid_results.get("end_to_end_integration", {})
            if integration:
                console.print("\n[bold]Integration Test Results[/bold]")
                if "wrapper_leakage" in integration:
                    leakage = integration["wrapper_leakage"]
                    leakage_rate = leakage.get('leakage_rate', 0)
                    console.print(f"[dim]Wrapper leakage rate: {leakage_rate:.1%}[/dim]")
                if "demographic_consistency" in integration:
                    consistency = integration["demographic_consistency"]
                    consistency_rate = consistency.get('consistency_rate', 0)
                    console.print(f"[dim]Demographic consistency: {consistency_rate:.1%}[/dim]")
            
            # Get dataset info from dataset component
            dataset_info = {
                "shape": {
                    "rows": len(pd.read_csv(str(config.dataset))),
                    "columns": len(pd.read_csv(str(config.dataset)).columns)
                }
            }
            
            # Continue with threshold checking and report writing (shared logic below)
            # category_metrics and metrics are already set above
        except Exception as e:
            console.print(f"[red]Error running hybrid audit: {e}[/red]")
            import traceback
            console.print(f"[red]{traceback.format_exc()}[/red]")
            return False
    
    # Calculate metrics (only for dataset_model mode - hybrid already calculated above)
    if config.mode == 'dataset_model':
        if not config.dataset:
            console.print("[red]Error: Dataset is required for dataset_model and hybrid modes[/red]")
            return False
        if not config.model:
            console.print("[red]Error: Model is required for dataset_model and hybrid modes[/red]")
            return False
            
        try:
            console.print("[cyan]Calculating audit metrics (category-based)...[/cyan]")
            category_metrics, dataset_info = calculate_metrics(
                str(config.dataset),
                str(config.model),
                config.label,
                config.sensitive,
                task_type=config.task_type,
                positive_class=config.positive_class
            )
            
            # Flatten category metrics for threshold checking (backward compatibility)
            # Extract all metric values from categories
            metrics = {}
            _flatten_category_metrics(category_metrics, metrics)
            
            # Log dataset info (counts only, no raw data)
            console.print(f"[dim]Dataset: {dataset_info['shape']['rows']} rows × {dataset_info['shape']['columns']} columns[/dim]")
            if dataset_info.get('sensitive_attribute_distributions'):
                console.print("[dim]Sensitive attribute distributions logged[/dim]")
            
            # Log category summary
            console.print(f"[dim]Fairness metrics: {len(category_metrics.get('fairness', {}).get('group_metrics', {}))} attributes analyzed[/dim]")
            console.print(f"[dim]Performance metrics: calculated[/dim]")
            console.print(f"[dim]Data quality checks: completed[/dim]")
            console.print(f"[dim]Explainability analysis: completed[/dim]")
        except Exception as e:
            console.print(f"[red]Error calculating metrics: {e}[/red]")
            import traceback
            console.print(f"[red]{traceback.format_exc()}[/red]")
            return False
    else:
        # Should not reach here, but handle gracefully
        console.print(f"[red]Error: Unknown audit mode: {config.mode}[/red]")
        return False
    
    # Determine plan from environment variable or API
    # IMPORTANT: Do this BEFORE threshold filtering so plan is correct
    plan = os.getenv("VERITY_PLAN", "starter").lower()  # Default to starter for local audits
    
    # If API key is available, assume Pro plan (Starter plan doesn't have API access)
    # This ensures thresholds are correct from the start
    api_key = os.getenv("VERITY_API_KEY")
    if api_key and plan == "starter":
        # Starter plan doesn't have API access, so if API key is present, must be Pro or Enterprise
        plan = "pro"
        console.print(f"[dim]API key detected: Assuming Pro plan (Starter plan has no API access)[/dim]")
    
    # Starter plan: Only check these 3 specific thresholds
    STARTER_THRESHOLDS = {
        "demographic_parity_diff": 0.10,
        "equal_opportunity_diff": 0.10,
        "disparate_impact_ratio_min": 0.80
    }
    
    # Filter thresholds based on plan
    if plan == "starter":
        # For starter, only use the 3 basic thresholds
        filtered_thresholds = {}
        for key in STARTER_THRESHOLDS.keys():
            if key in config.thresholds:
                filtered_thresholds[key] = config.thresholds[key]
            else:
                # Use default values if not in config
                filtered_thresholds[key] = STARTER_THRESHOLDS[key]
        # Also check for any threshold keys that match the starter pattern
        for key, value in config.thresholds.items():
            if any(starter_key in key.lower() for starter_key in ["demographic_parity_diff", "equal_opportunity_diff", "disparate_impact_ratio_min"]):
                filtered_thresholds[key] = value
        config.thresholds = filtered_thresholds
        console.print(f"[dim]Starter plan: Checking only basic fairness thresholds[/dim]")
    elif plan in ["pro", "enterprise"]:
        # Pro/Enterprise: Use all thresholds from config
        console.print(f"[dim]{plan.capitalize()} plan: Checking all configured thresholds + governance checks[/dim]")
    
    # Check thresholds (using flattened metrics)
    passed = True
    threshold_checks = {}
    
    for metric_name, value in metrics.items():
        # Try to find threshold by exact name first
        threshold = config.get_threshold(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
            # This handles cases where attributes have underscores (e.g., "age_group")
            for base_name in config.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 = config.get_threshold(base_name)
                    threshold_key = base_name
                    break
                # Check if threshold key starts with a suffix of the metric name
                # e.g., "disparate_impact_ratio_min" matches "gender_disparate_impact_ratio"
                if '_' in metric_name:
                    # Try removing attribute prefix by finding where the threshold might start
                    # Check all possible suffixes
                    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 = config.get_threshold(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
            
            threshold_checks[metric_name] = {
                "value": value,
                "threshold": threshold,
                "passed": metric_passed
            }
            
            if not metric_passed:
                passed = False
    
    # Print summary
    console.print("\n[bold]Audit Summary[/bold]\n")
    
    table = Table(box=box.ROUNDED)
    table.add_column("Metric", style="cyan")
    table.add_column("Value", style="magenta")
    table.add_column("Threshold", style="yellow")
    table.add_column("Status", style="bold")
    
    for metric_name, check in threshold_checks.items():
        status = "[green]PASS[/green]" if check["passed"] else "[red]FAIL[/red]"
        table.add_row(
            metric_name,
            f"{check['value']:.4f}",
            f"{check['threshold']:.4f}",
            status
        )
    
    # Add metrics without thresholds
    for metric_name, value in metrics.items():
        if metric_name not in threshold_checks:
            table.add_row(
                metric_name,
                f"{value:.4f}",
                "N/A",
                "[dim]No threshold[/dim]"
            )
    
    console.print(table)
    
    # Overall status
    overall_status = "[green]PASSED[/green]" if passed else "[red]FAILED[/red]"
    console.print(f"\n[bold]Overall Status: {overall_status}[/bold]\n")
    
    # Override evidence_dir if output_dir is specified
    evidence_dir = Path(output_dir) if output_dir else config.evidence_dir
    
    # Get API key for upload
    api_key = os.getenv("VERITY_API_KEY")
    
    # Check if this is a local audit (no API key = local only)
    is_local_audit = not api_key and not upload
    
    # Check if evidence storage is enabled (from environment or config)
    # Pro/Enterprise plans have evidence storage, but only when uploaded to platform
    has_evidence_storage = (plan in ["pro", "enterprise"] and not is_local_audit) or (os.getenv("VERITY_EVIDENCE_STORAGE", "false").lower() == "true")
    has_api_access = api_key is not None  # If API key exists, API access is available
    
    # Governance checks and ACM ethics mapping only for Pro/Enterprise AND when uploaded
    # Local audits are intentionally incomplete - no governance checks or ethics mapping
    include_governance = (plan in ["pro", "enterprise"]) and not is_local_audit
    include_ethics = (plan in ["pro", "enterprise"]) and not is_local_audit
    
    # Write report with category-based structure
    # Only include governance checks and ethics mapping for Pro/Enterprise
    try:
        report_path = write_report(
            evidence_dir,
            category_metrics,  # Pass category_metrics (includes hybrid structure if applicable)
            config.thresholds,
            passed,
            config_path=config.config_path,
            dataset_path=config.dataset if config.mode != 'api_model' else None,
            model_path=config.model if config.mode != 'api_model' else None,
            label_col=config.label if config.mode != 'api_model' else None,
            sensitive_attrs=config.sensitive if config.mode != 'api_model' else None,
            dataset_info=dataset_info if config.mode != 'api_model' else None,
            dataset_metadata={
                "description": config.dataset_description,
                "population": config.dataset_population,
                "jurisdiction": config.dataset_jurisdiction,
                "sample_size": config.dataset_sample_size
            } if config.dataset_metadata else None,
            has_evidence_storage=has_evidence_storage,
            has_api_access=has_api_access,
            include_governance=include_governance,
            include_ethics=include_ethics,
            task_type=config.task_type
        )
        console.print(f"[green]Report written to: {report_path}[/green]")
        
        # Reminder for local audits
        if is_local_audit:
            console.print("\n[yellow]⚠️  LOCAL AUDIT - VALIDATION ONLY[/yellow]")
            console.print("[dim]This local audit is for developer validation and debugging only.[/dim]")
            console.print("[dim]It does NOT provide:[/dim]")
            console.print("[dim]  • Audit history or trend analysis[/dim]")
            console.print("[dim]  • Comparison against past runs[/dim]")
            console.print("[dim]  • Regulator-ready summaries[/dim]")
            console.print("[dim]  • Signed evidence for compliance[/dim]")
            console.print("[dim]  • CI/CD enforcement capabilities[/dim]")
            console.print("[dim][/dim]")
            console.print("[dim]For governance enforcement, compliance evidence, and regulatory reporting,[/dim]")
            console.print("[dim]audits must be uploaded to Verity's platform.[/dim]")
            console.print("[dim]Local runs are advisory. Platform-recorded runs are authoritative.[/dim]\n")
    except Exception as e:
        console.print(f"[red]Error writing report: {e}[/red]")
        return False
    
    # Handle upload if requested or if API key is available
    should_upload = upload or (api_key is not None)
    
    if should_upload:
        # Use production API URL by default, not localhost
        api_url = os.getenv("VERITY_API_URL") or "https://verityai-production.up.railway.app"
        
        if not api_key:
            console.print("[yellow]VERITY_API_KEY not set - skipping upload (offline mode)[/yellow]")
        else:
            try:
                # Read the report we just wrote
                with open(report_path, 'r') as f:
                    report_data = json.load(f)
                
                # Extract commit context from GitHub Actions environment (if available)
                # These are set automatically by GitHub Actions
                commit_sha = os.getenv("GITHUB_SHA")
                branch = os.getenv("GITHUB_REF_NAME")
                repo = os.getenv("GITHUB_REPOSITORY")
                pr_number = os.getenv("GITHUB_PR_NUMBER") or os.getenv("GITHUB_EVENT_PULL_REQUEST_NUMBER")
                github_server = os.getenv("GITHUB_SERVER_URL", "https://github.com")
                run_id = os.getenv("GITHUB_RUN_ID")
                
                # Build workflow run URL if we have the necessary info
                workflow_run_url = None
                if github_server and repo and run_id:
                    workflow_run_url = f"{github_server}/{repo}/actions/runs/{run_id}"
                
                # Add commit context to report metadata if available
                if commit_sha or branch or repo or pr_number or workflow_run_url:
                    if "metadata" not in report_data:
                        report_data["metadata"] = {}
                    if commit_sha:
                        report_data["metadata"]["commit_sha"] = commit_sha
                        report_data["commit_sha"] = commit_sha  # Also at top level for convenience
                    if branch:
                        report_data["metadata"]["branch"] = branch
                        report_data["branch"] = branch
                    if repo:
                        report_data["metadata"]["repo"] = repo
                        report_data["repo"] = repo
                    if pr_number:
                        try:
                            report_data["metadata"]["pr_number"] = int(pr_number)
                            report_data["pr_number"] = int(pr_number)
                        except (ValueError, TypeError):
                            pass
                    if workflow_run_url:
                        report_data["metadata"]["workflow_run_url"] = workflow_run_url
                
                # Upload to API
                # API keys use X-API-Key header, not Bearer token
                headers = {
                    "X-API-Key": api_key,
                    "Content-Type": "application/json"
                }
                
                response = requests.post(
                    f"{api_url}/v1/audits",
                    json=report_data,
                    headers=headers,
                    timeout=10
                )
                
                if response.status_code == 200:
                    result = response.json()
                    console.print(f"[green]✓ Audit uploaded: {result.get('audit_id')}[/green]")
                    if commit_sha:
                        console.print(f"[dim]Commit: {commit_sha[:8]}[/dim]")
                    
                    # Get plan from response (either in body or headers)
                    org_plan = (
                        result.get("organization_plan") or 
                        response.headers.get("X-Organization-Plan") or 
                        response.headers.get("Organization-Plan")
                    )
                    if org_plan:
                        plan = org_plan.lower()
                        console.print(f"[dim]Detected plan from API: {plan}[/dim]")
                        # Update thresholds based on detected plan
                        if plan in ["pro", "enterprise"]:
                            # Pro/Enterprise: Use all thresholds from config
                            config.thresholds = config.data.get('thresholds', {})
                            console.print(f"[dim]{plan.capitalize()} plan: Checking all configured thresholds[/dim]")
                elif response.status_code == 409:
                    console.print(f"[yellow]Audit already exists (append-only storage)[/yellow]")
                else:
                    console.print(f"[red]Upload failed: {response.status_code} - {response.text}[/red]")
            
            except requests.exceptions.RequestException as e:
                console.print(f"[red]Upload error: {e}[/red]")
                console.print("[yellow]Continuing in offline mode[/yellow]")
    
    # Return pass/fail based on fail_on_threshold flag
    if fail_on_threshold:
        return passed
    else:
        # If not failing on threshold, always return True (but still report the status)
        return True


def main():
    """Main CLI entry point."""
    parser = argparse.ArgumentParser(
        description="AI Model Audit Tool",
        epilog="""
LICENSE NOTICE:
Local audits are allowed for personal or internal use only.
Redistribution, resale, or hosted use is prohibited.

Local audits are for developer validation only. They do NOT provide:
- Audit history or trend analysis
- Comparison against past runs
- Regulator-ready summaries
- Signed evidence for compliance
- CI/CD enforcement capabilities

For governance enforcement, compliance evidence, and regulatory reporting,
audits must be uploaded to Verity's platform. Local runs are advisory.
Platform-recorded runs are authoritative.

For licensing questions: licensing@verity.com
        """,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    subparsers = parser.add_subparsers(dest="command", help="Command to run")
    
    # Run command
    run_parser = subparsers.add_parser("run", help="Run audit")
    run_parser.add_argument(
        "--config",
        required=True,
        help="Path to configuration YAML file"
    )
    run_parser.add_argument(
        "--fail-on-threshold",
        action="store_true",
        default=True,
        help="Exit with code 1 if any metric fails threshold (default: True)"
    )
    run_parser.add_argument(
        "--no-fail-on-threshold",
        dest="fail_on_threshold",
        action="store_false",
        help="Do not exit with code 1 on threshold failure"
    )
    run_parser.add_argument(
        "--output-dir",
        type=str,
        help="Directory to write audit evidence (overrides config)"
    )
    run_parser.add_argument(
        "--upload",
        action="store_true",
        default=False,
        help="Upload audit results (currently offline mode - no-op)"
    )
    run_parser.add_argument(
        "--no-upload",
        dest="upload",
        action="store_false",
        help="Do not upload audit results (default)"
    )
    run_parser.add_argument(
        "--export",
        choices=["pdf", "json"],
        help="Export report in specified format (validation only for local audits)"
    )
    
    args = parser.parse_args()
    
    if args.command == "run":
        passed = run_audit(
            args.config,
            fail_on_threshold=args.fail_on_threshold,
            output_dir=args.output_dir,
            upload=args.upload
        )
        
        # Handle export if requested
        if args.export:
            # Reload config to get evidence_dir
            config = Config(args.config)
            evidence_dir = Path(args.output_dir) if args.output_dir else config.evidence_dir
            report_path = evidence_dir / "report.json"
            
            if not report_path.exists():
                console.print(f"[red]Report not found: {report_path}[/red]")
                sys.exit(1)
            
            # Load report to get audit_id
            with open(report_path, 'r') as f:
                report = json.load(f)
            
            try:
                # Check if this is a local audit
                is_local = not report.get("evidence_metadata", {}).get("retention_status") == "active"
                
                if args.export == "pdf":
                    export_path = evidence_dir / f"audit_report_{report.get('audit_id', 'unknown')[:8]}.pdf"
                    export_pdf(report_path, export_path, config.config_path)
                    console.print(f"[green]PDF export written to: {export_path}[/green]")
                    if is_local:
                        console.print("[yellow]⚠️  This is a local validation export. It does not provide signed evidence or regulator-ready summaries.[/yellow]")
                elif args.export == "json":
                    export_path = evidence_dir / f"audit_report_{report.get('audit_id', 'unknown')[:8]}.json"
                    export_json(report_path, export_path, config.config_path)
                    console.print(f"[green]JSON export written to: {export_path}[/green]")
                    if is_local:
                        console.print("[yellow]⚠️  This is a local validation export. It does not provide signed evidence or regulator-ready summaries.[/yellow]")
            except Exception as e:
                console.print(f"[red]Export error: {e}[/red]")
                import traceback
                console.print(f"[red]{traceback.format_exc()}[/red]")
                sys.exit(1)
        
        # Ensure exit codes: 0 = pass, 1 = fail
        sys.exit(0 if passed else 1)
    else:
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()

