import typer
import json
import subprocess
import time
from pathlib import Path
from rich.console import Console
from rich.table import Table
from robobench.core.initialization import ModelInitializationBenchmark
from robobench.core.runtime import RuntimeBenchmark
from robobench.core.advanced import AdvancedBenchmark
from robobench.core.resources import ResourceBenchmark
from robobench.core.workloads import WorkloadBenchmark
from robobench.core.quantization import QuantizationBenchmark, BenchmarkExample
from robobench.core.stability import StabilityBenchmark
from typing import List, Optional
import asyncio
from robobench.core.integration import IntegrationBenchmark
import statistics

app = typer.Typer(help="Local LLM benchmarking tool")
console = Console()

def run_cortex_command(
    command: list[str],
    timeout: int = 30,
    check_output: bool = True
) -> tuple[str, int]:
    """Run a cortex command and return its output and return code."""
    try:
        console.print(f"[yellow]Running command: {' '.join(command)}")
        result = subprocess.run(
            command,
            capture_output=True,
            text=True,
            check=check_output,
            timeout=timeout
        )
        return result.stdout, result.returncode
    except subprocess.TimeoutExpired:
        console.print(f"[red]Command timed out after {timeout} seconds: {' '.join(command)}")
        return "", 1
    except subprocess.CalledProcessError as e:
        console.print(f"[red]Command failed: {' '.join(command)}")
        console.print(f"[red]Error: {e.stderr}")
        return e.stderr, e.returncode

def reset_cortex_environment(timeout: int = 30):
    """Stop and restart cortex to ensure a clean environment."""
    console.print("[yellow]Resetting cortex environment...")

    # Stop any running instances (don't check output as it might fail if nothing is running)
    out, code = run_cortex_command(["cortex", "stop"], check_output=False)
    console.print("[green]Successfully stopped cortex")

    # Small delay to ensure everything is cleaned up
    time.sleep(2)

    # Start cortex in detached mode
    out, code = run_cortex_command(["cortex", "serve", "--detached"], timeout=timeout)
    if code != 0:
        console.print("[red]Failed to start cortex server")
        raise typer.Exit(code=1)
    console.print("[green]Successfully started cortex server")

    # Check if server is ready by pinging it
    max_retries = 5
    retry_delay = 2
    for i in range(max_retries):
        out, code = run_cortex_command(["cortex", "ps"], timeout=10)
        if code == 0:
            console.print("[green]Cortex server is ready")
            break
        else:
            if i < max_retries - 1:
                console.print(f"[yellow]Waiting for server to be ready (attempt {i+1}/{max_retries})")
                time.sleep(retry_delay)
            else:
                console.print("[red]Failed to verify server is ready")
                raise typer.Exit(code=1)

def display_init_metrics(metrics, table):
    """Display initialization metrics in a table."""
    # Basic metrics
    table.add_row("Model Name", metrics.model_name)
    table.add_row("Model Size", f"{metrics.model_size / 1024**2:.2f} MB")
    table.add_row("Quantization", metrics.quantization)
    table.add_row("Cold Start Time", f"{metrics.cold_start_time:.3f}s")
    table.add_row("Warm Start Time", f"{metrics.warm_start_time:.3f}s")
    table.add_row("Disk to RAM Time", f"{metrics.disk_to_ram_time:.3f}s")

    if metrics.model_switch_time:
        table.add_row("Model Switch Time", f"{metrics.model_switch_time:.3f}s")

    # Memory metrics
    table.add_row("Initial RAM Usage", f"{metrics.initial_memory.ram_used:.2f} MB")
    table.add_row("Peak RAM Usage", f"{metrics.peak_memory.ram_used:.2f} MB")
    table.add_row("Memory Spike", f"{metrics.memory_spike:.2f} MB")

    if metrics.initial_memory.gpu_memory_used:
        table.add_row("Initial GPU Memory",
                     f"{metrics.initial_memory.gpu_memory_used:.2f} MB")
        table.add_row("Peak GPU Memory",
                     f"{metrics.peak_memory.gpu_memory_used:.2f} MB")

