import argparse
import os
from pathlib import Path

from huggingface_hub import HfApi


def parse_args():
    parser = argparse.ArgumentParser(description="Upload a LoRA adapter folder to a Hugging Face model repo")
    parser.add_argument("--repo-id", default="ASTERIZER/LUNA-100M")
    parser.add_argument("--folder", default="Base/out/sft/rag_mcp_lora")
    parser.add_argument("--path-in-repo", default="rag_mcp_lora")
    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("HF_TOKEN is required in the environment to upload the LoRA adapter")

    folder = Path(args.folder)
    if not folder.exists():
        raise FileNotFoundError(f"LoRA output folder not found: {folder}")
    if not folder.is_dir():
        raise NotADirectoryError(f"Expected a folder, got: {folder}")

    candidate_files = [
        folder / "adapter_model.pt",
        folder / "adapter_bundle.pt",
        folder / "final" / "adapter_model.pt",
        folder / "final" / "adapter_bundle.pt",
    ]
    existing = [path for path in candidate_files if path.exists()]
    missing = [str(path) for path in candidate_files if not path.exists()]
    if not existing:
        raise FileNotFoundError(f"Missing expected adapter files: {', '.join(missing)}")

    api = HfApi(token=token)
    api.create_repo(
        repo_id=args.repo_id,
        repo_type="model",
        private=args.private,
        exist_ok=True,
    )
    api.upload_folder(
        repo_id=args.repo_id,
        repo_type="model",
        folder_path=str(folder),
        path_in_repo=args.path_in_repo,
    )
    print(
        "uploaded_lora "
        f"repo_id={args.repo_id} "
        f"folder={folder} "
        f"path_in_repo={args.path_in_repo} "
        f"url=https://huggingface.co/{args.repo_id}/tree/main/{args.path_in_repo}"
    )


if __name__ == "__main__":
    main()