#!/usr/bin/env python3
# ============================================================
# download_and_run_sinhalalm.py
#
# Downloads SinhalaLM (by HelaAI) from a private Hugging Face repo
# and runs text generation with it.
#
# ------------------------------------------------------------
# SETUP
# ------------------------------------------------------------
# 1. Create a .env file next to this script with:
#       HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxx      # a token with READ access
#                                              # to the private repo
#       HF_REPO_ID=your-username/sinhala-gpt  # same repo you uploaded to
#
# 2. Install deps:
#       pip install torch huggingface_hub safetensors python-dotenv
#
# 3. Run:
#       python download_and_run_sinhalalm.py
#       python download_and_run_sinhalalm.py --prompt "ආයුබෝවන්" --max-new-tokens 200
#
# The first run downloads and caches the repo locally (via HF's normal
# cache, under ~/.cache/huggingface). Later runs reuse the cache unless
# you pass --refresh.
# ============================================================

import argparse
import json
import sys
from pathlib import Path

import torch


def main():
    parser = argparse.ArgumentParser(description="Download & run SinhalaLM (by HelaAI) from Hugging Face")
    parser.add_argument("--repo-id", default=None, help="Overrides HF_REPO_ID from .env if given")
    parser.add_argument("--local-dir", default=None,
                         help="Where to download the repo to. Default: HF's own cache dir.")
    parser.add_argument("--refresh", action="store_true", help="Force re-download even if cached")
    parser.add_argument("--prompt", default=None,
                         help="Optional text prompt to continue. If omitted, generates from <BOS> only.")
    parser.add_argument("--max-new-tokens", type=int, default=200)
    parser.add_argument("--temperature", type=float, default=0.8)
    parser.add_argument("--top-k", type=int, default=200)
    args = parser.parse_args()

    # ---- load .env ----
    try:
        from dotenv import load_dotenv
    except ImportError:
        import subprocess
        subprocess.check_call([sys.executable, "-m", "pip", "install", "python-dotenv", "-q"])
        from dotenv import load_dotenv
    load_dotenv()

    import os
    HF_TOKEN = os.getenv("HF_TOKEN")
    HF_REPO_ID = args.repo_id or os.getenv("HF_REPO_ID")

    if not HF_TOKEN:
        sys.exit("❌ HF_TOKEN not found. Put it in a .env file (needs READ access to the repo).")
    if not HF_REPO_ID:
        sys.exit("❌ HF_REPO_ID not found. Put it in .env, or pass --repo-id you/repo-name")

    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        import subprocess
        subprocess.check_call([sys.executable, "-m", "pip", "install", "huggingface_hub", "-q"])
        from huggingface_hub import snapshot_download

    try:
        from safetensors.torch import load_file as load_safetensors
    except ImportError:
        import subprocess
        subprocess.check_call([sys.executable, "-m", "pip", "install", "safetensors", "-q"])
        from safetensors.torch import load_file as load_safetensors

    # ---- download the repo ----
    print(f"Downloading '{HF_REPO_ID}' from Hugging Face ...")
    download_kwargs = dict(repo_id=HF_REPO_ID, repo_type="model", token=HF_TOKEN)
    if args.local_dir:
        download_kwargs["local_dir"] = args.local_dir
    if args.refresh:
        download_kwargs["force_download"] = True

    local_dir = snapshot_download(**download_kwargs)
    local_dir = Path(local_dir)
    print(f"✓ Repo available at: {local_dir}")

    # ---- import the model + tokenizer classes shipped in the repo ----
    sys.path.insert(0, str(local_dir))
    from modeling_sinhalalm import SinhalaGPT, FastBPETokenizer

    # ---- load config ----
    config = json.loads((local_dir / "config.json").read_text(encoding="utf-8"))
    print(f"Model      : {config.get('model_name', 'SinhalaLM')} (by {config.get('developed_by', 'HelaAI')})")
    print(f"Params cfg : embed={config['n_embed']}, heads={config['n_heads']}, "
          f"layers={config['n_layers']}, block_size={config['block_size']}, "
          f"vocab={config['vocab_size']:,}")

    # ---- load tokenizer ----
    tokenizer = FastBPETokenizer.load(str(local_dir / "bpe_vocabulary.json"))

    # ---- build model + load weights ----
    print("Loading weights ...")
    state_dict = load_safetensors(str(local_dir / "model.safetensors"))

    model = SinhalaGPT(
        vocab_size=config["vocab_size"],
        embed_dim=config["n_embed"],
        n_heads=config["n_heads"],
        n_layers=config["n_layers"],
        block_size=config["block_size"],
        dropout=0.0,  # no dropout at inference
    )
    model.load_state_dict(state_dict)
    model.eval()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    print(f"Device     : {device}")

    # ---- build starting context ----
    bos_id = tokenizer.stoi["<BOS>"]
    if args.prompt:
        # Encode the prompt but drop the auto-added <EOS> from tokenizer.encode,
        # since we're continuing generation, not encoding a finished sequence.
        ids = tokenizer.encode(args.prompt)
        eos_id = tokenizer.stoi["<EOS>"]
        if ids and ids[-1] == eos_id:
            ids = ids[:-1]
        ctx = torch.tensor([ids], dtype=torch.long, device=device)
    else:
        ctx = torch.tensor([[bos_id]], dtype=torch.long, device=device)

    # ---- generate ----
    print("\nGenerating ...\n" + "-" * 70)
    with torch.no_grad():
        out = model.generate(
            ctx,
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature,
            top_k=args.top_k,
        )
    text = tokenizer.decode(out[0].tolist())
    print(text)
    print("-" * 70)


if __name__ == "__main__":
    main()