import argparse
import sys
from rich.console import Console
from rich.table import Table
from rich.panel import Panel

from .hardware import get_hardware_profile
from .hf_fetcher import get_model_parameters, get_model_config, ModelFetchError
from .calculator import calculate_memory_requirements, evaluate_hardware_compatibility, get_max_model_sizes, calculate_kv_cache_memory
from .recommend import generate_recommendations

def main():
    parser = argparse.ArgumentParser(
        description="Check if your local hardware can run a Hugging Face LLM.",
        epilog="Example: llm-check meta-llama/Llama-3-8b"
    )
    parser.add_argument("repo_id", nargs="?", help="Hugging Face Model ID (e.g., meta-llama/Llama-3-8b). Optional if using --recommend.")
    parser.add_argument("--recommend", action="store_true", help="Scan local hardware and recommend popular Hugging Face models you can run.")
    parser.add_argument("--context-length", "-c", type=int, default=8192, help="Context window size in tokens for accurate KV cache footprint (default: 8192).")
    parser.add_argument("--train", action="store_true", help="Calculate memory for training (Fine-Tuning/LoRA) instead of inference.")
    args = parser.parse_args()

    if not args.repo_id and not args.recommend:
        parser.print_help()
        sys.exit(1)

    console = Console()

    with console.status("[bold green]Profiling local hardware...", spinner="dots"):
        hw_profile = get_hardware_profile()

    if args.recommend:
        generate_recommendations(hw_profile, console)
        if not args.repo_id:
            return  # Exit here if they ONLY wanted recommendations
        console.print("\n" + "="*50 + "\n")
    
    with console.status(f"[bold green]Fetching metadata for {args.repo_id}...", spinner="dots"):
        try:
            params_b = get_model_parameters(args.repo_id)
            config = get_model_config(args.repo_id)
        except ModelFetchError as e:
            console.print(f"\n[bold red]Error:[/bold red] {str(e)}")
            sys.exit(1)

    console.print(f"[bold blue]Model:[/bold blue] {args.repo_id} (~{params_b:.1f}B Parameters)")

    # Create Hardware Table
    hw_table = Table(title="Detected Hardware Profile", show_header=True, header_style="bold magenta")
    hw_table.add_column("Component", style="cyan")
    hw_table.add_column("Details", style="green")

    hw_table.add_row("Operating System", hw_profile["os"])
    hw_table.add_row("System RAM", f"{hw_profile['system_ram_gb']:.1f} GB")
    
    if hw_profile["is_apple_silicon"]:
        hw_table.add_row("Apple Silicon", "Yes (Unified Memory Architect)")
        hw_table.add_row("Unified Memory", f"{hw_profile['unified_memory_gb']:.1f} GB")
    else:
        if hw_profile["has_nvidia_gpu"]:
            hw_table.add_row("NVIDIA GPU", "Detected")
            hw_table.add_row("Dedicated VRAM", f"{hw_profile['nvidia_vram_gb']:.1f} GB")
        else:
            hw_table.add_row("NVIDIA GPU", "Not Detected / CPU Only Inference")

    max_sizes = get_max_model_sizes(hw_profile)
    max_4bit = f"~{max_sizes.get('4-bit (INT4)', 0):.1f}B Parameters"
    hw_table.add_row("Supported Max Size (4-bit)", max_4bit)

    console.print(hw_table)
    console.print()

    # Calculate Memory & Compatibility
    kv_cache_gb = calculate_kv_cache_memory(config, args.context_length) if config else 0.0
    mem_reqs = calculate_memory_requirements(params_b, kv_cache_gb=kv_cache_gb, is_training=args.train)
    compats = evaluate_hardware_compatibility(mem_reqs, hw_profile)

    # Create Compatibility Table
    if args.train:
        table_title = "Training Memory Requirements & Verdict"
    elif kv_cache_gb > 0:
        table_title = f"Inference Memory Requirements & Verdict (+ {kv_cache_gb:.1f} GB for KV Cache)"
    else:
        table_title = "Inference Memory Requirements & Verdict (Includes 20% Overhead for KV Cache)"
        
    req_table = Table(title=table_title, show_header=True, header_style="bold magenta")
    req_table.add_column("Deployment Form", style="cyan")
    req_table.add_column("Required Memory", justify="right", style="green")
    req_table.add_column("Hardware Verdict", style="bold")

    for quant, req_gb in mem_reqs.items():
        verdict = compats[quant]
        
        # Colorize verdict inline
        verdict_colored = verdict
        if "🔴" in verdict:
            verdict_colored = f"[red]{verdict}[/red]"
        elif "🟡" in verdict:
            verdict_colored = f"[yellow]{verdict}[/yellow]"
        elif "🟢" in verdict:
            verdict_colored = f"[green]{verdict}[/green]"
            
        req_table.add_row(quant, f"{req_gb:.1f} GB", verdict_colored)

    console.print(req_table)
    
    # Final overall panel
    if any("🟢" in v for v in compats.values()):
        console.print(Panel("[bold green]Go:[/bold green] Your system can run this model.", border_style="green"))
    elif any("🟡" in v for v in compats.values()):
        console.print(Panel("[bold yellow]Warning:[/bold yellow] You can run this model, but only with heavy quantization or severe performance penalties (CPU offloading).", border_style="yellow"))
    else:
        console.print(Panel("[bold red]No-Go:[/bold red] Your system does not have enough memory to run this model even at 4-bit quantization.", border_style="red"))

if __name__ == "__main__":
    main()