def display_runtime_metrics(metrics, table):
    """Display runtime metrics in a table."""
    # Core Inference Metrics
    table.add_row("First Token Latency", f"{metrics.first_token_latency:.3f}s")
    table.add_row("Tokens per Second", f"{metrics.tokens_per_second:.2f}")
    table.add_row("Token Latency Variance", f"{metrics.token_latency_variance:.3f}ms")
    table.add_row("Streaming Throughput", f"{metrics.streaming_throughput:.2f} tokens/s")

    # Context Handling
    table.add_row("Context Window Size", str(metrics.context_window_size))
    table.add_row("KV Cache Usage", f"{metrics.kv_cache_usage:.2f}%")
    table.add_row("Memory per Token", f"{metrics.memory_per_token:.2f} MB")
    table.add_row("Context Fill Rate", f"{metrics.context_fill_rate:.2f} tokens/s")

    # Batching Performance
    table.add_row("Optimal Batch Size", str(metrics.optimal_batch_size))
    table.add_row("Queue Processing Time", f"{metrics.queue_processing_time:.3f}s")

def display_resource_metrics(metrics, table):
    """Display resource utilization metrics in a table."""
    # Memory Management
    table.add_row("Peak RAM Usage", f"{metrics.peak_memory_usage.ram_used:.2f} MB")
    table.add_row("Memory Growth Rate", f"{metrics.memory_growth_rate:.2f} MB/s")
    table.add_row("Memory Fragmentation", f"{metrics.fragmentation_ratio:.2%}")

    # CPU Metrics
    table.add_row("CPU Frequency", f"{metrics.cpu_frequency:.0f} MHz")
    if metrics.cpu_temperature:
        table.add_row("CPU Temperature", f"{metrics.cpu_temperature:.1f}°C")

    avg_cpu = sum(sum(usage)/len(usage) for usage in metrics.cpu_usage_timeline)/len(metrics.cpu_usage_timeline)
    table.add_row("Average CPU Usage", f"{avg_cpu:.1f}%")

    # GPU Metrics
    if metrics.gpu_metrics:
        for gpu_id, gpu_data in metrics.gpu_metrics.items():
            table.add_row(f"GPU {gpu_id} Utilization", f"{gpu_data['utilization']:.1f}%")
            table.add_row(f"GPU {gpu_id} Memory", f"{gpu_data['memory_used']} MB")
            table.add_row(f"GPU {gpu_id} Temperature", f"{gpu_data['temperature']}°C")

    if metrics.pcie_bandwidth_usage:
        table.add_row("PCIe Bandwidth", f"{metrics.pcie_bandwidth_usage:.2f} MB/s")

    # System Metrics
    table.add_row("Load Average (1m)", f"{metrics.system_load_average[0]:.2f}")
    table.add_row("IO Operations", f"R: {metrics.io_operations['read_count']}, W: {metrics.io_operations['write_count']}")
    table.add_row("Network Usage", f"Sent: {metrics.network_usage['bytes_sent']/1024**2:.2f} MB, Recv: {metrics.network_usage['bytes_recv']/1024**2:.2f} MB")

def display_advanced_metrics(metrics, table):
    """Display advanced processing metrics in a table."""
    # Model Sharing Metrics
    table.add_row("Memory Efficiency", f"{metrics.model_sharing.memory_efficiency:.2%}")
    table.add_row("Interference Score", f"{metrics.model_sharing.interference_score:.2f}")
    table.add_row("Context Switches", str(metrics.model_sharing.context_switches))

    # GPU Memory Sharing
    for model, memory in metrics.model_sharing.gpu_memory_sharing.items():
        table.add_row(f"Shared Memory ({model})", f"{memory:.2f} MB")

    # Concurrency Metrics
    table.add_row("Max Concurrent Requests", str(metrics.concurrency.max_concurrent_requests))
    table.add_row("Average Latency", f"{metrics.concurrency.average_latency:.3f}s")
    table.add_row("Throughput", f"{metrics.concurrency.throughput:.2f} req/s")
    table.add_row("Error Rate", f"{metrics.concurrency.error_rate:.2%}")

    # Latency Percentiles
    for p, v in metrics.concurrency.latency_percentiles.items():
        table.add_row(f"Latency {p}", f"{v:.3f}s")

    # Resource Contention
    for resource, value in metrics.concurrency.resource_contention.items():
        table.add_row(f"{resource.upper()} Contention", f"{value:.1f}%")

