"""
ACE-Step Model Downloader

This module provides functionality to download models from HuggingFace Hub or ModelScope.
It supports automatic downloading when models are not found locally,
with intelligent fallback between download sources.
"""

import os
import sys
import hashlib
import shutil
import argparse
from typing import Optional, List, Dict, Tuple
from pathlib import Path

from loguru import logger


# =============================================================================
# Model Code File Sync (GitHub repo -> checkpoint directories)
# =============================================================================

# Mapping from checkpoint directory name to source model variant in acestep/models/
_CHECKPOINT_TO_VARIANT: Dict[str, str] = {
    "acestep-v15-turbo": "turbo",
    "acestep-v15-sft": "sft",
    "acestep-v15-base": "base",
    # SFT variants (base-SFT uses the same model code as SFT)
    "acestep-v15-base-sft-fix-inst": "sft",
    # Turbo variants all share the turbo model code
    "acestep-v15-turbo-shift1": "turbo",
    "acestep-v15-turbo-shift3": "turbo",
    "acestep-v15-turbo-continuous": "turbo",
    "acestep-v15-turbo-fix-inst-shift3": "turbo",
    "acestep-v15-turbo-fix-inst-shift-continuous": "turbo",
    "acestep-v15-turbo-fix-inst-shift-dynamic": "turbo",
    "acestep-v15-turbo-rl": "turbo",
    # XL (4B DiT) variants have their own model code under acestep/models/xl_*/
    "acestep-v15-xl-base": "xl_base",
    "acestep-v15-xl-sft": "xl_sft",
    "acestep-v15-xl-turbo": "xl_turbo",
}


def _get_models_source_dir() -> Path:
    """Get the acestep/models/ directory (authoritative source for model code)."""
    return Path(__file__).resolve().parent / "models"


def _file_hash(filepath: Path) -> str:
    """Compute SHA-256 hash of a file's contents."""
    h = hashlib.sha256()
    with open(filepath, "rb") as f:
        for chunk in iter(lambda: f.read(8192), b""):
            h.update(chunk)
    return h.hexdigest()


def _check_code_mismatch(model_name: str, checkpoints_dir) -> List[str]:
    """
    Compare .py files in acestep/models/{variant}/ with those in the checkpoint directory.

    Args:
        model_name: Checkpoint directory name (e.g. "acestep-v15-turbo")
        checkpoints_dir: Path to the checkpoints root directory

    Returns:
        List of filenames that differ (empty list if all match or model_name is unknown)
    """
    variant = _CHECKPOINT_TO_VARIANT.get(model_name)
    if variant is None:
        return []

    source_dir = _get_models_source_dir() / variant
    if not source_dir.exists():
        return []

    if isinstance(checkpoints_dir, str):
        checkpoints_dir = Path(checkpoints_dir)
    target_dir = checkpoints_dir / model_name

    mismatched = []
    for src_file in source_dir.glob("*.py"):
        if src_file.name == "__init__.py":
            continue
        dst_file = target_dir / src_file.name
        if not dst_file.exists():
            mismatched.append(src_file.name)
        elif _file_hash(src_file) != _file_hash(dst_file):
            mismatched.append(src_file.name)

    return mismatched


def _sync_model_code_files(model_name: str, checkpoints_dir) -> List[str]:
    """
    Copy .py files from acestep/models/{variant}/ into the checkpoint directory,
    overwriting the HuggingFace-downloaded versions.

    Args:
        model_name: Checkpoint directory name (e.g. "acestep-v15-turbo")
        checkpoints_dir: Path to the checkpoints root directory

    Returns:
        List of filenames that were synced (empty if model_name is unknown or no source)
    """
    variant = _CHECKPOINT_TO_VARIANT.get(model_name)
    if variant is None:
        return []

    source_dir = _get_models_source_dir() / variant
    if not source_dir.exists():
        logger.warning(f"[Model Sync] Source directory not found: {source_dir}")
        return []

    if isinstance(checkpoints_dir, str):
        checkpoints_dir = Path(checkpoints_dir)
    target_dir = checkpoints_dir / model_name
    if not target_dir.exists():
        logger.warning(f"[Model Sync] Target directory not found: {target_dir}")
        return []

    synced = []
    for src_file in source_dir.glob("*.py"):
        if src_file.name == "__init__.py":
            continue
        dst_file = target_dir / src_file.name
        shutil.copy2(src_file, dst_file)
        synced.append(src_file.name)
        logger.debug(f"[Model Sync] Synced {src_file.name} -> {dst_file}")

    return synced


