"""
CHRONOS PQC Migration CLI Commands
==================================

Commands for post-quantum cryptography migration planning and execution
with performance tracking and rollback capabilities.
"""

import json
from pathlib import Path
from typing import Optional

import typer
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.progress import Progress

from chronos.core.defend.pqc_migrator import (
    PQCMigrator,
    PQCAlgorithm,
    ClassicalAlgorithm,
    MigrationMode,
    analyze_migration_feasibility,
)

app = typer.Typer(help="Post-quantum cryptography migration engine")
console = Console()


@app.command()
def plan(
    source: str = typer.Argument(
        "RSA-2048",
        help="Source classical algorithm (RSA-2048, ECDSA-P256, etc.)"
    ),
    target: str = typer.Argument(
        "Kyber-1024",
        help="Target PQC algorithm (Kyber-1024, Dilithium3, etc.)"
    ),
    mode: str = typer.Option(
        "full_replacement",
        "--mode",
        help="Migration mode: full_replacement, hybrid, dual_algorithm"
    ),
    output: Optional[str] = typer.Option(
        None,
        "--output",
        help="Output file for migration plan"
    ),
):
    """Plan a PQC migration."""
    console.print(f"\n[cyan]Planning migration: {source} → {target}[/cyan]\n")
    
    try:
        # Parse algorithms
        source_alg = ClassicalAlgorithm[source.replace("-", "_").upper()]
        target_alg = PQCAlgorithm[target.replace("-", "_").upper()]
        migration_mode = MigrationMode[mode.upper()]
    except KeyError as e:
        console.print(f"[red]✗ Invalid algorithm: {e}[/red]")
        raise typer.Exit(1)
    
    # Analyze feasibility
    feasibility = analyze_migration_feasibility(source_alg, target_alg)
    
    if not feasibility.get('feasible', False):
        console.print(f"[red]✗ Migration not feasible: {feasibility.get('reason')}[/red]")
        raise typer.Exit(1)
    
    # Display plan
    console.print("[bold cyan]Migration Feasibility Analysis[/bold cyan]\n")
    
    table = Table(show_header=False)
    table.add_row("Source Algorithm:", f"[yellow]{feasibility['source']}[/yellow]")
    table.add_row("Target Algorithm:", f"[green]{feasibility['target']}[/green]")
    table.add_row("Feasibility:", f"[green]✓ Feasible[/green]")
    table.add_row("Compatibility:", f"[cyan]{feasibility.get('compatibility', 'unknown')}[/cyan]")
    table.add_row("Effort Level:", f"[yellow]{feasibility.get('effort', 'unknown')}[/yellow]")
    table.add_row("Timeline:", f"[cyan]{feasibility.get('timeline_weeks', '?')} weeks[/cyan]")
    
    console.print(table)
    
    if feasibility.get('notes'):
        console.print(f"\n[dim]Notes: {feasibility['notes']}[/dim]")
    
    # Save plan if requested
    if output:
        with open(output, 'w') as f:
            json.dump(feasibility, f, indent=2)
        console.print(f"\n[green]✓ Plan saved to {output}[/green]")


