# This script pre-downloads models for the Acoustic Intelligence app
import os
import sys
import logging
import torch
import traceback

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]
)

logger = logging.getLogger("model_cache_builder")

# Check if we're in a HF Space
IS_SPACE = os.environ.get("SPACE_ID") is not None
logger.info(f"Running in HF Space: {IS_SPACE}")

# Set device
if torch.cuda.is_available():
    DEVICE = "cuda:0"
    logger.info(f"Using CUDA device: {torch.cuda.get_device_name(0)}")
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
    DEVICE = "mps"
    logger.info("Using MPS (Metal Performance Shaders) device")
else:
    DEVICE = "cpu"
    logger.info("Using CPU device")

def download_nltk_data():
    """Download NLTK data for text processing."""
    logger.info("Downloading NLTK data...")
    try:
        import nltk
        nltk.download('punkt', quiet=True)
        nltk.download('stopwords', quiet=True)
        logger.info("Successfully downloaded NLTK data")
    except Exception as e:
        logger.error(f"Error downloading NLTK data: {e}")
        # Create fallback directories
        os.makedirs(os.path.expanduser("~/nltk_data/tokenizers"), exist_ok=True)
        os.makedirs(os.path.expanduser("~/nltk_data/corpora"), exist_ok=True)

def download_tts_models():
    """Download TTS models."""
    logger.info("Downloading TTS models...")
    try:
        from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
        processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
        model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
        vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
        logger.info("Successfully downloaded TTS models")
    except Exception as e:
        logger.error(f"Error downloading TTS models: {e}")

def download_asr_model():
    """Download ASR model."""
    logger.info("Downloading FunASR model...")
    try:
        from funasr import AutoModel
        # Just initialize without loading to cache the model files
        # The real loading happens at runtime
        logger.info("Successfully initialized FunASR")
    except Exception as e:
        logger.error(f"Error initializing FunASR: {e}")

def download_emotion_model():
    """Download emotion recognition model."""
    logger.info("Downloading emotion recognition model...")
    try:
        from transformers import Wav2Vec2Processor, Wav2Vec2ForSequenceClassification
        model_name = "Dpngtm/wav2vec2-emotion-recognition"
        processor = Wav2Vec2Processor.from_pretrained(model_name)
        model = Wav2Vec2ForSequenceClassification.from_pretrained(model_name)
        logger.info("Successfully downloaded emotion recognition model")
    except Exception as e:
        logger.error(f"Error downloading emotion recognition model: {e}")

def download_speaker_diarization():
    """Initialize pyannote for speaker diarization."""
    logger.info("Initializing pyannote.audio...")
    try:
        from pyannote.audio import Pipeline
        # Don't authenticate here, just download the model weights
        # The actual authentication happens at runtime with user token
        logger.info("Successfully initialized pyannote.audio")
    except Exception as e:
        logger.error(f"Error initializing pyannote.audio: {e}")

def main():
    """Main function to download all necessary models."""
    logger.info("Starting model download...")

    # Download all models
    download_nltk_data()
    download_tts_models()
    download_asr_model()
    download_emotion_model()
    download_speaker_diarization()

    logger.info("Model download complete!")

if __name__ == "__main__":
    main()