# =============================================================================
# Network Detection & Smart Download
# =============================================================================

def _can_access_google(timeout: float = 3.0) -> bool:
    """
    Check if Google is accessible (to determine HuggingFace vs ModelScope).

    Args:
        timeout: Connection timeout in seconds

    Returns:
        True if Google is accessible, False otherwise
    """
    import socket
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        sock.settimeout(timeout)
        sock.connect(("www.google.com", 443))
        return True
    except (socket.timeout, socket.error, OSError):
        return False
    finally:
        sock.close()


def _download_from_huggingface_internal(
    repo_id: str,
    local_dir: Path,
    token: Optional[str] = None,
) -> None:
    """
    Internal function to download from HuggingFace Hub.

    Args:
        repo_id: HuggingFace repository ID (e.g., "ACE-Step/Ace-Step1.5")
        local_dir: Local directory to save the model
        token: HuggingFace token for private repos (optional)

    Raises:
        Exception: If download fails
    """
    from huggingface_hub import snapshot_download

    logger.info(f"[Model Download] Downloading from HuggingFace: {repo_id} -> {local_dir}")

    snapshot_download(
        repo_id=repo_id,
        local_dir=str(local_dir),
        local_dir_use_symlinks="auto",
        token=token,
    )


def _download_from_modelscope_internal(
    repo_id: str,
    local_dir: Path,
) -> None:
    """
    Internal function to download from ModelScope.

    Args:
        repo_id: ModelScope repository ID (e.g., "ACE-Step/Ace-Step1.5")
        local_dir: Local directory to save the model

    Raises:
        Exception: If download fails
    """
    from modelscope import snapshot_download

    logger.info(f"[Model Download] Downloading from ModelScope: {repo_id} -> {local_dir}")

    snapshot_download(
        model_id=repo_id,
        local_dir=str(local_dir),
    )


def _smart_download(
    repo_id: str,
    local_dir: Path,
    token: Optional[str] = None,
    prefer_source: Optional[str] = None,
) -> Tuple[bool, str]:
    """
    Smart download with automatic fallback between HuggingFace and ModelScope.

    Automatically detects network environment and chooses the best download source.
    If the primary source fails, automatically falls back to the alternative.

    Args:
        repo_id: Repository ID (same format for both HF and ModelScope)
        local_dir: Local directory to save the model
        token: HuggingFace token for private repos (optional)
        prefer_source: Preferred download source ("huggingface", "modelscope", or None for auto-detect)

    Returns:
        Tuple of (success, message)
    """
    # Ensure directory exists
    local_dir.mkdir(parents=True, exist_ok=True)

    # Determine primary source
    if prefer_source == "huggingface":
        use_huggingface_first = True
        logger.info("[Model Download] User preference: HuggingFace Hub")
    elif prefer_source == "modelscope":
        use_huggingface_first = False
        logger.info("[Model Download] User preference: ModelScope")
    else:
        # Auto-detect network environment
        can_access_google = _can_access_google()
        use_huggingface_first = can_access_google
        logger.info(f"[Model Download] Auto-detected: {'HuggingFace Hub' if can_access_google else 'ModelScope'}")

    if use_huggingface_first:
        logger.info("[Model Download] Using HuggingFace Hub...")
        try:
            _download_from_huggingface_internal(repo_id, local_dir, token)
            return True, f"Successfully downloaded from HuggingFace: {repo_id}"
        except Exception as e:
            logger.warning(f"[Model Download] HuggingFace download failed: {e}")
            logger.info("[Model Download] Falling back to ModelScope...")
            try:
                _download_from_modelscope_internal(repo_id, local_dir)
                return True, f"Successfully downloaded from ModelScope: {repo_id}"
            except Exception as e2:
                error_msg = f"Both HuggingFace and ModelScope downloads failed. HF: {e}, MS: {e2}"
                logger.error(error_msg)
                return False, error_msg
    else:
        logger.info("[Model Download] Using ModelScope...")
        try:
            _download_from_modelscope_internal(repo_id, local_dir)
            return True, f"Successfully downloaded from ModelScope: {repo_id}"
        except Exception as e:
            logger.warning(f"[Model Download] ModelScope download failed: {e}")
            logger.info("[Model Download] Falling back to HuggingFace Hub...")
            try:
                _download_from_huggingface_internal(repo_id, local_dir, token)
                return True, f"Successfully downloaded from HuggingFace: {repo_id}"
            except Exception as e2:
                error_msg = f"Both ModelScope and HuggingFace downloads failed. MS: {e}, HF: {e2}"
                logger.error(error_msg)
                return False, error_msg