def display_workload_metrics(metrics, table):
    """Display workload-specific metrics in a table."""
    # Prompt handling
    table.add_row("Short Prompt Latency", f"{metrics.short_prompt_latency:.3f}s")
    table.add_row("Long Prompt Latency", f"{metrics.long_prompt_latency:.3f}s")
    table.add_row("Input Processing Speed", f"{metrics.input_tokens_per_second:.2f} tokens/s")
    table.add_row("Output Generation Speed", f"{metrics.output_tokens_per_second:.2f} tokens/s")

    # Task-specific
    table.add_row("Code Generation Speed", f"{metrics.code_generation_speed:.3f}s")
    table.add_row("Math Computation Speed", f"{metrics.math_computation_speed:.3f}s")

    # Multilingual
    for lang, speed in metrics.multilingual_performance.items():
        table.add_row(f"{lang.upper()} Performance", f"{speed:.3f}s")

    # System prompt impact
    table.add_row("Baseline Latency", f"{metrics.baseline_latency:.3f}s")
    table.add_row("System Prompt Overhead", f"{metrics.system_prompt_overhead:.3f}s")

    # Generation characteristics
    table.add_row("Token Generation Consistency", f"±{metrics.token_consistency:.3f}s")
    table.add_row("Completion Success Rate", f"{metrics.completion_success_rate:.1%}")
    table.add_row("Average Generation Length", f"{metrics.average_generation_length} tokens")

def display_quantization_metrics(metrics, table):
    """Display quantization-related metrics in a table."""
    # Model Size Metrics
    table.add_row("Original Size", f"{metrics.original_size:.2f} MB")
    table.add_row("Quantized Size", f"{metrics.quantized_size:.2f} MB")
    table.add_row("Compression Ratio", f"{metrics.compression_ratio:.2f}x")

    # Performance Metrics
    table.add_row("Inference Speed", f"{metrics.inference_speed:.2f} tokens/s")
    table.add_row("RAM Usage", f"{metrics.memory_usage:.2f} MB")
    if metrics.gpu_memory_usage:
        table.add_row("GPU Memory Usage", f"{metrics.gpu_memory_usage:.2f} MB")

    # Accuracy Metrics
    table.add_row("Perplexity", f"{metrics.perplexity:.2f}")
    table.add_row("Accuracy Score", f"{metrics.accuracy_score:.2%}")
    table.add_row("Relative Quality", f"{metrics.relative_quality:.2%}")

    # Resource Impact
    table.add_row("CPU Utilization", f"{metrics.cpu_utilization:.1f}%")
    if metrics.power_consumption:
        table.add_row("Power Consumption", f"{metrics.power_consumption:.1f}W")
    if metrics.temperature_impact:
        table.add_row("Temperature Impact", f"{metrics.temperature_impact:.1f}°C")

    # Stability Metrics
    table.add_row("Numerical Stability", f"{metrics.numerical_stability:.2%}")
    table.add_row("Error Rate", f"{metrics.error_rate:.2%}")
    table.add_row("NaN/Inf Occurrences", str(metrics.nan_infinity_occurrences))

def display_integration_metrics(metrics, table):
    """Display system integration metrics in a table."""
    # Latency Metrics
    table.add_row("Average Latency", f"{metrics.average_latency:.3f}s")
    for p, v in metrics.latency_percentiles.items():
        table.add_row(f"Latency {p}", f"{v:.3f}s")
    table.add_row("Cold Call Latency", f"{metrics.cold_call_latency:.3f}s")
    table.add_row("Warm Call Latency", f"{metrics.warm_call_latency:.3f}s")

    # Bandwidth Metrics
    table.add_row("Avg Request Size", f"{metrics.average_request_size:.2f}KB")
    table.add_row("Avg Response Size", f"{metrics.average_response_size:.2f}KB")
    table.add_row("Bandwidth Usage", f"{metrics.bandwidth_utilization:.2f}MB/s")

    # Connection Metrics
    table.add_row("Max Connections", str(metrics.max_concurrent_connections))
    table.add_row("Connection Errors", str(metrics.connection_errors))
    table.add_row("Connection Timeouts", str(metrics.connection_timeouts))
    table.add_row("Avg Connection Lifetime", f"{metrics.average_connection_lifetime:.3f}s")

    # WebSocket Metrics
    table.add_row("WebSocket Latency", f"{metrics.websocket_latency:.3f}s")
    table.add_row("Streaming Stability", f"{metrics.streaming_stability:.2%}")
    table.add_row("Message Throughput", f"{metrics.message_throughput:.2f}/s")

    # Queue Metrics
    table.add_row("Avg Queue Length", f"{metrics.average_queue_length:.1f}")
    table.add_row("Queue Wait Time", f"{metrics.queue_wait_time:.3f}s")
    table.add_row("Queue Timeout Rate", f"{metrics.queue_timeout_rate:.2%}")
    table.add_row("Request Success Rate", f"{metrics.request_success_rate:.2%}")

