"""
setup.py
Pre-download model weights to local cache.
Run this ONCE before launching the app.

Usage:
    python setup.py                    # download default (gpt2)
    python setup.py --model gpt2-medium
    python setup.py --model Qwen/Qwen2.5-0.5B
"""

import argparse
import os
from pathlib import Path

def download_model(model_name: str):
    print(f"\n{'='*55}")
    print(f"  LLM Probability Inspector — Model Setup")
    print(f"{'='*55}")
    print(f"  Model   : {model_name}")
    print(f"  Cache   : {Path.home()}/.cache/huggingface/hub")
    print(f"{'='*55}\n")

    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer
    except ImportError:
        print("[ERROR] transformers not installed.")
        print("  Run: pip install -r requirements.txt\n")
        exit(1)

    print("[1/2] Downloading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    print(f"      ✓ Tokenizer ready. Vocab size: {tokenizer.vocab_size:,}")

    print("[2/2] Downloading model weights...")
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        low_cpu_mem_usage=True,
    )

    # Count parameters
    n_params = sum(p.numel() for p in model.parameters())
    print(f"      ✓ Model ready. Parameters: {n_params/1e6:.1f}M")

    print(f"\n✅ Setup complete! '{model_name}' is cached locally.")
    print("   You can now run the app without internet:\n")
    print("   streamlit run app.py\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Pre-download model weights for LLM Probability Inspector"
    )
    parser.add_argument(
        "--model",
        type=str,
        default="gpt2",
        help="HuggingFace model ID to download (default: gpt2)",
    )
    args = parser.parse_args()
    download_model(args.model)