# =============================================================================
# Model Registry
# =============================================================================
# Main model contains core components (vae, text_encoder, default DiT)
MAIN_MODEL_REPO = "ACE-Step/Ace-Step1.5"

# Sub-models that can be downloaded separately into the checkpoints directory
SUBMODEL_REGISTRY: Dict[str, str] = {
    # LM models
    "acestep-5Hz-lm-0.6B": "ACE-Step/acestep-5Hz-lm-0.6B",
    "acestep-5Hz-lm-4B": "ACE-Step/acestep-5Hz-lm-4B",
    # DiT models
    "acestep-v15-turbo-shift3": "ACE-Step/acestep-v15-turbo-shift3",
    "acestep-v15-sft": "ACE-Step/acestep-v15-sft",
    "acestep-v15-base": "ACE-Step/acestep-v15-base",
    "acestep-v15-turbo-shift1": "ACE-Step/acestep-v15-turbo-shift1",
    "acestep-v15-turbo-continuous": "ACE-Step/acestep-v15-turbo-continuous",
    # XL (4B DiT) models
    "acestep-v15-xl-base": "ACE-Step/acestep-v15-xl-base",
    "acestep-v15-xl-sft": "ACE-Step/acestep-v15-xl-sft",
    "acestep-v15-xl-turbo": "ACE-Step/acestep-v15-xl-turbo",
}

# Components that come from the main model repo (ACE-Step/Ace-Step1.5)
MAIN_MODEL_COMPONENTS = [
    "acestep-v15-turbo",      # Default DiT model
    "vae",                     # VAE for audio encoding/decoding
    "Qwen3-Embedding-0.6B",    # Text encoder
    "acestep-5Hz-lm-1.7B",     # Default LM model (1.7B)
]

# Default LM model (included in main model)
DEFAULT_LM_MODEL = "acestep-5Hz-lm-1.7B"

# Optional community-finetuned VAE checkpoints. Each entry maps a short
# variant id (also used as the on-disk subdirectory under
# <checkpoints>/) to its HuggingFace repo id. The bundled official VAE
# stays at <checkpoints>/vae/ and is referenced as variant id "official".
VAE_REGISTRY: Dict[str, str] = {
    "scragvae": "scragnog/Ace-Step-1.5-ScragVAE",
}

# Variant id used for the bundled VAE that ships with the main model repo.
DEFAULT_VAE_VARIANT = "official"


def get_project_root() -> Path:
    """Get the project root directory.

    Returns the directory set by the ``ACESTEP_PROJECT_ROOT`` environment
    variable when present, otherwise the current working directory.  Using
    the working directory (rather than ``__file__``) keeps the checkpoints
    folder next to where the user launched the process, regardless of whether
    the package was installed via ``pip install .`` or run from source.
    """
    env_root = os.environ.get("ACESTEP_PROJECT_ROOT")
    if env_root:
        return Path(env_root).resolve()
    return Path(os.getcwd())


def get_checkpoints_dir(custom_dir: Optional[str] = None) -> Path:
    """Get the checkpoints directory path.

    Resolution order:
    1. *custom_dir* argument (passed programmatically)
    2. ``ACESTEP_CHECKPOINTS_DIR`` environment variable – allows users to
       share a single model directory across multiple ACE-Step installations,
       avoiding duplicate downloads that waste disk space.
    3. ``<project_root>/checkpoints`` (original default)
    """
    if custom_dir:
        return Path(custom_dir)
    env_dir = os.environ.get("ACESTEP_CHECKPOINTS_DIR")
    if env_dir:
        return Path(env_dir).expanduser().resolve()
    return get_project_root() / "checkpoints"


