"""
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 argparse
from typing import Optional, List, Dict, Tuple
from pathlib import Path

from loguru import logger


# =============================================================================
# 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
    try:
        socket.setdefaulttimeout(timeout)
        socket.socket(socket.AF_INET, socket.SOCK_STREAM).connect(("www.google.com", 443))
        return True
    except (socket.timeout, socket.error, OSError):
        return False


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),
        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 series (4B DiT, v0.1.6)
    "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"

# Auxiliary models stored directly under ACESTEP_MODELS_DIR.
MELBAND_ROFORMER_MODEL_NAME = "melband-roformer"
MELBAND_ROFORMER_REPO = "daydreamlive/MelBandRoFormer"
MELBAND_ROFORMER_MODEL_FILE = "MelBandRoformer_fp16.safetensors"


def get_project_root() -> Path:
    """Get the project root directory."""
    from acestep.paths import project_root
    return project_root()


def get_checkpoints_dir(custom_dir: Optional[str] = None) -> Path:
    """Get the checkpoints directory path."""
    if custom_dir:
        return Path(custom_dir)
    from acestep.paths import checkpoints_dir
    return checkpoints_dir()


def get_melband_roformer_dir(custom_dir: Optional[str] = None) -> Path:
    """Get the Mel-Band RoFormer model directory path."""
    if custom_dir:
        return Path(custom_dir)
    from acestep.paths import melband_roformer_dir
    return melband_roformer_dir()


def get_melband_roformer_model_path(
    model_dir: Optional[Path] = None,
    filename: str = MELBAND_ROFORMER_MODEL_FILE,
) -> Path:
    """Get the expected local Mel-Band RoFormer checkpoint path."""
    if model_dir is None:
        model_dir = get_melband_roformer_dir()
    return Path(model_dir) / filename


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 exist, False otherwise.
    """
    if checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()
    
    for component in MAIN_MODEL_COMPONENTS:
        component_path = checkpoints_dir / component
        if not component_path.exists():
            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 checkpoints_dir is None:
        checkpoints_dir = get_checkpoints_dir()
    
    model_path = checkpoints_dir / model_name
    return model_path.exists()


def check_melband_roformer_model_exists(
    model_dir: Optional[Path] = None,
    filename: str = MELBAND_ROFORMER_MODEL_FILE,
) -> bool:
    """Check if the Mel-Band RoFormer checkpoint exists locally."""
    return get_melband_roformer_model_path(model_dir, filename).exists()


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,
        MELBAND_ROFORMER_MODEL_NAME: MELBAND_ROFORMER_REPO,
    }
    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()

    # 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
    return _smart_download(MAIN_MODEL_REPO, checkpoints_dir, token, prefer_source)


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()

    # 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
    return _smart_download(repo_id, model_path, token, prefer_source)


def download_melband_roformer_model(
    model_dir: Optional[Path] = None,
    force: bool = False,
    token: Optional[str] = None,
) -> Tuple[bool, str]:
    """
    Download the Mel-Band RoFormer stem-separation checkpoint from HuggingFace.

    This is a single-file public model, so it intentionally avoids the
    HuggingFace/ModelScope fallback machinery used for ACE-Step checkpoints.
    The file is materialized directly under ACESTEP_MODELS_DIR.
    """
    if model_dir is None:
        model_dir = get_melband_roformer_dir()

    model_dir.mkdir(parents=True, exist_ok=True)
    model_path = get_melband_roformer_model_path(model_dir)

    if not force and model_path.exists():
        return True, f"Mel-Band RoFormer already exists at {model_path}"

    print(f"Downloading Mel-Band RoFormer from {MELBAND_ROFORMER_REPO}...")
    print(f"Destination: {model_path}")

    try:
        from huggingface_hub import hf_hub_download

        hf_hub_download(
            repo_id=MELBAND_ROFORMER_REPO,
            filename=MELBAND_ROFORMER_MODEL_FILE,
            local_dir=str(model_dir),
            token=token,
            force_download=force,
        )
    except Exception as exc:
        error_msg = f"Mel-Band RoFormer download failed: {exc}"
        logger.error(error_msg)
        return False, error_msg

    if not model_path.exists():
        return False, f"Mel-Band RoFormer download did not create {model_path}"
    return True, f"Successfully downloaded Mel-Band RoFormer to {model_path}"


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()
    
    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

    success, msg = download_melband_roformer_model(force=force, token=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()

    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()

    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)

    return False, f"Unknown DiT model: {model_name}"


def ensure_melband_roformer_model(
    model_dir: Optional[Path] = None,
    token: Optional[str] = None,
) -> Tuple[bool, str]:
    """
    Ensure the Mel-Band RoFormer checkpoint is available, downloading if needed.
    """
    if model_dir is None:
        model_dir = get_melband_roformer_dir()

    if check_melband_roformer_model_exists(model_dir):
        return True, "Mel-Band RoFormer model is available"

    print("\n" + "=" * 60)
    print("Mel-Band RoFormer model not found. Starting automatic download...")
    print("=" * 60 + "\n")

    return download_melband_roformer_model(model_dir, token=token)


def resolve_melband_roformer_model_path(
    model_dir: Optional[Path] = None,
    token: Optional[str] = None,
) -> Path:
    """Return the local Mel-Band RoFormer checkpoint path, downloading if needed."""
    if model_dir is None:
        model_dir = get_melband_roformer_dir()

    success, msg = ensure_melband_roformer_model(model_dir, token=token)
    if not success:
        raise RuntimeError(msg)
    return get_melband_roformer_model_path(model_dir)


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}")

    print("\n[Stem Separation Models]")
    print(f"  {MELBAND_ROFORMER_MODEL_NAME} -> {MELBAND_ROFORMER_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)

Default destination:
  ~/.daydream-scope/models/demon/checkpoints (override with the
  ACESTEP_MODELS_DIR environment variable or --dir)

Alternative using huggingface-cli (substitute your checkpoints dir):
  huggingface-cli download ACE-Step/Ace-Step1.5 --local-dir ~/.daydream-scope/models/demon/checkpoints
        """
    )
    
    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: "
             "~/.daydream-scope/models/demon/checkpoints, or "
             "$ACESTEP_MODELS_DIR/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 == MELBAND_ROFORMER_MODEL_NAME:
            model_dir = get_melband_roformer_dir(args.dir) if args.dir else get_melband_roformer_dir()
            success, msg = download_melband_roformer_model(
                model_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())
