#!/usr/bin/env python3
import argparse
import os
from pathlib import Path

from huggingface_hub import HfApi, upload_folder


def main():
    parser = argparse.ArgumentParser(description="Upload PixDLM release files to Hugging Face.")
    parser.add_argument("--namespace", default="WhynotHug")
    parser.add_argument("--model-repo", default="PixDLM")
    parser.add_argument("--dataset-repo", default="DRSeg")
    parser.add_argument("--private", action="store_true")
    args = parser.parse_args()

    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    if not token:
        raise SystemExit("Set HF_TOKEN in the environment before uploading.")

    root = Path(__file__).resolve().parents[1]
    api = HfApi(token=token)

    model_id = f"{args.namespace}/{args.model_repo}"
    dataset_id = f"{args.namespace}/{args.dataset_repo}"

    api.create_repo(model_id, repo_type="model", private=args.private, exist_ok=True)
    api.create_repo(dataset_id, repo_type="dataset", private=args.private, exist_ok=True)
    api.update_repo_settings(model_id, private=args.private, repo_type="model")
    api.update_repo_settings(dataset_id, private=args.private, repo_type="dataset")

    upload_folder(
        repo_id=model_id,
        repo_type="model",
        folder_path=root,
        token=token,
        ignore_patterns=[
            ".git/*",
            ".DS_Store",
            ".cache/*",
            "*/.cache/*",
            "__pycache__/*",
            "*/__pycache__/*",
            "data/*",
            "checkpoints/*",
            "outputs/*",
            "logs/*",
            "release/tmp/*",
            ".hf_token",
            "hf_token*",
            "*.pt",
            "*.pth",
            "*.bin",
            "*.safetensors",
            "*.zip",
            "*.tar",
            "*.tar.gz",
        ],
    )

    upload_folder(
        repo_id=dataset_id,
        repo_type="dataset",
        folder_path=root / "release" / "huggingface" / "dataset",
        token=token,
    )

    print(f"Uploaded model repo: https://huggingface.co/{model_id}")
    print(f"Uploaded dataset repo: https://huggingface.co/datasets/{dataset_id}")


if __name__ == "__main__":
    main()