import os
import time
import numpy as np
import gradio as gr
import librosa
import soundfile as sf
import torch
import traceback
import threading
from spaces import GPU
from datetime import datetime

from vibevoice.modular.modeling_vibevoice_inference import VibeVoiceForConditionalGenerationInference
from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor
from vibevoice.modular.streamer import AudioStreamer
from transformers.utils import logging
from transformers import set_seed

logging.set_verbosity_info()
logger = logging.get_logger(__name__)



class VibeVoiceDemo:
     def __init__(self, model_paths: dict, device: str = "cuda", inference_steps: int = 5):
        """
        model_paths: dict like {"VibeVoice-1.5B": "microsoft/VibeVoice-1.5B",
                                "VibeVoice-1.1B": "microsoft/VibeVoice-1.1B"}
        """
        self.model_paths = model_paths
        self.device = device
        self.inference_steps = inference_steps

        self.is_generating = False

        # Multi-model holders
        self.models = {}        # name -> model
        self.processors = {}    # name -> processor
        self.current_model_name = None

        self.available_voices = {}

        self.load_models()          # load all on CPU
        self.setup_voice_presets()
        self.load_example_scripts()

    def load_models(self):
        print("Loading processors and models on CPU...")
        for name, path in self.model_paths.items():
            print(f" - {name} from {path}")
            proc = VibeVoiceProcessor.from_pretrained(path)
            mdl = VibeVoiceForConditionalGenerationInference.from_pretrained(
                path, torch_dtype=torch.bfloat16
            )
            # Keep on CPU initially
            self.processors[name] = proc
            self.models[name] = mdl
        # choose default
        self.current_model_name = next(iter(self.models))
        print(f"Default model is {self.current_model_name}")

    def _place_model(self, target_name: str):
        """
        Move the selected model to CUDA and push all others back to CPU.
        """
        for name, mdl in self.models.items():
            if name == target_name:
                self.models[name] = mdl.to(self.device)
            else:
                self.models[name] = mdl.to("cpu")
        self.current_model_name = target_name
        print(f"Model {target_name} is now on {self.device}. Others moved to CPU.")

    def setup_voice_presets(self):
        voices_dir = os.path.join(os.path.dirname(__file__), "voices")
        if not os.path.exists(voices_dir):
            print(f"Warning: Voices directory not found at {voices_dir}")
            return
        wav_files = [f for f in os.listdir(voices_dir)
                     if f.lower().endswith(('.wav', '.mp3', '.flac', '.ogg', '.m4a', '.aac'))]
        for wav_file in wav_files:
            name = os.path.splitext(wav_file)[0]
            self.available_voices[name] = os.path.join(voices_dir, wav_file)
        print(f"Voices loaded: {list(self.available_voices.keys())}")

    def read_audio(self, audio_path: str, target_sr: int = 24000) -> np.ndarray:
        try:
            wav, sr = sf.read(audio_path)
            if len(wav.shape) > 1:
                wav = np.mean(wav, axis=1)
            if sr != target_sr:
                wav = librosa.resample(wav, orig_sr=sr, target_sr=target_sr)
            return wav
        except Exception as e:
            print(f"Error reading audio {audio_path}: {e}")
            return np.array([])

    @GPU(duration=60)
    def generate_podcast(self,
                         num_speakers: int,
                         script: str,
                         speaker_1: str = None,
                         speaker_2: str = None,
                         speaker_3: str = None,
                         speaker_4: str = None,
                         cfg_scale: float = 1.3,
                         model_name: str = None):
        """
        Generates a podcast as a single audio file from a script and saves it.
        Non-streaming.
        """
        try:
            # pick model
            model_name = model_name or self.current_model_name
            if model_name not in self.models:
                raise gr.Error(f"Unknown model: {model_name}")

            # place models on devices
            self._place_model(model_name)
            model = self.models[model_name]
            processor = self.processors[model_name]

            print(f"Using model {model_name} on {self.device}")

            model.eval()
            model.set_ddpm_inference_steps(num_steps=self.inference_steps)

            self.is_generating = True

            if not script.strip():
                raise gr.Error("Error: Please provide a script.")

            script = script.replace("’", "'")

            if not 1 <= num_speakers <= 4:
                raise gr.Error("Error: Number of speakers must be between 1 and 4.")

            selected_speakers = [speaker_1, speaker_2, speaker_3, speaker_4][:num_speakers]
            for i, speaker_name in enumerate(selected_speakers):
                if not speaker_name or speaker_name not in self.available_voices:
                    raise gr.Error(f"Error: Please select a valid speaker for Speaker {i+1}.")

            log = f"🎙️ Generating podcast with {num_speakers} speakers\n"
            log += f"🧠 Model: {model_name}\n"
            log += f"📊 Parameters: CFG Scale={cfg_scale}\n"
            log += f"🎭 Speakers: {', '.join(selected_speakers)}\n"

            voice_samples = []
            for speaker_name in selected_speakers:
                audio_path = self.available_voices[speaker_name]
                audio_data = self.read_audio(audio_path)
                if len(audio_data) == 0:
                    raise gr.Error(f"Error: Failed to load audio for {speaker_name}")
                voice_samples.append(audio_data)

            log += f"✅ Loaded {len(voice_samples)} voice samples\n"

            lines = script.strip().split('\n')
            formatted_script_lines = []
            for line in lines:
                line = line.strip()
                if not line:
                    continue
                if line.startswith('Speaker ') and ':' in line:
                    formatted_script_lines.append(line)
                else:
                    speaker_id = len(formatted_script_lines) % num_speakers
                    formatted_script_lines.append(f"Speaker {speaker_id}: {line}")

            formatted_script = '\n'.join(formatted_script_lines)
            log += f"📝 Formatted script with {len(formatted_script_lines)} turns\n"
            log += "🔄 Processing with VibeVoice...\n"

            inputs = processor(
                text=[formatted_script],
                voice_samples=[voice_samples],
                padding=True,
                return_tensors="pt",
                return_attention_mask=True,
            )

            start_time = time.time()
            outputs = model.generate(
                **inputs,
                max_new_tokens=None,
                cfg_scale=cfg_scale,
                tokenizer=processor.tokenizer,
                generation_config={'do_sample': False},
                verbose=False,
            )
            generation_time = time.time() - start_time

            if hasattr(outputs, 'speech_outputs') and outputs.speech_outputs[0] is not None:
                audio_tensor = outputs.speech_outputs[0]
                audio = audio_tensor.cpu().float().numpy()
            else:
                raise gr.Error("❌ Error: No audio was generated by the model. Please try again.")

            if audio.ndim > 1:
                audio = audio.squeeze()

            sample_rate = 24000

            output_dir = "outputs"
            os.makedirs(output_dir, exist_ok=True)
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            file_path = os.path.join(output_dir, f"podcast_{timestamp}.wav")
            sf.write(file_path, audio, sample_rate)
            print(f"💾 Podcast saved to {file_path}")

            total_duration = len(audio) / sample_rate
            log += f"⏱️ Generation completed in {generation_time:.2f} seconds\n"
            log += f"🎵 Final audio duration: {total_duration:.2f} seconds\n"
            log += f"✅ Successfully saved podcast to: {file_path}\n"

            self.is_generating = False
            return (sample_rate, audio), log

        except gr.Error as e:
            self.is_generating = False
            error_msg = f"❌ Input Error: {str(e)}"
            print(error_msg)
            return None, error_msg

        except Exception as e:
            self.is_generating = False
            error_msg = f"❌ An unexpected error occurred: {str(e)}"
            print(error_msg)
            traceback.print_exc()
            return None, error_msg




    def load_example_scripts(self):
        """Load example scripts from the text_examples directory."""
        examples_dir = os.path.join(os.path.dirname(__file__), "text_examples")
        self.example_scripts = []
        
        # Check if text_examples directory exists
        if not os.path.exists(examples_dir):
            print(f"Warning: text_examples directory not found at {examples_dir}")
            return
        
        # Get all .txt files in the text_examples directory
        txt_files = sorted([f for f in os.listdir(examples_dir) 
                          if f.lower().endswith('.txt') and os.path.isfile(os.path.join(examples_dir, f))])
        
        for txt_file in txt_files:
            file_path = os.path.join(examples_dir, txt_file)
            
            import re
            # Check if filename contains a time pattern like "45min", "90min", etc.
            time_pattern = re.search(r'(\d+)min', txt_file.lower())
            if time_pattern:
                minutes = int(time_pattern.group(1))
                if minutes > 15:
                    print(f"Skipping {txt_file}: duration {minutes} minutes exceeds 15-minute limit")
                    continue

            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    script_content = f.read().strip()
                
                # Remove empty lines and lines with only whitespace
                script_content = '\n'.join(line for line in script_content.split('\n') if line.strip())
                
                if not script_content:
                    continue
                
                # Parse the script to determine number of speakers
                num_speakers = self._get_num_speakers_from_script(script_content)
                
                # Add to examples list as [num_speakers, script_content]
                self.example_scripts.append([num_speakers, script_content])
                print(f"Loaded example: {txt_file} with {num_speakers} speakers")
                
            except Exception as e:
                print(f"Error loading example script {txt_file}: {e}")
        
        if self.example_scripts:
            print(f"Successfully loaded {len(self.example_scripts)} example scripts")
        else:
            print("No example scripts were loaded")