def _contains_model_weights(model_path: Path) -> bool:
    """Return whether a model directory contains at least one weights artifact.

    Args:
        model_path: Candidate model directory path.

    Returns:
        `True` when a known model weights file exists in the directory.
    """
    weight_filenames = (
        "model.safetensors",
        "model.safetensors.index.json",
        "pytorch_model.bin",
        "pytorch_model.bin.index.json",
        "diffusion_pytorch_model.safetensors",
        "diffusion_pytorch_model.safetensors.index.json",
        "diffusion_pytorch_model.bin",
        "diffusion_pytorch_model.bin.index.json",
    )
    if not model_path.is_dir():
        return False
    return any((model_path / filename).exists() for filename in weight_filenames)


def check_main_model_exists(checkpoints_dir: Optional[Path] = None) -> bool:
    """
    Check if the main model components exist in the checkpoints directory.

    Returns:
        True if all main model components contain weights, False otherwise.
    """
    if checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()
    elif isinstance(checkpoints_dir, str):
        checkpoints_dir = Path(checkpoints_dir)

    for component in MAIN_MODEL_COMPONENTS:
        component_path = checkpoints_dir / component
        if not _contains_model_weights(component_path):
            return False
    return True


def check_model_exists(model_name: str, checkpoints_dir: Optional[Path] = None) -> bool:
    """
    Check if a specific model exists in the checkpoints directory.
    
    Args:
        model_name: Name of the model to check
        checkpoints_dir: Custom checkpoints directory (optional)
    
    Returns:
        True if the model exists, False otherwise.
    """
    if not model_name:
        logger.warning("[check_model_exists] Empty model_name; treating as missing.")
        return False
    if checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()
    elif isinstance(checkpoints_dir, str):
        checkpoints_dir = Path(checkpoints_dir)

    model_path = checkpoints_dir / model_name
    return _contains_model_weights(model_path)


def list_available_models() -> Dict[str, str]:
    """
    List all available models for download.
    
    Returns:
        Dictionary mapping local names to HuggingFace repo IDs.
    """
    models = {
        "main": MAIN_MODEL_REPO,
        **SUBMODEL_REGISTRY
    }
    return models


def download_main_model(
    checkpoints_dir: Optional[Path] = None,
    force: bool = False,
    token: Optional[str] = None,
    prefer_source: Optional[str] = None,
) -> Tuple[bool, str]:
    """
    Download the main ACE-Step model from HuggingFace or ModelScope.

    The main model includes:
    - acestep-v15-turbo (default DiT model)
    - vae (audio encoder/decoder)
    - Qwen3-Embedding-0.6B (text encoder)
    - acestep-5Hz-lm-1.7B (default LM model)

    Args:
        checkpoints_dir: Custom checkpoints directory (optional)
        force: Force re-download even if model exists
        token: HuggingFace token for private repos (optional)
        prefer_source: Preferred download source ("huggingface", "modelscope", or None for auto-detect)

    Returns:
        Tuple of (success, message)
    """
    if checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()
    elif isinstance(checkpoints_dir, str):
        checkpoints_dir = Path(checkpoints_dir)

    # Ensure checkpoints directory exists
    checkpoints_dir.mkdir(parents=True, exist_ok=True)

    if not force and check_main_model_exists(checkpoints_dir):
        return True, f"Main model already exists at {checkpoints_dir}"

    print(f"Downloading main model from {MAIN_MODEL_REPO}...")
    print(f"Destination: {checkpoints_dir}")
    print("This may take a while depending on your internet connection...")

    # Use smart download with automatic fallback
    success, msg = _smart_download(MAIN_MODEL_REPO, checkpoints_dir, token, prefer_source)
    if success:
        # Sync model code files for all DiT components in the main model
        for component in MAIN_MODEL_COMPONENTS:
            if component in _CHECKPOINT_TO_VARIANT:
                synced = _sync_model_code_files(component, checkpoints_dir)
                if synced:
                    logger.info(f"[Model Download] Synced code files for {component}: {synced}")
    return success, msg


