import os
import subprocess
import torch
import gradio as gr
import scipy.io.wavfile
import numpy as np
from datetime import datetime
import logging

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

# Function to install dependencies
def install_dependencies():
    """Install required dependencies, including Microsoft's custom transformers fork."""
    try:
        # Install core dependencies from requirements.txt
        subprocess.run(["pip", "install", "-r", "requirements.txt"], check=True)
        
        # Install Microsoft's custom transformers fork for VibeVoice
        subprocess.run([
            "pip", "install", "git+https://github.com/microsoft/VibeVoice.git#subdirectory=transformers"
        ], check=True)
        
        # Install flash-attn only if GPU is available
        if torch.cuda.is_available():
            subprocess.run(["pip", "install", "flash-attn==2.6.3", "--no-build-isolation"], check=True)
        
        # Install ffmpeg for audio processing
        if os.name == "posix":
            subprocess.run(["apt", "update"], check=True)
            subprocess.run(["apt", "install", "ffmpeg", "-y"], check=True)
        elif os.name == "nt":
            logger.warning("FFmpeg must be installed manually on Windows. Download from https://ffmpeg.org/download.html")
    except subprocess.CalledProcessError as e:
        logger.error(f"Failed to install dependencies: {e}")
        raise

# Verify transformers fork installation
def verify_transformers_fork():
    """Verify that the custom transformers fork is installed correctly."""
    try:
        from transformers import AutoModelForTextToSpeech
        logger.info("Custom transformers fork with AutoModelForTextToSpeech loaded successfully")
    except ImportError as e:
        logger.error(
            f"Failed to import AutoModelForTextToSpeech. Ensure the custom transformers fork is installed.\n"
            f"Run: pip install git+https://github.com/microsoft/VibeVoice.git#subdirectory=transformers\n"
            f"Error: {e}"
        )
        raise

# Determine device and attention implementation
def get_device():
    """Determine the device (CPU/GPU) and attention implementation."""
    if torch.cuda.is_available():
        logger.info(f"Using GPU: {torch.cuda.get_device_name(0)}")
        return "cuda", "flash_attention_2"
    logger.info("Using CPU. Inference may be slow and is not officially supported.")
    return "cpu", "eager"

# Load model and tokenizer
def load_model(model_name="microsoft/VibeVoice-1.5B"):
    """Load VibeVoice model and tokenizer with memory optimization."""
    from transformers import AutoModelForTextToSpeech, AutoTokenizer, AutoConfig
    
    device, attn_impl = get_device()
    try:
        # Load config with specified attention implementation
        config = AutoConfig.from_pretrained(model_name, attn_implementation=attn_impl)
        # Enable memory-efficient loading
        model = AutoModelForTextToSpeech.from_pretrained(
            model_name, config=config, torch_dtype=torch.float16 if device == "cuda" else torch.float32
        )
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = model.to(device)
        logger.info(f"Model loaded successfully on {device}")
        return model, tokenizer, device
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        raise

# Generate audio
def generate_audio(text_input, speaker_names, model, tokenizer, device, output_path=None):
    """Generate audio from text input with multi-speaker support."""
    try:
        # Validate input
        if not text_input.strip():
            raise ValueError("Input text cannot be empty")
        speaker_names = [s.strip() for s in speaker_names.split(",") if s.strip()] or ["Speaker"]
        
        # Format input text for multi-speaker dialogue
        lines = text_input.strip().split("\n")
        formatted_lines = [f"{speaker_names[i % len(speaker_names)]}: {line.strip()}" 
                         for i, line in enumerate(lines) if line.strip()]
        if not formatted_lines:
            raise ValueError("No valid text lines provided")
        formatted_text = "\n".join(formatted_lines)

        # Tokenize input
        inputs = tokenizer(formatted_text, return_tensors="pt").to(device)

        # Generate audio with memory management
        with torch.no_grad():
            if device == "cuda":
                with torch.cuda.amp.autocast():
                    audio_output = model.generate(**inputs)
            else:
                audio_output = model.generate(**inputs)

        # Convert to numpy array
        audio_np = audio_output.cpu().numpy().squeeze()
        if audio_np.ndim > 1:
            audio_np = audio_np[0]

        # Normalize audio to [-1, 1]
        max_val = np.max(np.abs(audio_np))
        if max_val > 0:
            audio_np = audio_np / max_val

        # Save or return audio
        if output_path:
            scipy.io.wavfile.write(output_path, rate=24000, data=(audio_np * 32767).astype(np.int16))
            return output_path
        return (24000, (audio_np * 32767).astype(np.int16))

    except Exception as e:
        logger.error(f"Error during audio generation: {e}")
        raise ValueError(f"Audio generation failed: {e}")

# Gradio interface function
def gradio_interface(text_input, speakers_str):
    """Gradio interface for generating audio."""
    try:
        # Generate unique output path
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = f"vibevoice_output_{timestamp}.wav"
        
        # Generate audio
        result = generate_audio(text_input, speakers_str, model, tokenizer, device, output_path)
        return result
    except ValueError as e:
        return str(e)

# Global model and tokenizer
try:
    verify_transformers_fork()
    model, tokenizer, device = load_model()
except Exception as e:
    logger.error(f"Failed to initialize: {e}")
    exit(1)

# Gradio interface setup
if __name__ == "__main__":
    install_dependencies()
    demo = gr.Interface(
        fn=gradio_interface,
        inputs=[
            gr.Textbox(
                label="Input Text (multi-line for dialogue)",
                lines=10,
                placeholder="Enter text here...\nLine 2 for next speaker...",
            ),
            gr.Textbox(
                label="Speaker Names (comma-separated, e.g., Alice,Bob)",
                value="Alice,Bob",
            ),
        ],
        outputs=gr.Audio(label="Generated Audio", type="filepath"),
        title="VibeVoice-1.5B Text-to-Speech Demo",
        description=(
            "Generate expressive multi-speaker audio from text (English/Chinese). "
            "Model adds AI disclaimer and watermark. For research use only. "
            "CPU inference is slow and not officially supported."
        ),
        examples=[
            ["Hello, how are you?\nI'm doing great, thanks!", "Alice,Bob"],
            ["欢迎来到我们的播客。\n很高兴在这里讨论AI进展。", "Chen,Li"],
        ],
    )
    demo.launch(share=False)  # Set share=True for public hosting