def display_stability_metrics(metrics, table):
    """Display stability and reliability metrics in a table."""
    # Performance Degradation
    table.add_row("Initial Throughput", f"{metrics.initial_throughput:.2f} tokens/s")
    table.add_row("Final Throughput", f"{metrics.final_throughput:.2f} tokens/s")
    table.add_row("Degradation Rate", f"{metrics.degradation_rate:.2f}%/hour")

    # Memory Management
    table.add_row("Memory Leak Detected", "Yes" if metrics.memory_leak_detected else "No")
    table.add_row("Memory Growth Rate", f"{metrics.memory_growth_rate:.2f} MB/hour")
    table.add_row("Baseline Memory", f"{metrics.baseline_memory:.2f} MB")
    table.add_row("Peak Memory", f"{metrics.peak_memory:.2f} MB")

    # Error Analysis
    table.add_row("Error Rate", f"{metrics.error_rate:.2f}/hour")
    for error_type, count in metrics.error_types.items():
        table.add_row(f"Errors - {error_type}", str(count))
    table.add_row("Crash Count", str(metrics.crash_count))

    # Thermal Management
    if metrics.temperature_profile:
        table.add_row("Temperature (avg)", f"{statistics.mean(metrics.temperature_profile):.1f}°C")
    table.add_row("Thermal Throttling", "Yes" if metrics.thermal_throttling else "No")
    table.add_row("Throttling Impact", f"{metrics.throttling_impact:.2f}%")

    # System Stability
    table.add_row("Uptime", f"{metrics.uptime:.2f} hours")
    table.add_row("Restart Count", str(metrics.restart_count))
    table.add_row("Resource Exhaustion Events", str(metrics.resource_exhaustion_events))

