#!/usr/bin/env python3
"""
Upload quantized models and samples to HuggingFace Hub.
Creates one repo per quantization variant + one combined samples repo.
"""
import os, json, argparse
from huggingface_hub import HfApi, create_repo

def upload_phase(output_dir, phase_id, repo_prefix="Swagcrew/fish-speech-s2-quant"):
    """Upload a single phase's quantized model + sample to HF Hub."""
    api = HfApi()
    phase_dir = f"{output_dir}/{phase_id}"
    samples_dir = f"{output_dir}/samples"

    if not os.path.exists(phase_dir):
        print(f"Phase dir not found: {phase_dir}")
        return

    # Create repo name
    method_map = {
        "phase1a": "fp8",
        "phase1b": "int4",
        "phase2a": "int4-all",
        "phase2b": "int8",
        "phase2c": "int3",
        "phase3a": "int2",
        "phase3b": "int2-all",
    }
    short = method_map.get(phase_id, phase_id)
    repo_id = f"{repo_prefix}-{short}"

    # Create repo
    create_repo(repo_id, repo_type="model", exist_ok=True, private=False)

    # Upload model
    model_path = f"{phase_dir}/model.safetensors"
    if os.path.exists(model_path):
        api.upload_file(model_path, f"{repo_id}/model.safetensors", repo_type="model")
        print(f"Uploaded model to {repo_id}/model.safetensors")

    # Upload results
    results_path = f"{phase_dir}/results.json"
    if os.path.exists(results_path):
        api.upload_file(results_path, f"{repo_id}/results.json", repo_type="model")

    # Upload samples
    for suffix in ["_tts.wav", "_clone.wav"]:
        sample = f"{samples_dir}/{phase_id}{suffix}"
        if os.path.exists(sample):
            api.upload_file(sample, f"{repo_id}/samples/{phase_id}{suffix}", repo_type="model")
            print(f"Uploaded sample: {phase_id}{suffix}")

    # Upload codec (shared)
    codec_path = f"{output_dir}/codec.pth"
    if os.path.exists(codec_path):
        api.upload_file(codec_path, f"{repo_id}/codec.pth", repo_type="model")

    print(f"Phase {phase_id} uploaded to {repo_id}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", default="./output", help="Output directory")
    parser.add_argument("--phase", default="all", help="Phase to upload")
    parser.add_argument("--repo-prefix", default="Swagcrew/fish-speech-s2-quant")
    args = parser.parse_args()

    phases = ["phase1a", "phase1b", "phase2a", "phase2b", "phase2c", "phase3a", "phase3b"]
    if args.phase != "all":
        phases = [f"phase{args.phase}"]

    for p in phases:
        upload_phase(args.output, p, args.repo_prefix)

if __name__ == "__main__":
    main()