@app.command()
def execute(
    source: str = typer.Argument(
        "RSA-2048",
        help="Source classical algorithm"
    ),
    target: str = typer.Argument(
        "Kyber-1024",
        help="Target PQC algorithm"
    ),
    benchmark: bool = typer.Option(
        True,
        "--benchmark",
        help="Run performance benchmarks"
    ),
    output: Optional[str] = typer.Option(
        None,
        "--output",
        help="Output file for migration report"
    ),
):
    """Execute a PQC migration."""
    console.print(f"\n[cyan]Executing migration: {source} → {target}[/cyan]\n")
    
    try:
        source_alg = ClassicalAlgorithm[source.replace("-", "_").upper()]
        target_alg = PQCAlgorithm[target.replace("-", "_").upper()]
    except KeyError:
        console.print(f"[red]✗ Invalid algorithm specified[/red]")
        raise typer.Exit(1)
    
    migrator = PQCMigrator()
    
    with Progress() as progress:
        task = progress.add_task("Creating migration plan...", total=4)
        
        # Create plan
        plan = migrator.create_migration_plan(
            source=source_alg,
            target=target_alg,
            mode=MigrationMode.FULL_REPLACEMENT,
        )
        progress.update(task, advance=1)
        
        # Start migration
        migrator.start_migration(plan.id)
        progress.update(task, advance=1, description="Starting migration...")
        
        # Benchmark if requested
        if benchmark:
            progress.update(task, advance=1, description="Benchmarking performance...")
            migrator.benchmark_migration(plan.id, iterations=10)
        else:
            progress.update(task, advance=1)
        
        # Complete migration
        migrator.complete_migration(plan.id)
        progress.update(task, advance=1, description="Completing migration...")
    
    # Display results
    console.print("\n[bold cyan]Migration Results[/bold cyan]\n")
    
    report = migrator.get_migration_report(plan.id)
    plan_info = report['plan']
    
    table = Table(show_header=False)
    table.add_row("Migration ID:", f"[cyan]{plan_info['id']}[/cyan]")
    table.add_row("Status:", f"[green]{plan_info['status'].upper()}[/green]")
    table.add_row("Source:", f"[yellow]{plan_info['source_algorithm']}[/yellow]")
    table.add_row("Target:", f"[green]{plan_info['target_algorithm']}[/green]")
    table.add_row("Duration:", f"[cyan]{plan_info['duration_seconds']:.2f}s[/cyan]" if plan_info['duration_seconds'] else "N/A")
    table.add_row("Benchmarks:", f"[cyan]{len(report['benchmarks'])}[/cyan]")
    
    console.print(table)
    
    # Display benchmarks if available
    if report['benchmarks']:
        console.print("\n[bold cyan]Performance Benchmarks[/bold cyan]\n")
        
        bench_table = Table(show_header=True, header_style="bold magenta")
        bench_table.add_column("Operation", style="cyan")
        bench_table.add_column("Duration (ms)", style="yellow")
        bench_table.add_column("Memory (MB)", style="green")
        
        for bench in report['benchmarks']:
            bench_table.add_row(
                bench['operation'],
                f"{bench['duration_ms']:.3f}",
                f"{bench['memory_mb']:.1f}",
            )
        
        console.print(bench_table)
    
    # Save report if requested
    if output:
        with open(output, 'w') as f:
            json.dump(report, f, indent=2)
        console.print(f"\n[green]✓ Report saved to {output}[/green]")
    
    console.print(f"\n[green]✓ Migration completed successfully[/green]")


@app.command()
def rollback(
    migration_id: str = typer.Argument(
        help="ID of migration to rollback"
    ),
):
    """Rollback a failed migration."""
    console.print(f"\n[yellow]Rolling back migration: {migration_id}[/yellow]\n")
    
    migrator = PQCMigrator()
    
    # Try to find migration
    migrations = migrator.list_migrations()
    target_migration = None
    for m in migrations:
        if m['id'] == migration_id:
            target_migration = m
            break
    
    if not target_migration:
        console.print(f"[red]✗ Migration not found: {migration_id}[/red]")
        raise typer.Exit(1)
    
    # Attempt rollback
    success = migrator.rollback_migration(migration_id)
    
    if success:
        console.print(f"[green]✓ Migration rolled back successfully[/green]")
        console.print(f"[dim]Snapshot restored from {migration_id}[/dim]")
    else:
        console.print(f"[red]✗ Rollback failed[/red]")
        raise typer.Exit(1)


@app.command()
def status(
    migration_id: Optional[str] = typer.Argument(
        None,
        help="Specific migration ID (optional)"
    ),
    format: str = typer.Option(
        "console",
        "--format",
        help="Output format: console, json"
    ),
):
    """Show migration status."""
    migrator = PQCMigrator()
    
    if migration_id:
        # Show specific migration
        report = migrator.get_migration_report(migration_id)
        
        if not report:
            console.print(f"[red]✗ Migration not found: {migration_id}[/red]")
            raise typer.Exit(1)
        
        if format == "json":
            console.print(json.dumps(report, indent=2))
        else:
            _display_migration_details(report)
    else:
        # Show all migrations
        migrations = migrator.list_migrations()
        
        if not migrations:
            console.print("[dim]No migrations found[/dim]")
            return
        
        if format == "json":
            console.print(json.dumps(migrations, indent=2))
        else:
            _display_migration_list(migrations)