@app.command()
def benchmark(
    model_name: str,
    benchmark_type: str = typer.Option(
        "all",
        "--type", "-t",
        help="ype of benchmark to run (init, runtime, resources, advanced, workload, quantization, integration, stability, all)"
    ),
    reference_model: Optional[str] = typer.Option(
        None,
        "--reference-model",
        help="Reference model for quality comparison"
    ),
    benchmark_dataset: Optional[Path] = typer.Option(
        None,
        "--dataset",
        help="Path to JSON file containing benchmark examples"
    ),
    custom_prompts: Optional[Path] = typer.Option(
        None,
        "--prompts",
        help="Path to JSON file containing custom prompts"
    ),
    short_prompt_len: int = typer.Option(
        100,
        "--short-len",
        help="Target length for short prompts"
    ),
    long_prompt_len: int = typer.Option(
        2000,
        "--long-len",
        help="Target length for long prompts"
    ),
    output_len: int = typer.Option(
        500,
        "--output-len",
        help="Target length for generated outputs"
    ),
    system_prompt: str = typer.Option(
        "You are a helpful AI assistant.",
        "--system-prompt",
        help="System prompt to use for testing"
    ),
    secondary_models: List[str] = typer.Option(
            None,
            "--secondary-models",
            help="Additional models for sharing tests"
        ),
    concurrent_users: int = typer.Option(
        10,
        "--users", "-u",
        help="Number of concurrent users for load testing"
    ),
    host: str = typer.Option(
        "localhost",
        "--host",
        help="Host address for integration tests"
    ),
    port: int = typer.Option(
        39281,
        "--port",
        help="Port number for integration tests"
    ),
    stability_duration: float = typer.Option(
        1.0,
        "--stability-duration",
        help="Duration in hours for stability test"
    ),
    stability_interval: int = typer.Option(
        60,
        "--stability-interval",
        help="Sampling interval in seconds for stability test"
    ),
    verbose: bool = typer.Option(False, "--verbose", "-v", help="Show detailed metrics"),
    output_json: str = typer.Option(None, "--json", help="Save results to JSON file"),
    timeout: int = typer.Option(120, help="Timeout in seconds for operations"),
    duration: int = typer.Option(60, "--duration", "-d", help="Duration for resource monitoring in seconds")
):
    """Benchmark model initialization and performance."""
    try:
        if verbose:
            console.print(f"[blue]Starting {benchmark_type} benchmark for model: {model_name}")

        # Reset cortex environment with the specified timeout
        reset_cortex_environment(timeout=timeout)

        if verbose:
            console.print("[blue]Environment reset complete. Starting benchmarks...")

        # Wait for server to be fully ready
        time.sleep(5)

        # Check model availability
        console.print("[yellow]Checking model availability...")
        out, code = run_cortex_command(
            ["cortex", "models", "list"],
            timeout=timeout
        )
        if code != 0:
            console.print("[red]Failed to list models")
            raise typer.Exit(code=1)

        results = {}

        # Run initialization benchmarks if requested
        if benchmark_type in ["init", "all"]:
            console.print("[yellow]Running initialization benchmarks...")
            init_benchmark = ModelInitializationBenchmark(model_name)
            init_metrics = init_benchmark.benchmark_initialization()
            results["initialization"] = init_metrics

            init_table = Table(title=f"Model Initialization Metrics - {model_name}")
            init_table.add_column("Metric", justify="right", style="cyan")
            init_table.add_column("Value", justify="left", style="green")
            display_init_metrics(init_metrics, init_table)
            console.print(init_table)

        # Run runtime benchmarks if requested
        if benchmark_type in ["runtime", "all"]:
            console.print("[yellow]Running runtime benchmarks...")
            runtime_benchmark = RuntimeBenchmark(model_name)
            runtime_metrics = runtime_benchmark.benchmark_runtime()
            results["runtime"] = runtime_metrics

            runtime_table = Table(title=f"Runtime Performance Metrics - {model_name}")
            runtime_table.add_column("Metric", justify="right", style="cyan")
            runtime_table.add_column("Value", justify="left", style="green")
            display_runtime_metrics(runtime_metrics, runtime_table)
            console.print(runtime_table)

        # Run resource benchmarks if requested
        if benchmark_type in ["resources", "all"]:
            console.print(f"[yellow]Running resource benchmarks (duration: {duration}s)...")
            resource_benchmark = ResourceBenchmark(model_name, duration=duration)
            resource_metrics = resource_benchmark.benchmark_resources()
            results["resources"] = resource_metrics

            resource_table = Table(title=f"Resource Utilization Metrics - {model_name}")
            resource_table.add_column("Metric", justify="right", style="cyan")
            resource_table.add_column("Value", justify="left", style="green")
            display_resource_metrics(resource_metrics, resource_table)
            console.print(resource_table)

        # Run advanced benchmarks if requested
        if benchmark_type in ["advanced", "all"]:
            console.print("[yellow]Running advanced processing benchmarks...")
            secondary_models = secondary_models or ["tinyllama:1b-gguf-q4"]  # Default secondary model
            advanced_benchmark = AdvancedBenchmark(
                model_name,
                secondary_models,
                num_concurrent_users=concurrent_users,
                test_duration=duration
            )
            advanced_metrics = advanced_benchmark.benchmark_advanced()
            results["advanced"] = advanced_metrics

            advanced_table = Table(title=f"Advanced Processing Metrics - {model_name}")
            advanced_table.add_column("Metric", justify="right", style="cyan")
            advanced_table.add_column("Value", justify="left", style="green")
            display_advanced_metrics(advanced_metrics, advanced_table)
            console.print(advanced_table)

        # Run workload benchmarks if requested
        if benchmark_type in ["workload", "all"]:
            console.print("[yellow]Running workload-specific benchmarks...")

            # Load custom prompts if provided
            custom_prompts_dict = None
            if custom_prompts:
                try:
                    with open(custom_prompts) as f:
                        custom_prompts_dict = json.load(f)
                except (json.JSONDecodeError, FileNotFoundError) as e:
                    console.print(f"[red]Error loading custom prompts: {e}")
                    raise typer.Exit(code=1)

            # Create workload configuration
            config = WorkloadConfig(
                short_prompt_len=short_prompt_len,
                long_prompt_len=long_prompt_len,
                target_output_len=output_len,
                system_prompt=system_prompt,
                custom_prompts=custom_prompts_dict
            )

            workload_benchmark = WorkloadBenchmark(model_name, config)
            workload_metrics = workload_benchmark.benchmark_workloads()
            results["workload"] = workload_metrics

            workload_table = Table(title=f"Workload-Specific Metrics - {model_name}")
            workload_table.add_column("Metric", justify="right", style="cyan")
            workload_table.add_column("Value", justify="left", style="green")
            display_workload_metrics(workload_metrics, workload_table)
            console.print(workload_table)

        # Run quantization benchmarks if requested
        if benchmark_type in ["quantization", "all"]:
            console.print("[yellow]Running quantization benchmarks...")

            # Load custom benchmark dataset if provided
            custom_dataset = None
            if benchmark_dataset:
                try:
                    with open(benchmark_dataset) as f:
                        data = json.load(f)
                        custom_dataset = [BenchmarkExample(**example) for example in data]
                except (json.JSONDecodeError, FileNotFoundError) as e:
                    console.print(f"[red]Error loading benchmark dataset: {e}")
                    raise typer.Exit(code=1)

            quant_benchmark = QuantizationBenchmark(
                model_name,
                reference_model=reference_model,
                benchmark_dataset=custom_dataset
            )
            quant_metrics = quant_benchmark.benchmark_quantization()
            results["quantization"] = quant_metrics

            quant_table = Table(title=f"Quantization Metrics - {model_name}")
            quant_table.add_column("Metric", justify="right", style="cyan")
            quant_table.add_column("Value", justify="left", style="green")
            display_quantization_metrics(quant_metrics, quant_table)
            console.print(quant_table)

        # Run integration benchmarks if requested
        if benchmark_type in ["integration", "all"]:
            console.print("[yellow]Running system integration benchmarks...")

            integration_benchmark = IntegrationBenchmark(
                model_name,
                host=host,
                port=port,
                num_concurrent=concurrent_users,
                test_duration=duration
            )

            # Since integration benchmark uses async functions, we need to run it in an event loop
            integration_metrics = asyncio.run(
                integration_benchmark.benchmark_integration()
            )
            results["integration"] = integration_metrics

            integration_table = Table(title=f"System Integration Metrics - {model_name}")
            integration_table.add_column("Metric", justify="right", style="cyan")
            integration_table.add_column("Value", justify="left", style="green")
            display_integration_metrics(integration_metrics, integration_table)
            console.print(integration_table)

        # Run stability benchmarks if requested
        if benchmark_type in ["stability", "all"]:
            console.print("[yellow]Running stability and reliability benchmarks...")
            stability_benchmark = StabilityBenchmark(
                model_name,
                duration_hours=stability_duration,
                sampling_interval=stability_interval
            )
            stability_metrics = stability_benchmark.benchmark_stability()
            results["stability"] = stability_metrics

            stability_table = Table(title=f"Stability and Reliability Metrics - {model_name}")
            stability_table.add_column("Metric", justify="right", style="cyan")
            stability_table.add_column("Value", justify="left", style="green")
            display_stability_metrics(stability_metrics, stability_table)
            console.print(stability_table)

        # Save results to JSON if requested
        if output_json:
            with open(output_json, 'w') as f:
                json.dump(results, f, indent=2, default=str)
            console.print(f"\n[green]Results saved to: {output_json}")

    except Exception as e:
        console.print(f"[red]Error: {str(e)}")
        if verbose:
            console.print_exception()
        run_cortex_command(["cortex", "stop"], check_output=False)
        raise typer.Exit(code=1)

    finally:
        if verbose:
            console.print("[yellow]Cleaning up...")
        run_cortex_command(["cortex", "stop"], check_output=False)

def main():
    app()

if __name__ == "__main__":
    app()
