#!/usr/bin/env python3
"""
Setup script for EmoVoice Hugging Face Space
This script prepares the environment and downloads necessary models
"""

import os
import sys
import subprocess
import logging
from pathlib import Path

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def run_command(command, description):
    """Run a command and log the result"""
    logger.info(f"Running: {description}")
    try:
        result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True)
        logger.info(f"✅ {description} completed successfully")
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"❌ {description} failed: {e}")
        logger.error(f"Error output: {e.stderr}")
        return False

def check_python_version():
    """Check if Python version is compatible"""
    if sys.version_info < (3, 8):
        logger.error("❌ Python 3.8 or higher is required")
        return False
    logger.info(f"✅ Python version: {sys.version}")
    return True

def install_requirements():
    """Install required packages"""
    requirements_file = Path(__file__).parent / "requirements.txt"
    if not requirements_file.exists():
        logger.error("❌ requirements.txt not found")
        return False
    
    return run_command(
        f"pip install -r {requirements_file}",
        "Installing requirements"
    )

def create_directories():
    """Create necessary directories"""
    directories = [
        "models",
        "cache",
        "temp"
    ]
    
    for directory in directories:
        Path(directory).mkdir(exist_ok=True)
        logger.info(f"✅ Created directory: {directory}")
    
    return True

def setup_environment():
    """Set up environment variables"""
    env_vars = {
        "TOKENIZERS_PARALLELISM": "false",
        "OMP_NUM_THREADS": "1",
        "PYDEVD_WARN_SLOW_RESOLVE_TIMEOUT": "2",
        "CUDA_LAUNCH_BLOCKING": "1"
    }
    
    for key, value in env_vars.items():
        os.environ[key] = value
        logger.info(f"✅ Set environment variable: {key}={value}")
    
    return True

def verify_installation():
    """Verify that all required packages are installed"""
    required_packages = [
        "gradio",
        "torch",
        "torchaudio", 
        "soundfile",
        "numpy",
        "librosa",
        "huggingface_hub",
        "transformers"
    ]
    
    missing_packages = []
    for package in required_packages:
        try:
            __import__(package)
            logger.info(f"✅ {package} is installed")
        except ImportError:
            missing_packages.append(package)
            logger.error(f"❌ {package} is missing")
    
    if missing_packages:
        logger.error(f"Missing packages: {missing_packages}")
        return False
    
    return True

def main():
    """Main setup function"""
    logger.info("🚀 Starting EmoVoice Hugging Face Space setup...")
    
    # Check Python version
    if not check_python_version():
        sys.exit(1)
    
    # Create directories
    if not create_directories():
        sys.exit(1)
    
    # Set up environment
    if not setup_environment():
        sys.exit(1)
    
    # Install requirements
    if not install_requirements():
        logger.warning("⚠️ Some packages failed to install, but continuing...")
    
    # Verify installation
    if not verify_installation():
        logger.warning("⚠️ Some packages are missing, but continuing...")
    
    logger.info("✅ Setup completed successfully!")
    logger.info("🎭 You can now run the EmoVoice demo with: python app.py")

if __name__ == "__main__":
    main()