"""
Benchmarking & Performance Profiler for Nova AI - ChatGPT Astra Edition.
Measures Time-To-First-Token (TTFT), Tokens-Per-Second (TPS),
and memory footprint across Nova Astra models.
"""

import time
import psutil
from typing import List, Dict, Any
from rich.table import Table
from rich.panel import Panel

from src.config import AVAILABLE_MODELS, GenerationConfig
from src.model_engine import NovaAstraEngine
from src.utils import console, get_system_hardware_info


def run_benchmark(
    models_to_test: List[str] = None,
    test_prompt: str = "Explain the fundamentals of quantum computing in three structured bullet points."
) -> List[Dict[str, Any]]:
    """Profile latency, generation throughput, and RAM usage for Nova Astra models."""
    if not models_to_test:
        from src.config import config
        models_to_test = [config.default_model_id]

    hw = get_system_hardware_info()
    console.print(Panel(
        f"[bold]System Hardware:[/bold] {hw['device_name']} | "
        f"[bold]Logical Cores:[/bold] {hw['cpu_count_logical']} | "
        f"[bold]Available RAM:[/bold] {hw['available_ram_gb']} GB",
        title="🌟 Nova Astra Hardware Profiler",
        border_style="cyan"
    ))

    results = []

    for model_id in models_to_test:
        console.print(f"\n[bold yellow]⏱️ Profiling Astra Model:[/bold yellow] [bold cyan]{model_id}[/bold cyan]")
        engine = NovaAstraEngine(model_id=model_id)

        # Measure Load Time & RAM
        ram_before = psutil.virtual_memory().used / (1024 ** 2)
        load_start = time.perf_counter()
        engine.load_model()
        load_time = time.perf_counter() - load_start
        ram_after = psutil.virtual_memory().used / (1024 ** 2)
        model_ram_mb = max(ram_after - ram_before, 0.0)

        # Warm-up run
        messages = [
            {"role": "system", "content": "You are Nova Astra, a concise assistant."},
            {"role": "user", "content": "Hi"}
        ]
        engine.generate_response(messages, GenerationConfig(max_new_tokens=10))

        # Actual Benchmark Run
        messages = [
            {"role": "system", "content": "You are Nova Astra, an expert assistant."},
            {"role": "user", "content": test_prompt}
        ]

        cfg = GenerationConfig(max_new_tokens=128, temperature=0.7)
        response, metrics = engine.generate_response(messages, cfg)

        record = {
            "model_id": model_id.split("/")[-1] if "/" in model_id else model_id,
            "load_time_sec": round(load_time, 2),
            "ram_mb": round(model_ram_mb, 1),
            "ttft_ms": metrics.get("time_to_first_token_ms", 0.0),
            "tokens_per_sec": metrics.get("tokens_per_second", 0.0),
            "tokens_generated": metrics.get("tokens_generated", 0),
            "total_latency_sec": metrics.get("total_latency_seconds", 0.0)
        }
        results.append(record)

    # Display Rich Comparison Table
    table = Table(title="🌟 Nova AI (ChatGPT Astra Edition) Benchmark Results", border_style="green")
    table.add_column("Model Name", style="bold cyan")
    table.add_column("Load Time (s)", justify="right")
    table.add_column("Model RAM (MB)", justify="right")
    table.add_column("TTFT (ms)", justify="right")
    table.add_column("Throughput (Tokens/s)", justify="right", style="bold green")
    table.add_column("Total Latency (s)", justify="right")

    for r in results:
        table.add_row(
            r["model_id"],
            str(r["load_time_sec"]),
            str(r["ram_mb"]),
            str(r["ttft_ms"]),
            str(r["tokens_per_sec"]),
            str(r["total_latency_sec"])
        )

    console.print(table)
    return results