from __future__ import annotations

import argparse
import os
from pathlib import Path

from huggingface_hub import HfApi

PROJECT_ROOT = Path(__file__).resolve().parents[1]
SHARE_ROOT = Path("/project/llmsvgen/share/data/kmw_layered_checkpoint")

USED_FILE_MAP = {
    PROJECT_ROOT / "demo" / "model_card.md": "README.md",
    PROJECT_ROOT / "demo" / "__init__.py": "demo/__init__.py",
    PROJECT_ROOT / "demo" / "app.py": "demo/app.py",
    PROJECT_ROOT / "demo" / "hf_repo_assets.py": "demo/hf_repo_assets.py",
    PROJECT_ROOT / "demo" / "README.md": "demo/README.md",
    PROJECT_ROOT / "demo" / "publish_space.py": "demo/publish_space.py",
    PROJECT_ROOT / "demo" / "real_world_pipeline.py": "demo/real_world_pipeline.py",
    PROJECT_ROOT / "demo" / "requirements-hf-space.txt": "demo/requirements-hf-space.txt",
    PROJECT_ROOT / "demo" / "upload_used_bundle_to_hf.py": "demo/upload_used_bundle_to_hf.py",
    PROJECT_ROOT / "demo" / "infer" / "__init__.py": "demo/infer/__init__.py",
    PROJECT_ROOT / "demo" / "infer" / "run_caption_bbox_infer.py": "demo/infer/run_caption_bbox_infer.py",
    PROJECT_ROOT / "demo" / "infer" / "vlm_bbox_inference.py": "demo/infer/vlm_bbox_inference.py",
    PROJECT_ROOT / "infer" / "__init__.py": "infer/__init__.py",
    PROJECT_ROOT / "infer" / "common_infer.py": "infer/common_infer.py",
    PROJECT_ROOT / "infer" / "infer.py": "infer/infer.py",
    PROJECT_ROOT / "infer" / "infer.yaml": "infer/infer.yaml",
    PROJECT_ROOT / "models" / "__init__.py": "models/__init__.py",
    PROJECT_ROOT / "models" / "multiLayer_adapter.py": "models/multiLayer_adapter.py",
    PROJECT_ROOT / "models" / "mmdit.py": "models/mmdit.py",
    PROJECT_ROOT / "models" / "pipeline.py": "models/pipeline.py",
    PROJECT_ROOT / "models" / "transp_vae.py": "models/transp_vae.py",
    PROJECT_ROOT / "tools" / "__init__.py": "tools/__init__.py",
    PROJECT_ROOT / "tools" / "tools.py": "tools/tools.py",
    PROJECT_ROOT / "dataset_scaleup" / "dataset_construction.sh": "dataset/dataset_construction.sh",
    PROJECT_ROOT / "dataset_scaleup" / "scaleup_api.py": "dataset/scaleup_api.py",
    PROJECT_ROOT / "dataset_scaleup" / "scaleup_dataset.py": "dataset/scaleup_dataset.py",
    PROJECT_ROOT / "dataset_scaleup" / "scaleup_utils.py": "dataset/scaleup_utils.py",
    PROJECT_ROOT / "environment.yml": "environment.yml",
}

REQUIRED_ASSET_FILES = {
    PROJECT_ROOT / "ckpt" / "trans_vae" / "0008000.pt": "ckpt/trans_vae/0008000.pt",
    PROJECT_ROOT / "ckpt" / "pre_trained_LoRA" / "pytorch_lora_weights.safetensors": "ckpt/pre_trained_LoRA/pytorch_lora_weights.safetensors",
    PROJECT_ROOT / "ckpt" / "prism_ft_LoRA" / "pytorch_lora_weights.safetensors": "ckpt/prism_ft_LoRA/pytorch_lora_weights.safetensors",
}

REQUIRED_ASSET_FOLDERS = {
    SHARE_ROOT / "SynLayers_checkpoints" / "FLUX.1-dev": "SynLayers_checkpoints/FLUX.1-dev",
    SHARE_ROOT / "SynLayers_ckpt" / "step_120000": "SynLayers_ckpt/step_120000",
    SHARE_ROOT / "SynLayers_checkpoints" / "FLUX.1-dev-Controlnet-Inpainting-Alpha": "SynLayers_checkpoints/FLUX.1-dev-Controlnet-Inpainting-Alpha",
}


def upload_files(api: HfApi, repo_id: str, repo_type: str):
    for local_path, remote_path in USED_FILE_MAP.items():
        if not local_path.exists():
            print(f"Skipping missing file: {local_path}")
            continue
        print(f"Uploading file {local_path} -> {remote_path}")
        api.upload_file(
            path_or_fileobj=str(local_path),
            path_in_repo=remote_path,
            repo_id=repo_id,
            repo_type=repo_type,
        )

    for local_path, remote_path in REQUIRED_ASSET_FILES.items():
        if not local_path.exists():
            print(f"Skipping missing asset: {local_path}")
            continue
        print(f"Uploading asset {local_path} -> {remote_path}")
        api.upload_file(
            path_or_fileobj=str(local_path),
            path_in_repo=remote_path,
            repo_id=repo_id,
            repo_type=repo_type,
        )


def upload_folders(api: HfApi, repo_id: str, repo_type: str):
    for local_path, remote_path in REQUIRED_ASSET_FOLDERS.items():
        if not local_path.exists():
            print(f"Skipping missing folder: {local_path}")
            continue
        print(f"Uploading folder {local_path} -> {remote_path}")
        api.upload_folder(
            folder_path=str(local_path),
            path_in_repo=remote_path,
            repo_id=repo_id,
            repo_type=repo_type,
        )


def main():
    parser = argparse.ArgumentParser(
        description="Upload the minimal used SynLayers demo bundle to a Hugging Face repo."
    )
    parser.add_argument("--repo-id", type=str, default="SynLayers/Bbox-caption-8b")
    parser.add_argument("--repo-type", type=str, default="model", choices=["model", "space", "dataset"])
    parser.add_argument("--token", type=str, default=os.environ.get("HF_TOKEN"))
    args = parser.parse_args()

    if not args.token:
        raise ValueError("Missing Hugging Face token. Pass --token or set HF_TOKEN.")

    api = HfApi(token=args.token)
    upload_files(api, args.repo_id, args.repo_type)
    upload_folders(api, args.repo_id, args.repo_type)

    print("")
    print(f"Finished uploading used demo bundle to https://huggingface.co/{args.repo_id}")


if __name__ == "__main__":
    main()