def download_submodel(
    model_name: str,
    checkpoints_dir: Optional[Path] = None,
    force: bool = False,
    token: Optional[str] = None,
    prefer_source: Optional[str] = None,
) -> Tuple[bool, str]:
    """
    Download a specific sub-model from HuggingFace or ModelScope.

    Args:
        model_name: Name of the model to download (must be in SUBMODEL_REGISTRY)
        checkpoints_dir: Custom checkpoints directory (optional)
        force: Force re-download even if model exists
        token: HuggingFace token for private repos (optional)
        prefer_source: Preferred download source ("huggingface", "modelscope", or None for auto-detect)

    Returns:
        Tuple of (success, message)
    """
    if model_name not in SUBMODEL_REGISTRY:
        available = ", ".join(SUBMODEL_REGISTRY.keys())
        return False, f"Unknown model '{model_name}'. Available models: {available}"

    if checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()
    elif isinstance(checkpoints_dir, str):
        checkpoints_dir = Path(checkpoints_dir)

    # Ensure checkpoints directory exists
    checkpoints_dir.mkdir(parents=True, exist_ok=True)

    model_path = checkpoints_dir / model_name

    if not force and model_path.exists():
        return True, f"Model '{model_name}' already exists at {model_path}"

    repo_id = SUBMODEL_REGISTRY[model_name]

    print(f"Downloading {model_name} from {repo_id}...")
    print(f"Destination: {model_path}")

    # Use smart download with automatic fallback
    success, msg = _smart_download(repo_id, model_path, token, prefer_source)
    if success and model_name in _CHECKPOINT_TO_VARIANT:
        # Sync model code files after successful download
        synced = _sync_model_code_files(model_name, checkpoints_dir)
        if synced:
            logger.info(f"[Model Download] Synced code files for {model_name}: {synced}")
    return success, msg


def download_all_models(
    checkpoints_dir: Optional[Path] = None,
    force: bool = False,
    token: Optional[str] = None,
) -> Tuple[bool, List[str]]:
    """
    Download all available models.
    
    Args:
        checkpoints_dir: Custom checkpoints directory (optional)
        force: Force re-download even if models exist
        token: HuggingFace token for private repos (optional)
    
    Returns:
        Tuple of (all_success, list of messages)
    """
    if checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()
    elif isinstance(checkpoints_dir, str):
        checkpoints_dir = Path(checkpoints_dir)

    messages = []
    all_success = True
    
    # Download main model first
    success, msg = download_main_model(checkpoints_dir, force, token)
    messages.append(msg)
    if not success:
        all_success = False
    
    # Download all sub-models
    for model_name in SUBMODEL_REGISTRY:
        success, msg = download_submodel(model_name, checkpoints_dir, force, token)
        messages.append(msg)
        if not success:
            all_success = False
    
    return all_success, messages


def ensure_main_model(
    checkpoints_dir: Optional[Path] = None,
    token: Optional[str] = None,
    prefer_source: Optional[str] = None,
) -> Tuple[bool, str]:
    """
    Ensure the main model is available, downloading if necessary.

    This function is designed to be called during initialization.
    It will only download if the model doesn't exist.

    Args:
        checkpoints_dir: Custom checkpoints directory (optional)
        token: HuggingFace token for private repos (optional)
        prefer_source: Preferred download source ("huggingface", "modelscope", or None for auto-detect)

    Returns:
        Tuple of (success, message)
    """
    if checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()

    if check_main_model_exists(checkpoints_dir):
        return True, "Main model is available"

    print("\n" + "=" * 60)
    print("Main model not found. Starting automatic download...")
    print("=" * 60 + "\n")

    return download_main_model(checkpoints_dir, token=token, prefer_source=prefer_source)


