"""
Upload the final LoRA adapter bundle from a GPU instance to the Hugging Face model repo.

Default behavior:
  - Reads local files from Base/out/sft/rag_mcp_lora/final/
  - Uploads them into ASTERIZER/LUNA-100M under rag_mcp_lora/final/

Usage:
    python upload_lora_final_to_hf.py
    python upload_lora_final_to_hf.py --repo ASTERIZER/LUNA-100M
    python upload_lora_final_to_hf.py --source Base/out/sft/rag_mcp_lora/final --target adapters/rag_mcp_lora/final
"""

import argparse
import os
from pathlib import Path

from huggingface_hub import HfApi


DEFAULT_REPO = "ASTERIZER/LUNA-100M"
DEFAULT_SOURCE = "Base/out/sft/rag_mcp_lora/final"
DEFAULT_TARGET = "rag_mcp_lora/final"


def parse_args():
    parser = argparse.ArgumentParser(description="Upload final LoRA adapter artifacts to Hugging Face")
    parser.add_argument("--repo", default=DEFAULT_REPO, help="Destination Hugging Face model repo")
    parser.add_argument("--source", default=DEFAULT_SOURCE, help="Local final adapter directory")
    parser.add_argument("--target", default=DEFAULT_TARGET, help="Target folder inside the HF repo")
    parser.add_argument("--token", default=os.environ.get("HF_TOKEN"), help="HF token, defaults to HF_TOKEN env var")
    return parser.parse_args()


def iter_files(root: Path):
    for path in sorted(root.rglob("*")):
        if path.is_file():
            yield path


def main():
    args = parse_args()
    if not args.token:
        raise RuntimeError("Set HF_TOKEN or pass --token")

    source_dir = Path(args.source)
    if not source_dir.exists() or not source_dir.is_dir():
        raise FileNotFoundError(f"Source directory not found: {source_dir}")

    files = list(iter_files(source_dir))
    if not files:
        raise RuntimeError(f"No files found under: {source_dir}")

    api = HfApi(token=args.token)

    print(f"Uploading {len(files)} files from {source_dir} -> https://huggingface.co/{args.repo}/tree/main/{args.target}")
    for file_path in files:
        relative = file_path.relative_to(source_dir).as_posix()
        path_in_repo = f"{args.target}/{relative}" if relative != "." else args.target
        api.upload_file(
            path_or_fileobj=str(file_path),
            path_in_repo=path_in_repo,
            repo_id=args.repo,
            repo_type="model",
            token=args.token,
        )
        print(f"  OK: {relative} -> {path_in_repo}")

    print("Upload complete.")


if __name__ == "__main__":
    main()