import os
import subprocess

# Optional Hugging Face token
HF_TOKEN = "hf_TOLKEIN"  # leave empty if not needed

# Target folders
base_dirs = {
    "Wan2.2-Fun": "models/loras/Wan2.2-Fun-Reward-LoRAs",
    "Wan2_2-T2V": "models/diffusion_models/T2V/HoloCine"
}

# Model URLs
urls = [
    "https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-high-noise-MPS.safetensors",
    "https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-low-noise-MPS.safetensors",
    "https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-low-noise-HPS2.1.safetensors",
    "https://huggingface.co/alibaba-pai/Wan2.2-Fun-Reward-LoRAs/resolve/main/Wan2.2-Fun-A14B-InP-high-noise-HPS2.1.safetensors",
    "https://huggingface.co/Kijai/WanVideo_comfy_fp8_scaled/resolve/main/T2V/HoloCine/Wan2_2-T2V-A14B-LOW-HoloCine-full_fp8_e4m3fn_scaled_KJ.safetensors",
    "https://huggingface.co/Kijai/WanVideo_comfy_fp8_scaled/resolve/main/T2V/HoloCine/Wan2_2-T2V-A14B-HIGH-HoloCine-full_fp8_e4m3fn_scaled_KJ.safetensors",
    "https://huggingface.co/QuantStack/HoloCine-GGUF/resolve/main/Full/HighNoise/HoloCine-Full-HighNoise-Q8_0.gguf",
    "https://huggingface.co/QuantStack/HoloCine-GGUF/resolve/main/Full/LowNoise/HoloCine-Full-LowNoise-Q8_0.gguf",
    "https://huggingface.co/RicketsRandom/Wan_2_2_Lora/resolve/main/Light/Wan22_A14B_T2V_LOW_Lightning_4steps_lora_250928_rank64_fp16.safetensors",
    "https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/resolve/main/split_files/loras/wan2.2_t2v_lightx2v_4steps_lora_v1.1_low_noise.safetensors",
    "https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/resolve/main/split_files/loras/wan2.2_t2v_lightx2v_4steps_lora_v1.1_high_noise.safetensors",
]

def download_model(url):
    filename = url.split("/")[-1]
    target_dir = None

    # Match destination folder
    for key, path in base_dirs.items():
        if key in filename:
            target_dir = path
            break
    if not target_dir:
        target_dir = "models/others"

    os.makedirs(target_dir, exist_ok=True)
    output_path = os.path.join(target_dir, filename)

    # Skip if file already exists *and* is not empty
    if os.path.exists(output_path) and os.path.getsize(output_path) > 1024:
        print(f"⏭️  Skipping {filename} — file already exists ({os.path.getsize(output_path)/1_000_000:.2f} MB).")
        return

    print(f"⬇️  Downloading {filename} → {target_dir}")

    cmd = [
        "wget",
        "--continue",
        "--no-check-certificate",
        "-O", output_path,
    ]

    if HF_TOKEN:
        cmd += ["--header", f"Authorization: Bearer {HF_TOKEN}"]

    cmd.append(url)

    try:
        subprocess.run(cmd, check=True)
        size = os.path.getsize(output_path) / 1_000_000
        print(f"✅  Finished {filename} ({size:.2f} MB)\n")
    except subprocess.CalledProcessError as e:
        print(f"❌  Failed {filename}: {e}\n")
        if os.path.exists(output_path) and os.path.getsize(output_path) == 0:
            os.remove(output_path)

if __name__ == "__main__":
    for link in urls:
        download_model(link)
    print("✅ All downloads attempted (skipping complete files only).")