def ensure_lm_model(
    model_name: Optional[str] = None,
    checkpoints_dir: Optional[Path] = None,
    token: Optional[str] = None,
    prefer_source: Optional[str] = None,
) -> Tuple[bool, str]:
    """
    Ensure an LM model is available, downloading if necessary.

    Args:
        model_name: Name of the LM model (defaults to DEFAULT_LM_MODEL)
        checkpoints_dir: Custom checkpoints directory (optional)
        token: HuggingFace token for private repos (optional)
        prefer_source: Preferred download source ("huggingface", "modelscope", or None for auto-detect)

    Returns:
        Tuple of (success, message)
    """
    if model_name is None:
        model_name = DEFAULT_LM_MODEL

    if checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()
    elif isinstance(checkpoints_dir, str):
        checkpoints_dir = Path(checkpoints_dir)

    if check_model_exists(model_name, checkpoints_dir):
        return True, f"LM model '{model_name}' is available"

    # Check if this is a known LM model
    if model_name not in SUBMODEL_REGISTRY:
        # Check if it might be a variant name
        for known_model in SUBMODEL_REGISTRY:
            if "lm" in known_model.lower() and model_name.lower() in known_model.lower():
                model_name = known_model
                break
        else:
            return False, f"Unknown LM model: {model_name}"

    print("\n" + "=" * 60)
    print(f"LM model '{model_name}' not found. Starting automatic download...")
    print("=" * 60 + "\n")

    return download_submodel(model_name, checkpoints_dir, token=token, prefer_source=prefer_source)


def ensure_dit_model(
    model_name: str,
    checkpoints_dir: Optional[Path] = None,
    token: Optional[str] = None,
    prefer_source: Optional[str] = None,
) -> Tuple[bool, str]:
    """
    Ensure a DiT model is available, downloading if necessary.

    Args:
        model_name: Name of the DiT model
        checkpoints_dir: Custom checkpoints directory (optional)
        token: HuggingFace token for private repos (optional)
        prefer_source: Preferred download source ("huggingface", "modelscope", or None for auto-detect)

    Returns:
        Tuple of (success, message)
    """
    if checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()
    elif isinstance(checkpoints_dir, str):
        checkpoints_dir = Path(checkpoints_dir)

    if check_model_exists(model_name, checkpoints_dir):
        return True, f"DiT model '{model_name}' is available"

    # Check if this is the default turbo model (part of main)
    if model_name == "acestep-v15-turbo":
        return ensure_main_model(checkpoints_dir, token, prefer_source)

    # Check if it's a known sub-model
    if model_name in SUBMODEL_REGISTRY:
        print("\n" + "=" * 60)
        print(f"DiT model '{model_name}' not found. Starting automatic download...")
        print("=" * 60 + "\n")
        return download_submodel(model_name, checkpoints_dir, token=token, prefer_source=prefer_source)

    if not model_name:
        return False, "Unknown DiT model: '' (pass None for default or choose a valid model)"
    return False, f"Unknown DiT model: {model_name}"


def list_available_vae_variants() -> List[str]:
    """Return all selectable VAE variant ids (official first)."""
    return [DEFAULT_VAE_VARIANT, *VAE_REGISTRY.keys()]


def resolve_vae_path(checkpoint_dir: "str | Path", vae_variant: Optional[str]) -> Path:
    """Resolve a VAE variant id (or absolute path) to its on-disk directory.

    Args:
        checkpoint_dir: Root checkpoints directory.
        vae_variant: Variant id (``"official"`` or a key in ``VAE_REGISTRY``)
            or an absolute filesystem path. ``None`` / ``""`` is treated as
            ``"official"``.

    Returns:
        Absolute path of the VAE checkpoint directory.

    Raises:
        ValueError: If ``vae_variant`` is not recognized.
    """
    if isinstance(checkpoint_dir, str):
        checkpoint_dir = Path(checkpoint_dir)
    if not vae_variant:
        return checkpoint_dir / "vae"
    if os.path.isabs(vae_variant):
        return Path(vae_variant)
    if vae_variant == DEFAULT_VAE_VARIANT:
        return checkpoint_dir / "vae"
    if vae_variant in VAE_REGISTRY:
        return checkpoint_dir / vae_variant
    raise ValueError(
        f"Unknown VAE variant '{vae_variant}'. Available: "
        f"{', '.join(list_available_vae_variants())}"
    )


