"""
Download ACE-Step model from HuggingFace
"""

import os
from pathlib import Path
from huggingface_hub import snapshot_download
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def download_model():
    """Download ACE-Step model from HuggingFace."""
    
    model_id = "ACE-Step/ACE-Step-v1-3.5B"
    cache_dir = Path("models")
    cache_dir.mkdir(exist_ok=True)
    
    logger.info(f"Downloading ACE-Step model: {model_id}")
    logger.info("This may take a while (model is ~7GB)...")
    
    try:
        model_path = snapshot_download(
            repo_id=model_id,
            cache_dir=str(cache_dir),
            resume_download=True,
            local_files_only=False
        )
        
        logger.info(f"✅ Model downloaded successfully to: {model_path}")
        return model_path
        
    except Exception as e:
        logger.error(f"❌ Failed to download model: {e}")
        logger.info("\nTroubleshooting:")
        logger.info("1. Check your internet connection")
        logger.info("2. Verify HuggingFace Hub access")
        logger.info("3. Ensure sufficient disk space (~10GB)")
        raise


if __name__ == "__main__":
    download_model()