import argparse
import os
from pathlib import Path

from huggingface_hub import HfApi


def parse_args():
    parser = argparse.ArgumentParser(description="Upload a full SFT run 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_full_sft")
    parser.add_argument("--path-in-repo", default="rag_mcp_full_sft")
    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 full SFT folder")

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

    required_files = [folder / "final" / "model.pth", folder / "latest.pt"]
    missing = [str(path) for path in required_files if not path.exists()]
    if missing:
        raise FileNotFoundError("Missing expected full SFT 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_full_sft "
        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()