def _display_migration_list(migrations):
    """Display list of migrations in table format."""
    console.print("\n[bold cyan]PQC Migrations[/bold cyan]\n")
    
    table = Table(show_header=True, header_style="bold magenta")
    table.add_column("ID", style="cyan", width=20)
    table.add_column("Source", style="yellow")
    table.add_column("Target", style="green")
    table.add_column("Status", width=12)
    table.add_column("Created", style="dim")
    
    for migration in migrations:
        status_color = {
            'completed': 'green',
            'in_progress': 'yellow',
            'failed': 'red',
            'rolled_back': 'orange1',
            'planned': 'dim',
        }.get(migration['status'], 'white')
        
        table.add_row(
            migration['id'][:16] + "...",
            migration['source_algorithm'],
            migration['target_algorithm'],
            f"[{status_color}]{migration['status'].upper()}[/{status_color}]",
            migration['created_at'][:10],
        )
    
    console.print(table)


def _display_migration_details(report):
    """Display detailed migration information."""
    plan = report['plan']
    
    console.print(f"\n[bold cyan]Migration Details[/bold cyan]\n")
    
    table = Table(show_header=False)
    table.add_row("ID:", f"[cyan]{plan['id']}[/cyan]")
    table.add_row("Status:", f"[green]{plan['status'].upper()}[/green]")
    table.add_row("Mode:", f"[cyan]{plan['mode']}[/cyan]")
    table.add_row("Source:", f"[yellow]{plan['source_algorithm']}[/yellow]")
    table.add_row("Target:", f"[green]{plan['target_algorithm']}[/green]")
    
    if plan['started_at']:
        table.add_row("Started:", f"[dim]{plan['started_at'][:19]}[/dim]")
    
    if plan['completed_at']:
        table.add_row("Completed:", f"[dim]{plan['completed_at'][:19]}[/dim]")
    
    if plan['duration_seconds']:
        table.add_row("Duration:", f"[cyan]{plan['duration_seconds']:.2f}s[/cyan]")
    
    table.add_row("Benchmarks:", f"[cyan]{len(report['benchmarks'])}[/cyan]")
    table.add_row("Errors:", f"[red]{plan['error_count']}[/red]" if plan['error_count'] > 0 else "[green]0[/green]")
    
    console.print(table)
    
    # Show benchmarks
    if report['benchmarks']:
        console.print("\n[bold cyan]Performance Benchmarks[/bold cyan]\n")
        
        bench_table = Table(show_header=True, header_style="bold magenta")
        bench_table.add_column("Metric", style="cyan")
        bench_table.add_column("Operation", style="yellow")
        bench_table.add_column("Duration (ms)", style="green")
        
        for bench in report['benchmarks']:
            bench_table.add_row(
                bench['metric'],
                bench['operation'],
                f"{bench['duration_ms']:.3f}",
            )
        
        console.print(bench_table)


@app.command()
def hybrid(
    source: str = typer.Argument(
        "RSA-2048",
        help="Source classical algorithm"
    ),
):
    """Create hybrid classical + PQC keypair."""
    console.print(f"\n[cyan]Creating hybrid keypair for {source}[/cyan]\n")
    
    from chronos.core.defend.pqc_migrator import HybridMigrator
    
    try:
        migrator = HybridMigrator()
        console.print("[green]✓ Hybrid migrator initialized[/green]")
    except ImportError:
        console.print("[red]✗ liboqs not installed. Install with: pip install liboqs[/red]")
        raise typer.Exit(1)
    
    # Create hybrid keypair
    if "RSA" in source:
        mock_classical = b"mock_rsa_public_key_" + source.replace("-", "_").encode()
        hybrid = migrator.create_hybrid_kex_keypair(mock_classical)
        console.print("[green]✓ Hybrid KEX keypair created[/green]")
        console.print(f"  Classical: {hybrid['classical_key_type']}")
        console.print(f"  PQC: Kyber1024")
        console.print(f"  Created: {hybrid['created_at'][:19]}")
    
    elif "ECDSA" in source:
        mock_classical = b"mock_ecdsa_public_key_" + source.replace("-", "_").encode()
        hybrid = migrator.create_hybrid_signature_keypair(mock_classical)
        console.print("[green]✓ Hybrid signature keypair created[/green]")
        console.print(f"  Classical: {hybrid['classical_key_type']}")
        console.print(f"  PQC: Dilithium3")
        console.print(f"  Created: {hybrid['created_at'][:19]}")


if __name__ == "__main__":
    app()