def check_vae_exists(vae_variant: str, checkpoints_dir: Optional[Path] = None) -> bool:
    """Return whether the requested VAE variant has weights on disk."""
    if checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()
    elif isinstance(checkpoints_dir, str):
        checkpoints_dir = Path(checkpoints_dir)
    try:
        path = resolve_vae_path(checkpoints_dir, vae_variant)
    except ValueError:
        return False
    return _contains_model_weights(path)


def download_vae(
    vae_variant: str,
    checkpoints_dir: Optional[Path] = None,
    force: bool = False,
    token: Optional[str] = None,
    prefer_source: Optional[str] = None,
) -> Tuple[bool, str]:
    """Download a community VAE variant into ``<checkpoints>/<variant>/``.

    The bundled ``"official"`` VAE is *not* downloaded here — it ships
    with the main model and is fetched by ``download_main_model``.
    """
    if vae_variant == DEFAULT_VAE_VARIANT:
        return False, (
            f"VAE variant '{DEFAULT_VAE_VARIANT}' ships with the main model; "
            "use download_main_model() instead."
        )
    if vae_variant not in VAE_REGISTRY:
        available = ", ".join(VAE_REGISTRY.keys()) or "(none)"
        return False, f"Unknown VAE variant '{vae_variant}'. Available: {available}"

    if checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()
    elif isinstance(checkpoints_dir, str):
        checkpoints_dir = Path(checkpoints_dir)
    checkpoints_dir.mkdir(parents=True, exist_ok=True)

    target_path = checkpoints_dir / vae_variant
    if not force and _contains_model_weights(target_path):
        return True, f"VAE variant '{vae_variant}' already exists at {target_path}"

    repo_id = VAE_REGISTRY[vae_variant]
    print(f"Downloading VAE '{vae_variant}' from {repo_id}...")
    print(f"Destination: {target_path}")

    return _smart_download(repo_id, target_path, token, prefer_source)


def ensure_vae_model(
    vae_variant: Optional[str] = None,
    checkpoints_dir: Optional[Path] = None,
    token: Optional[str] = None,
    prefer_source: Optional[str] = None,
) -> Tuple[bool, str]:
    """Ensure the requested VAE variant is on disk, downloading if needed.

    For ``"official"`` (or ``None``) this defers to ``ensure_main_model``
    since the bundled VAE travels with the main model. For registered
    community variants this calls ``download_vae`` when the directory
    is missing.
    """
    if not vae_variant or vae_variant == DEFAULT_VAE_VARIANT:
        return ensure_main_model(checkpoints_dir, token, prefer_source)

    if checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()
    elif isinstance(checkpoints_dir, str):
        checkpoints_dir = Path(checkpoints_dir)

    if check_vae_exists(vae_variant, checkpoints_dir):
        return True, f"VAE variant '{vae_variant}' is available"

    # Absolute paths are user-supplied and cannot be downloaded. Fail with a
    # clear, path-specific diagnostic instead of routing through download_vae,
    # which would surface a misleading "Unknown VAE variant" error.
    if os.path.isabs(vae_variant):
        path = Path(vae_variant)
        if not path.exists():
            return False, f"VAE path '{vae_variant}' does not exist."
        return False, (
            f"VAE path '{vae_variant}' does not contain VAE weights "
            "(expected diffusion_pytorch_model.safetensors)."
        )

    print("\n" + "=" * 60)
    print(f"VAE variant '{vae_variant}' not found. Starting automatic download...")
    print("=" * 60 + "\n")

    return download_vae(vae_variant, checkpoints_dir, token=token, prefer_source=prefer_source)


def print_model_list():
    """Print formatted list of available models."""
    print("\nAvailable Models for Download:")
    print("=" * 60)
    print("\nSupported Sources: HuggingFace Hub <-> ModelScope (auto-fallback)")

    print("\n[Main Model]")
    print(f"  main -> {MAIN_MODEL_REPO}")
    print("  Contains: vae, Qwen3-Embedding-0.6B, acestep-v15-turbo, acestep-5Hz-lm-1.7B")

    print("\n[Optional LM Models]")
    for name, repo in SUBMODEL_REGISTRY.items():
        if "lm" in name.lower():
            print(f"  {name} -> {repo}")

    print("\n[Optional DiT Models]")
    for name, repo in SUBMODEL_REGISTRY.items():
        if "lm" not in name.lower():
            print(f"  {name} -> {repo}")

    if VAE_REGISTRY:
        print("\n[Optional VAEs]")
        print(f"  official -> bundled in {MAIN_MODEL_REPO}")
        for name, repo in VAE_REGISTRY.items():
            print(f"  {name} -> {repo}")

    print("\n" + "=" * 60)