def convert_to_16_bit_wav(data):
    if torch.is_tensor(data):
        data = data.detach().cpu().numpy()
    data = np.array(data)
    if np.max(np.abs(data)) > 1.0:
        data = data / np.max(np.abs(data))
    return (data * 32767).astype(np.int16)


def create_demo_interface(demo_instance: VibeVoiceDemo):
    custom_css = """ /* Modern light theme with gradients */
                    .gradio-container {
                        background: linear-gradient(135deg, #f8fafc 0%, #e2e8f0 100%);
                        font-family: 'SF Pro Display', -apple-system, BlinkMacSystemFont, sans-serif;
                    }
                    
                    /* Header styling */
                    .main-header {
                        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
                        padding: 2rem;
                        border-radius: 20px;
                        margin-bottom: 2rem;
                        text-align: center;
                        box-shadow: 0 10px 40px rgba(102, 126, 234, 0.3);
                    }
                    
                    .main-header h1 {
                        color: white;
                        font-size: 2.5rem;
                        font-weight: 700;
                        margin: 0;
                        text-shadow: 0 2px 4px rgba(0,0,0,0.3);
                    }
                    
                    .main-header p {
                        color: rgba(255,255,255,0.9);
                        font-size: 1.1rem;
                        margin: 0.5rem 0 0 0;
                    }
                    
                    /* Card styling */
                    .settings-card, .generation-card {
                        background: rgba(255, 255, 255, 0.8);
                        backdrop-filter: blur(10px);
                        border: 1px solid rgba(226, 232, 240, 0.8);
                        border-radius: 16px;
                        padding: 1.5rem;
                        margin-bottom: 1rem;
                        box-shadow: 0 8px 32px rgba(0, 0, 0, 0.1);
                    }
                    
                    /* Speaker selection styling */
                    .speaker-grid {
                        display: grid;
                        gap: 1rem;
                        margin-bottom: 1rem;
                    }
                    
                    .speaker-item {
                        background: linear-gradient(135deg, #e2e8f0 0%, #cbd5e1 100%);
                        border: 1px solid rgba(148, 163, 184, 0.4);
                        border-radius: 12px;
                        padding: 1rem;
                        color: #374151;
                        font-weight: 500;
                    }
                    
                    /* Streaming indicator */
                    .streaming-indicator {
                        display: inline-block;
                        width: 10px;
                        height: 10px;
                        background: #22c55e;
                        border-radius: 50%;
                        margin-right: 8px;
                        animation: pulse 1.5s infinite;
                    }
                    
                    @keyframes pulse {
                        0% { opacity: 1; transform: scale(1); }
                        50% { opacity: 0.5; transform: scale(1.1); }
                        100% { opacity: 1; transform: scale(1); }
                    }
                    
                    /* Queue status styling */
                    .queue-status {
                        background: linear-gradient(135deg, #f0f9ff 0%, #e0f2fe 100%);
                        border: 1px solid rgba(14, 165, 233, 0.3);
                        border-radius: 8px;
                        padding: 0.75rem;
                        margin: 0.5rem 0;
                        text-align: center;
                        font-size: 0.9rem;
                        color: #0369a1;
                    }
                    
                    .generate-btn {
                        background: linear-gradient(135deg, #059669 0%, #0d9488 100%);
                        border: none;
                        border-radius: 12px;
                        padding: 1rem 2rem;
                        color: white;
                        font-weight: 600;
                        font-size: 1.1rem;
                        box-shadow: 0 4px 20px rgba(5, 150, 105, 0.4);
                        transition: all 0.3s ease;
                    }
                    
                    .generate-btn:hover {
                        transform: translateY(-2px);
                        box-shadow: 0 6px 25px rgba(5, 150, 105, 0.6);
                    }
                    
                    .stop-btn {
                        background: linear-gradient(135deg, #ef4444 0%, #dc2626 100%);
                        border: none;
                        border-radius: 12px;
                        padding: 1rem 2rem;
                        color: white;
                        font-weight: 600;
                        font-size: 1.1rem;
                        box-shadow: 0 4px 20px rgba(239, 68, 68, 0.4);
                        transition: all 0.3s ease;
                    }
                    
                    .stop-btn:hover {
                        transform: translateY(-2px);
                        box-shadow: 0 6px 25px rgba(239, 68, 68, 0.6);
                    }
                    
                    /* Audio player styling */
                    .audio-output {
                        background: linear-gradient(135deg, #f1f5f9 0%, #e2e8f0 100%);
                        border-radius: 16px;
                        padding: 1.5rem;
                        border: 1px solid rgba(148, 163, 184, 0.3);
                    }
                    
                    .complete-audio-section {
                        margin-top: 1rem;
                        padding: 1rem;
                        background: linear-gradient(135deg, #f0fdf4 0%, #dcfce7 100%);
                        border: 1px solid rgba(34, 197, 94, 0.3);
                        border-radius: 12px;
                    }
                    
                    /* Text areas */
                    .script-input, .log-output {
                        background: rgba(255, 255, 255, 0.9) !important;
                        border: 1px solid rgba(148, 163, 184, 0.4) !important;
                        border-radius: 12px !important;
                        color: #1e293b !important;
                        font-family: 'JetBrains Mono', monospace !important;
                    }
                    
                    .script-input::placeholder {
                        color: #64748b !important;
                    }
                    
                    /* Sliders */
                    .slider-container {
                        background: rgba(248, 250, 252, 0.8);
                        border: 1px solid rgba(226, 232, 240, 0.6);
                        border-radius: 8px;
                        padding: 1rem;
                        margin: 0.5rem 0;
                    }
                    
                    /* Labels and text */
                    .gradio-container label {
                        color: #374151 !important;
                        font-weight: 600 !important;
                    }
                    
                    .gradio-container .markdown {
                        color: #1f2937 !important;
                    }
                    
                    /* Responsive design */
                    @media (max-width: 768px) {
                        .main-header h1 { font-size: 2rem; }
                        .settings-card, .generation-card { padding: 1rem; }
                    }
                    
                    /* Random example button styling - more subtle professional color */
                    .random-btn {
                        background: linear-gradient(135deg, #64748b 0%, #475569 100%);
                        border: none;
                        border-radius: 12px;
                        padding: 1rem 1.5rem;
                        color: white;
                        font-weight: 600;
                        font-size: 1rem;
                        box-shadow: 0 4px 20px rgba(100, 116, 139, 0.3);
                        transition: all 0.3s ease;
                        display: inline-flex;
                        align-items: center;
                        gap: 0.5rem;
                    }
                    
                    .random-btn:hover {
                        transform: translateY(-2px);
                        box-shadow: 0 6px 25px rgba(100, 116, 139, 0.4);
                        background: linear-gradient(135deg, #475569 0%, #334155 100%);
                    }
                    """

    with gr.Blocks(
        title="VibeVoice - AI Podcast Generator",
        css=custom_css,
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="purple",
            neutral_hue="slate",
        )
    ) as interface:

        gr.HTML("""
        <div class="main-header">
            <h1>🎙️ Vibe Podcasting</h1>
            <p>Generating Long-form Multi-speaker AI Podcast with VibeVoice</p>
        </div>
        """)

        with gr.Row():
            with gr.Column(scale=1, elem_classes="settings-card"):
                gr.Markdown("### 🎛️ Podcast Settings")

                # NEW - model dropdown
                model_dropdown = gr.Dropdown(
                    choices=list(demo_instance.models.keys()),
                    value=demo_instance.current_model_name,
                    label="Model",
                )

                num_speakers = gr.Slider(
                    minimum=1, maximum=4, value=2, step=1,
                    label="Number of Speakers",
                    elem_classes="slider-container"
                )

                gr.Markdown("### 🎭 Speaker Selection")
                available_speaker_names = list(demo_instance.available_voices.keys())
                default_speakers = ['en-Alice_woman', 'en-Carter_man', 'en-Frank_man', 'en-Maya_woman']

                speaker_selections = []
                for i in range(4):
                    default_value = default_speakers[i] if i < len(default_speakers) else None
                    speaker = gr.Dropdown(
                        choices=available_speaker_names,
                        value=default_value,
                        label=f"Speaker {i+1}",
                        visible=(i < 2),
                        elem_classes="speaker-item"
                    )
                    speaker_selections.append(speaker)

                gr.Markdown("### ⚙️ Advanced Settings")
                with gr.Accordion("Generation Parameters", open=False):
                    cfg_scale = gr.Slider(
                        minimum=1.0, maximum=2.0, value=1.3, step=0.05,
                        label="CFG Scale (Guidance Strength)",
                        elem_classes="slider-container"
                    )

            with gr.Column(scale=2, elem_classes="generation-card"):
                gr.Markdown("### 📝 Script Input")
                script_input = gr.Textbox(
                    label="Conversation Script",
                    placeholder="Enter your podcast script here...",
                    lines=12,
                    max_lines=20,
                    elem_classes="script-input"
                )

                with gr.Row():
                    random_example_btn = gr.Button(
                        "🎲 Random Example", size="lg",
                        variant="secondary", elem_classes="random-btn", scale=1
                    )
                    generate_btn = gr.Button(
                        "🚀 Generate Podcast", size="lg",
                        variant="primary", elem_classes="generate-btn", scale=2
                    )

                gr.Markdown("### 🎵 Generated Podcast")
                complete_audio_output = gr.Audio(
                    label="Complete Podcast (Download)",
                    type="numpy",
                    elem_classes="audio-output complete-audio-section",
                    autoplay=False,
                    show_download_button=True,
                    visible=True
                )

                log_output = gr.Textbox(
                    label="Generation Log",
                    lines=8, max_lines=15,
                    interactive=False,
                    elem_classes="log-output"
                )

        def update_speaker_visibility(num_speakers):
            return [gr.update(visible=(i < num_speakers)) for i in range(4)]

        num_speakers.change(
            fn=update_speaker_visibility,
            inputs=[num_speakers],
            outputs=speaker_selections
        )

        def generate_podcast_wrapper(model_choice, num_speakers, script, *speakers_and_params):
            try:
                speakers = speakers_and_params[:4]
                cfg_scale_val = speakers_and_params[4]
                audio, log = demo_instance.generate_podcast(
                    num_speakers=int(num_speakers),
                    script=script,
                    speaker_1=speakers[0],
                    speaker_2=speakers[1],
                    speaker_3=speakers[2],
                    speaker_4=speakers[3],
                    cfg_scale=cfg_scale_val,
                    model_name=model_choice
                )
                return audio, log
            except Exception as e:
                traceback.print_exc()
                return None, f"❌ Error: {str(e)}"

        generate_btn.click(
            fn=generate_podcast_wrapper,
            inputs=[model_dropdown, num_speakers, script_input] + speaker_selections + [cfg_scale],
            outputs=[complete_audio_output, log_output],
            queue=True
        )

        def load_random_example():
            import random
            examples = getattr(demo_instance, "example_scripts", [])
            if not examples:
                examples = [
                    [2, "Speaker 0: Welcome to our AI podcast demo!\nSpeaker 1: Thanks, excited to be here!"]
                ]
            num_speakers_value, script_value = random.choice(examples)
            return num_speakers_value, script_value

        random_example_btn.click(
            fn=load_random_example,
            inputs=[],
            outputs=[num_speakers, script_input],
            queue=False
        )

        gr.Markdown("### 📚 Example Scripts")
        examples = getattr(demo_instance, "example_scripts", []) or [
            [1, "Speaker 1: Welcome to our AI podcast demo. This is a sample script."]
        ]
        gr.Examples(
            examples=examples,
            inputs=[num_speakers, script_input],
            label="Try these example scripts:"
        )

    return interface




def run_demo(
    model_paths: dict = None,
    device: str = "cuda",
    inference_steps: int = 5,
    share: bool = True,
):
    """
    model_paths default includes two entries. Replace paths as needed.
    """
    if model_paths is None:
        model_paths = {
            "VibeVoice-Large": "microsoft/VibeVoice-Large",
            "VibeVoice-1.1B": "microsoft/VibeVoice-1.1B"
        }

    set_seed(42)
    demo_instance = VibeVoiceDemo(model_paths, device, inference_steps)
    interface = create_demo_interface(demo_instance)
    interface.queue().launch(
        share=share,
        server_name="0.0.0.0" if share else "127.0.0.1",
        show_error=True,
        show_api=False
    )



if __name__ == "__main__":
    run_demo()