import argparse
import os
import secrets
from pathlib import Path

from huggingface_hub import HfApi


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Create/update the Hugging Face Docker Space for MinerU."
    )
    parser.add_argument("--repo-id", required=True, help="Example: username/mineru-api")
    parser.add_argument("--private", action="store_true")
    parser.add_argument("--hardware", help="Example: cpu-basic, t4-small, l4")
    parser.add_argument("--sleep-time", type=int, help="Seconds before paid Space sleeps")
    parser.add_argument(
        "--mineru-api-args",
        default=os.environ.get("MINERU_API_ARGS"),
        help=(
            "Extra args passed to mineru-api. Defaults to '--device cpu' for "
            "cpu-* hardware and '--device cuda' for GPU hardware."
        ),
    )
    parser.add_argument(
        "--api-token",
        default=os.environ.get("MINERU_SPACE_API_TOKEN"),
        help="Bearer token stored as the Space API_TOKEN secret.",
    )
    parser.add_argument(
        "--hf-token",
        default=os.environ.get("HF_TOKEN"),
        help="Hugging Face write token. Defaults to HF_TOKEN.",
    )
    args = parser.parse_args()

    if not args.hf_token:
        raise SystemExit("Set HF_TOKEN or pass --hf-token.")

    api_token = args.api_token or secrets.token_urlsafe(32)
    root = Path(__file__).resolve().parents[1]
    api = HfApi(token=args.hf_token)

    api.create_repo(
        repo_id=args.repo_id,
        repo_type="space",
        space_sdk="docker",
        private=args.private,
        exist_ok=True,
    )
    api.add_space_secret(repo_id=args.repo_id, key="API_TOKEN", value=api_token)
    api.add_space_variable(
        repo_id=args.repo_id, key="ALLOW_UNAUTHENTICATED", value="false"
    )
    api.add_space_variable(
        repo_id=args.repo_id, key="MINERU_MODEL_SOURCE", value="huggingface"
    )
    api.add_space_variable(repo_id=args.repo_id, key="MAX_BODY_BYTES", value="209715200")
    api.add_space_variable(
        repo_id=args.repo_id, key="UPSTREAM_TIMEOUT_SECONDS", value="1800"
    )
    mineru_api_args = args.mineru_api_args
    if mineru_api_args is None and args.hardware:
        mineru_api_args = (
            "--device cpu" if args.hardware.startswith("cpu") else "--device cuda"
        )
    if mineru_api_args:
        api.add_space_variable(
            repo_id=args.repo_id, key="MINERU_API_ARGS", value=mineru_api_args
        )

    if args.hardware:
        api.request_space_hardware(repo_id=args.repo_id, hardware=args.hardware)
    if args.sleep_time is not None:
        api.set_space_sleep_time(repo_id=args.repo_id, sleep_time=args.sleep_time)

    api.upload_folder(
        repo_id=args.repo_id,
        repo_type="space",
        folder_path=str(root),
        ignore_patterns=[
            ".git/*",
            ".env",
            ".venv/*",
            ".DS_Store",
            "**/.DS_Store",
            "__pycache__/*",
            "data/*",
            "outputs/*",
            "client_examples/*.pdf",
            "client_examples/*.zip",
            "*.pdf",
            "*.zip",
            "*.pyc",
        ],
    )

    print(f"Uploaded {root} to https://huggingface.co/spaces/{args.repo_id}")
    if not args.api_token:
        print(f"Generated API_TOKEN: {api_token}")


if __name__ == "__main__":
    main()