import argparse
import os
from huggingface_hub import upload_file, upload_folder


PROJECT_FILES = [
    "handler.py",
    "infer_local.py",
    "infer_cloud.py",
    "requirements.txt",
    "README.md",
    "IMPLEMENTATION.md",
]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--force", action="store_true")
    parser.add_argument("--model-dir", type=str, default="./model")
    parser.add_argument("--repo-id", type=str, required=True)
    parser.add_argument(
        "--model-only",
        action="store_true",
        help="Upload only model artifacts and skip endpoint/helper project files.",
    )
    args = parser.parse_args()

    if not os.path.exists(args.model_dir):
        raise FileNotFoundError(
            f"Model directory not found: {args.model_dir}. Train first before upload."
        )

    upload_folder(
        folder_path=args.model_dir,
        repo_id=args.repo_id,
        repo_type="model",
        commit_message="update model",
        ignore_patterns=["checkpoint-*"] if not args.force else [],
    )
    if not args.model_only:
        for path in PROJECT_FILES:
            if os.path.exists(path):
                upload_file(
                    path_or_fileobj=path,
                    path_in_repo=path,
                    repo_id=args.repo_id,
                    repo_type="model",
                    commit_message="update endpoint helper files",
                )
    print(f"Uploaded model to: {args.repo_id}")


if __name__ == "__main__":
    main()