#!/usr/bin/env python3
"""
Push trained LUNA model to HuggingFace model repo.

Target: https://huggingface.co/ASTERIZER/LUNA-100M

Uploads:
  - Model weights (lit_model.pth, latest.pt)
  - Tokenizer files
  - Training config
  - Quantized GGUF files (if present)

Usage:
    HF_TOKEN=hf_xxx python push_model_to_hf.py
    HF_TOKEN=hf_xxx python push_model_to_hf.py --model_dir out/pretrain/luna-100m-english-1b
    HF_TOKEN=hf_xxx python push_model_to_hf.py --include_gguf
"""

import argparse
import os
from pathlib import Path

from huggingface_hub import HfApi


MODEL_REPO = "ASTERIZER/LUNA-100M"


def parse_args():
    parser = argparse.ArgumentParser(description="Push LUNA model to HuggingFace")
    parser.add_argument("--repo_id", default=MODEL_REPO)
    parser.add_argument("--model_dir", default="out/pretrain/luna-100m-english-1b",
                        help="Directory containing trained model")
    parser.add_argument("--path_in_repo", default="english_1b_continued",
                        help="Subfolder in HF repo")
    parser.add_argument("--include_gguf", action="store_true",
                        help="Also upload GGUF quantisations")
    parser.add_argument("--include_tokenizer", action="store_true", default=True,
                        help="Upload tokenizer files")
    parser.add_argument("--private", action="store_true")
    return parser.parse_args()


def main():
    args = parse_args()
    token = os.environ.get("HF_TOKEN")
    if not token:
        raise RuntimeError("Set HF_TOKEN environment variable")

    api = HfApi(token=token)

    # Create model repo
    api.create_repo(
        repo_id=args.repo_id,
        repo_type="model",
        private=args.private,
        exist_ok=True,
    )
    print(f"Model repo: https://huggingface.co/{args.repo_id}\n")

    total = 0
    model_dir = Path(args.model_dir)

    # ── 1. Model weights ──────────────────────────────────────────────────
    print("── Model weights ──")
    if model_dir.exists():
        # Upload the full model directory
        api.upload_folder(
            repo_id=args.repo_id,
            repo_type="model",
            folder_path=str(model_dir),
            path_in_repo=args.path_in_repo,
        )
        file_count = len(list(model_dir.rglob("*")))
        print(f"  Uploaded {file_count} files from {model_dir}")
        total += file_count
    else:
        print(f"  SKIP: {model_dir} not found")
        # Try common alternative paths
        for alt in [
            "Base/out/pretrain/luna_100m/final",
            "Base/out/pretrain/custom-100m-english/final_raw",
        ]:
            alt_path = Path(alt)
            if alt_path.exists():
                print(f"  Found alternative: {alt_path}")
                api.upload_folder(
                    repo_id=args.repo_id,
                    repo_type="model",
                    folder_path=str(alt_path),
                    path_in_repo="pretrained",
                )
                total += len(list(alt_path.rglob("*")))
                break

    # ── 2. Tokenizer ──────────────────────────────────────────────────────
    if args.include_tokenizer:
        print("\n── Tokenizer ──")
        tok_dir = Path("Base/checkpoints/EleutherAI/pythia-160m")
        tok_files = ["config.json", "tokenizer_config.json", "tokenizer.json"]
        for tf in tok_files:
            fpath = tok_dir / tf
            if fpath.exists():
                api.upload_file(
                    path_or_fileobj=str(fpath),
                    path_in_repo=f"tokenizer/{tf}",
                    repo_id=args.repo_id,
                    repo_type="model",
                )
                print(f"  OK: {fpath}")
                total += 1

    # ── 3. Training config ────────────────────────────────────────────────
    print("\n── Config ──")
    for cfg in ["train_continue_english_1b.yaml", "train_config.yaml"]:
        if os.path.exists(cfg):
            api.upload_file(
                path_or_fileobj=cfg,
                path_in_repo=f"config/{cfg}",
                repo_id=args.repo_id,
                repo_type="model",
            )
            print(f"  OK: {cfg}")
            total += 1

    # ── 4. GGUF quantisations ────────────────────────────────────────────
    if args.include_gguf:
        print("\n── GGUF quantisations ──")
        gguf_dir = Path("quantisations")
        if gguf_dir.exists():
            for gf in gguf_dir.glob("*.gguf"):
                api.upload_file(
                    path_or_fileobj=str(gf),
                    path_in_repo=f"gguf/{gf.name}",
                    repo_id=args.repo_id,
                    repo_type="model",
                )
                print(f"  OK: {gf}")
                total += 1

    print(f"\nDone! Uploaded {total} items to https://huggingface.co/{args.repo_id}")


if __name__ == "__main__":
    main()