#!/usr/bin/env python3
import os
from typing import Dict
from huggingface_hub import hf_hub_download

# Enable hf_transfer for faster downloads
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

DEFAULT_MODELS_DIR = os.environ.get("QWEN_IMAGE_MODELS_DIR",
                                    "/Qwen-Image_models")
# Temporary download root (requested: use /tmp instead of /workspace)
TMP_DOWNLOAD_ROOT = os.environ.get("QWEN_IMAGE_TMP_DIR", "/tmp/qie_downloads")


_QWEN_IMAGE_TYPES = ("edit-2509", "edit-2511", "layered")


def _get_qwen_image_type() -> str:
    raw = os.environ.get("QWEN_IMAGE_TYPE", "edit-2509")
    v = raw.strip().lower()
    if v not in _QWEN_IMAGE_TYPES:
        print(f"[QIE] Unsupported QWEN_IMAGE_TYPE: {v}. Defaulting to edit-2509.")
        return "edit-2509"
    return v


def _ensure_dirs(root: str) -> None:
    os.makedirs(root, exist_ok=True)
    os.makedirs(os.path.join(root, "dit"), exist_ok=True)
    os.makedirs(os.path.join(root, "vae"), exist_ok=True)
    os.makedirs(os.path.join(root, "text_encoder"), exist_ok=True)


def _download_then_place(*, repo_id: str, filename: str, subfolder: str,
                         component: str, models_dir: str) -> str:
    """Download to /tmp and move to the final models_dir/component/filename.

    Returns absolute path to the final placed file.
    """
    # Ensure temp and final directories
    tmp_dir = os.path.join(TMP_DOWNLOAD_ROOT, component)
    os.makedirs(tmp_dir, exist_ok=True)
    final_dir = os.path.join(models_dir, component)
    os.makedirs(final_dir, exist_ok=True)

    print(f"[QIE] Downloading {component}: {repo_id}/{subfolder}/{filename}")
    tmp_path = hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        subfolder=subfolder,
        local_dir=tmp_dir,
    )

    final_path = os.path.join(final_dir, filename)
    try:
        # If target exists, replace it
        if os.path.exists(final_path):
            try:
                os.remove(final_path)
            except FileNotFoundError:
                pass
        # Move downloaded file into place
        if os.path.abspath(tmp_path) != os.path.abspath(final_path):
            os.replace(tmp_path, final_path)
    except Exception:
        # Fallback: copy if replace fails (e.g., cross-device)
        import shutil
        shutil.copy2(tmp_path, final_path)
    return final_path


def download_all_models(models_dir: str = DEFAULT_MODELS_DIR) -> Dict[str, str]:
    """Download required Qwen-Image models into models_dir.

    Returns a dict of component -> local file path.
    """
    _ensure_dirs(models_dir)

    print(f"[QIE] Models directory: {models_dir}")
    image_type = _get_qwen_image_type()
    print(f"[QIE] Model type: {image_type}")

    print("[QIE] Download dir (tmp):", TMP_DOWNLOAD_ROOT)
    print("[QIE] Final models dir:", models_dir)

    if image_type == "layered":
        dit_filename = "qwen_image_layered_bf16.safetensors"
        dit_path = _download_then_place(
            repo_id="Comfy-Org/Qwen-Image-Layered_ComfyUI",
            filename=dit_filename,
            subfolder="split_files/diffusion_models",
            component="dit",
            models_dir=models_dir,
        )
        print("[QIE] Downloading VAE model(s)...")
        vae_main = _download_then_place(
            repo_id="Comfy-Org/Qwen-Image-Layered_ComfyUI",
            filename="qwen_image_layered_vae.safetensors",
            subfolder="split_files/vae",
            component="vae",
            models_dir=models_dir,
        )
        vae_alt = ""
    else:
        edit_suffix = image_type.split("-", 1)[1]
        dit_filename = f"qwen_image_edit_{edit_suffix}_bf16.safetensors"
        # Download to /tmp then move to final path
        dit_path = _download_then_place(
            repo_id="Comfy-Org/Qwen-Image-Edit_ComfyUI",
            filename=dit_filename,
            subfolder="split_files/diffusion_models",
            component="dit",
            models_dir=models_dir,
        )

        print("[QIE] Downloading VAE model(s)...")
        vae_main = _download_then_place(
            repo_id="Qwen/Qwen-Image-Edit",
            filename="diffusion_pytorch_model.safetensors",
            subfolder="vae",
            component="vae",
            models_dir=models_dir,
        )
        vae_alt = _download_then_place(
            repo_id="Comfy-Org/Qwen-Image_ComfyUI",
            filename="qwen_image_vae.safetensors",
            subfolder="split_files/vae",
            component="vae",
            models_dir=models_dir,
        )

    print("[QIE] Downloading Text Encoder...")
    te_path = _download_then_place(
        repo_id="Comfy-Org/Qwen-Image_ComfyUI",
        filename="qwen_2.5_vl_7b.safetensors",
        subfolder="split_files/text_encoders",
        component="text_encoder",
        models_dir=models_dir,
    )

    print("[QIE] All models downloaded successfully!")
    return {
        "dit": dit_path,
        "vae_main": vae_main,
        "vae_alt": vae_alt,
        "text_encoder": te_path,
        "root": models_dir,
    }


if __name__ == "__main__":
    download_all_models()