def main():
    """CLI entry point for model downloading."""
    parser = argparse.ArgumentParser(
        description="Download ACE-Step models with automatic fallback (HuggingFace <-> ModelScope)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  acestep-download                          # Download main model (includes LM 1.7B)
  acestep-download --all                    # Download all available models
  acestep-download --model acestep-v15-sft  # Download a specific model
  acestep-download --list                   # List all available models

Network Detection:
  Automatically detects network environment and chooses the best download source:
  - Google accessible -> HuggingFace (fallback to ModelScope)
  - Google blocked -> ModelScope (fallback to HuggingFace)

Shared checkpoints directory:
  Set ACESTEP_CHECKPOINTS_DIR to share models across multiple installations:
  export ACESTEP_CHECKPOINTS_DIR=~/ace-step-models

Alternative using huggingface-cli:
  huggingface-cli download ACE-Step/Ace-Step1.5 --local-dir ./checkpoints
  huggingface-cli download ACE-Step/acestep-5Hz-lm-0.6B --local-dir ./checkpoints/acestep-5Hz-lm-0.6B
        """
    )
    
    parser.add_argument(
        "--model", "-m",
        type=str,
        help="Specific model to download (use --list to see available models)"
    )
    parser.add_argument(
        "--all", "-a",
        action="store_true",
        help="Download all available models"
    )
    parser.add_argument(
        "--list", "-l",
        action="store_true",
        help="List all available models"
    )
    parser.add_argument(
        "--dir", "-d",
        type=str,
        default=None,
        help="Custom checkpoints directory (default: ./checkpoints)"
    )
    parser.add_argument(
        "--force", "-f",
        action="store_true",
        help="Force re-download even if model exists"
    )
    parser.add_argument(
        "--token", "-t",
        type=str,
        default=None,
        help="HuggingFace token for private repos"
    )
    parser.add_argument(
        "--skip-main",
        action="store_true",
        help="Skip downloading the main model (only download specified sub-model)"
    )
    
    args = parser.parse_args()
    
    # Handle --list
    if args.list:
        print_model_list()
        return 0
    
    # Get checkpoints directory
    checkpoints_dir = get_checkpoints_dir(args.dir) if args.dir else get_checkpoints_dir()
    print(f"Checkpoints directory: {checkpoints_dir}")
    
    # Handle --all
    if args.all:
        success, messages = download_all_models(checkpoints_dir, args.force, args.token)
        for msg in messages:
            print(msg)
        return 0 if success else 1
    
    # Handle --model
    if args.model:
        if args.model == "main":
            success, msg = download_main_model(checkpoints_dir, args.force, args.token)
        elif args.model in SUBMODEL_REGISTRY:
            # Download main model first if needed (unless --skip-main)
            if not args.skip_main and not check_main_model_exists(checkpoints_dir):
                print("Main model not found. Downloading main model first...")
                main_success, main_msg = download_main_model(checkpoints_dir, args.force, args.token)
                print(main_msg)
                if not main_success:
                    return 1
            
            success, msg = download_submodel(args.model, checkpoints_dir, args.force, args.token)
        else:
            print(f"Unknown model: {args.model}")
            print("Use --list to see available models")
            return 1
        
        print(msg)
        return 0 if success else 1
    
    # Default: download main model (includes default LM 1.7B)
    print("Downloading main model (includes vae, text encoder, DiT, and LM 1.7B)...")
    
    # Download main model
    success, msg = download_main_model(checkpoints_dir, args.force, args.token)
    print(msg)
    
    if success:
        print("\nDownload complete!")
        print(f"Models are available at: {checkpoints_dir}")
    
    return 0 if success else 1


if __name__ == "__main__":
    sys.exit(main())