import gradio as gr
import io
import wave
import numpy as np

# Lazy imports for optional dependencies
try:
    import torch  # type: ignore
except Exception:  # pragma: no cover
    torch = None  # type: ignore

try:
    from pocket_tts import TTSModel  # type: ignore
except Exception:  # pragma: no cover
    TTSModel = None  # type: ignore

# Global state for lazy initialization
_POCKET_STATE = {
    "initialized": False,
    "model": None,
    "voice_states": {},
    "sample_rate": 24000,
}

def _get_available_voices() -> dict[str, str]:
    """Get available voices from the local ./voices/ directory.
    
    Scans ./voices/ directory for audio files (WAV, MP3, etc.)
    """
    import os
    
    voices_dir = os.path.join(os.path.dirname(__file__), "voices")
    local_voices = {}
    
    if os.path.exists(voices_dir):
        for f in os.listdir(voices_dir):
            # Support common audio formats
            if f.lower().endswith(('.wav', '.mp3', '.flac', '.ogg', '.m4a')):
                voice_name = os.path.splitext(f)[0]
                local_voices[voice_name] = os.path.join(voices_dir, f)
    
    if local_voices:
        print(f"Found {len(local_voices)} local voice(s): {list(local_voices.keys())}")
    else:
        print("WARNING: No voices found in voices/ directory")
    
    return local_voices


# Scan voices at import time
PRESET_VOICES = _get_available_voices()


def _init_pocket(
    temp: float = 0.7,
    lsd_decode_steps: int = 1,
    noise_clamp: float | None = None,
    eos_threshold: float = -4.0,
) -> None:
    """Lazy initialization of the Pocket TTS model."""
    if _POCKET_STATE["initialized"]:
        return

    if TTSModel is None:
        raise gr.Error(
            "pocket-tts is not installed. Please install with: pip install pocket-tts"
        )

    if torch is None:
        raise gr.Error("PyTorch is not installed. Please install torch>=2.5.0")

    print("Initializing Pocket TTS...")

    # Log in to HuggingFace if token is available (enables voice cloning on Spaces)
    import os
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        print("HF_TOKEN found, using for authentication")


    # Auto-detect device: CPU by default, CUDA if available
    # Note: The pocket-tts docs mention GPU doesn't provide speedup for this model
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {device}")

    try:
        model = TTSModel.load_model(
            temp=float(temp),
            lsd_decode_steps=int(lsd_decode_steps),
            noise_clamp=float(noise_clamp) if noise_clamp is not None else None,
            eos_threshold=float(eos_threshold),
        )
        _POCKET_STATE.update({
            "initialized": True,
            "model": model,
            "sample_rate": model.sample_rate,
        })
        print(f"Pocket TTS initialized. Sample rate: {model.sample_rate} Hz")
        
        # Auto-create missing embeddings if voice cloning is available
        if model.has_voice_cloning:
            _create_missing_embeddings(model)
        else:
            print("Voice cloning not available - using pre-computed embeddings only")
            
    except Exception as e:
        raise gr.Error(f"Failed to initialize Pocket TTS model: {str(e)}")


def _create_missing_embeddings(model) -> None:
    """Create embeddings for any voices that have audio files but no embedding."""
    import os
    from pocket_tts.data.audio import audio_read
    from pocket_tts.data.audio_utils import convert_audio
    import safetensors.torch
    
    voices_dir = os.path.join(os.path.dirname(__file__), "voices")
    embeddings_dir = os.path.join(os.path.dirname(__file__), "embeddings")
    
    if not os.path.exists(voices_dir):
        return
    
    os.makedirs(embeddings_dir, exist_ok=True)
    
    audio_extensions = ('.wav', '.mp3', '.flac', '.ogg', '.m4a')
    
    for voice_name, voice_path in PRESET_VOICES.items():
        embedding_path = os.path.join(embeddings_dir, f"{voice_name}.safetensors")
        
        # Skip if embedding already exists or no local file
        if os.path.exists(embedding_path) or voice_path is None:
            continue
        
        # Skip fallback HuggingFace voices
        if voice_path.startswith("hf://"):
            continue
            
        print(f"Creating embedding for '{voice_name}'...")
        
        try:
            # Convert to WAV if needed
            audio_path = voice_path
            if not voice_path.lower().endswith('.wav'):
                from pydub import AudioSegment
                import tempfile
                audio = AudioSegment.from_file(voice_path)
                temp_wav = tempfile.NamedTemporaryFile(suffix='.wav', delete=False)
                audio.export(temp_wav.name, format='wav')
                audio_path = temp_wav.name
            
            # Read and encode audio
            audio, sr = audio_read(audio_path)
            audio_tensor = convert_audio(audio, sr, model.config.mimi.sample_rate, 1)
            
            with torch.no_grad():
                audio_prompt = model._encode_audio(audio_tensor.unsqueeze(0).to(model.device))
            
            # Save embedding
            safetensors.torch.save_file(
                {"audio_prompt": audio_prompt.cpu()},
                embedding_path
            )
            print(f"  Saved: {embedding_path}")
            
        except Exception as e:
            print(f"  Error creating embedding for {voice_name}: {e}")


def _convert_to_wav(audio_path: str) -> str:
    """Convert audio file to WAV format if needed.
    
    Returns the path to a WAV file (original if already WAV, or converted temp file).
    Uses pydub for MP3 (requires ffmpeg), soundfile for other formats.
    """
    import tempfile
    
    # Check if already WAV
    if audio_path.lower().endswith('.wav'):
        return audio_path
    
    print(f"Converting {audio_path} to WAV format...")
    
    # Create temp file path
    import os
    tmp_fd, wav_path = tempfile.mkstemp(suffix=".wav")
    os.close(tmp_fd)
    
    # Try pydub first (better MP3 support via ffmpeg)
    try:
        from pydub import AudioSegment
        audio = AudioSegment.from_file(audio_path)
        audio.export(wav_path, format="wav")
        print(f"Converted via pydub to: {wav_path}")
        return wav_path
    except ImportError:
        pass  # pydub not installed, try soundfile
    except Exception as e:
        print(f"pydub conversion failed: {e}, trying soundfile...")
    
    # Fall back to soundfile
    try:
        import soundfile as sf
        audio_data, sample_rate = sf.read(audio_path)
        sf.write(wav_path, audio_data, sample_rate)
        print(f"Converted via soundfile to: {wav_path}")
        return wav_path
    except Exception as e:
        raise gr.Error(f"Failed to convert audio file: {str(e)}. Please upload a WAV file directly or install pydub+ffmpeg for MP3 support.")


def _get_voice_state(voice_name: str | None, custom_audio_path: str | None):
    """Get or create voice state for generation.
    
    Args:
        voice_name: Name of preset voice (alba, marius, etc.)
        custom_audio_path: Path to custom audio file for voice cloning
    
    Returns:
        Voice state dict for the model
    """
    model = _POCKET_STATE["model"]

    # Custom audio takes priority
    if custom_audio_path:
        print(f"Loading custom voice from: {custom_audio_path}")
        # Convert to WAV if needed
        wav_path = _convert_to_wav(custom_audio_path)
        return model.get_state_for_audio_prompt(wav_path)

    # Use preset voice
    if not voice_name or voice_name not in PRESET_VOICES:
        # Default to first available voice
        voice_name = list(PRESET_VOICES.keys())[0] if PRESET_VOICES else None
        if not voice_name:
            raise gr.Error("No voices available. Add audio files to the voices/ directory.")

    # Check cache
    if voice_name in _POCKET_STATE["voice_states"]:
        return _POCKET_STATE["voice_states"][voice_name]

    # Check for pre-computed embedding first (no voice cloning needed)
    import os
    embeddings_dir = os.path.join(os.path.dirname(__file__), "embeddings")
    embedding_path = os.path.join(embeddings_dir, f"{voice_name}.safetensors")
    
    if os.path.exists(embedding_path):
        print(f"Loading pre-computed embedding for '{voice_name}' from: {embedding_path}")
        import safetensors.torch
        from pocket_tts.modules.stateful_module import init_states
        
        # Load the audio prompt embedding
        state_dict = safetensors.torch.load_file(embedding_path)
        audio_prompt = state_dict["audio_prompt"].to(model.device)
        
        # Create fresh model state and condition it with the audio prompt
        # (same logic as model.get_state_for_audio_prompt uses internally)
        voice_state = init_states(model.flow_lm, batch_size=1, sequence_length=1000)
        model._run_flow_lm_and_increment_step(model_state=voice_state, audio_conditioning=audio_prompt)
        
        # Detach all tensors to make them leaf tensors (required for deepcopy)
        def detach_tensors(obj):
            if isinstance(obj, torch.Tensor):
                return obj.detach().clone()
            elif isinstance(obj, dict):
                return {k: detach_tensors(v) for k, v in obj.items()}
            else:
                return obj
        
        voice_state = detach_tensors(voice_state)
        
        _POCKET_STATE["voice_states"][voice_name] = voice_state
        return voice_state

    # Fall back to voice cloning (requires auth)
    if not model.has_voice_cloning:
        raise gr.Error(
            f"No embedding found for voice '{voice_name}'. "
            f"Voice cloning is not available (requires HF auth). "
            f"Run the app locally first to create embeddings."
        )
    
    voice_path = PRESET_VOICES[voice_name]
    print(f"Loading preset voice '{voice_name}' from: {voice_path}")
    
    # Convert to WAV if needed (local files may be MP3, etc.)
    wav_path = _convert_to_wav(voice_path)
    voice_state = model.get_state_for_audio_prompt(wav_path)
    _POCKET_STATE["voice_states"][voice_name] = voice_state
    return voice_state


def _audio_np_to_int16(audio_np: np.ndarray) -> np.ndarray:
    """Convert float audio array to int16."""
    audio_clipped = np.clip(audio_np, -1.0, 1.0)
    return (audio_clipped * 32767.0).astype(np.int16)


def _wav_bytes_from_int16(audio_int16: np.ndarray, sample_rate: int) -> bytes:
    """Create WAV bytes from int16 audio array."""
    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        wf.writeframes(audio_int16.tobytes())
    return buffer.getvalue()


def _split_into_sentences(text: str) -> list[str]:
    """Split text into sentences for chunk-by-chunk generation.
    
    Uses simple punctuation-based splitting for natural speech chunks.
    """
    import re
    # Split on sentence-ending punctuation, keeping the punctuation
    # Handle common patterns: . ! ? and combinations like "..." or "?!"
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    # Filter out empty strings and strip whitespace
    return [s.strip() for s in sentences if s.strip()]


def pocket_tts_stream(
    text: str,
    voice: str,
    custom_audio,
    temperature: float,
    lsd_decode_steps: int,
    noise_clamp: float | None,
    eos_threshold: float,
    frames_after_eos: int,
):
    """Generate speech with sentence-level streaming.
    
    Splits text into sentences and yields complete audio for each sentence,
    matching Kokoro's smooth streaming pattern.
    """
    if not text or not text.strip():
        raise gr.Error("Please enter text to synthesize.")

    # Initialize model with current parameters
    _init_pocket(
        temp=temperature,
        lsd_decode_steps=lsd_decode_steps,
        noise_clamp=noise_clamp if noise_clamp and noise_clamp > 0 else None,
        eos_threshold=eos_threshold,
    )

    model = _POCKET_STATE["model"]
    sample_rate = _POCKET_STATE["sample_rate"]

    # Get voice state
    custom_path = custom_audio if custom_audio else None
    voice_state = _get_voice_state(voice, custom_path)

    # Split text into sentences for natural chunking
    sentences = _split_into_sentences(text)
    if not sentences:
        raise gr.Error("No valid sentences found in text.")

    produced_any = False
    
    # Buffer for initial audio - wait for ~5 seconds before yielding first chunk
    # This prevents stuttering from short first sentences
    min_initial_samples = int(sample_rate * 5)  # 5 seconds of audio
    audio_buffer = []
    buffer_samples = 0
    initial_buffer_yielded = False

    try:
        for idx, sentence in enumerate(sentences):
            # Generate complete audio for this sentence (non-streaming per sentence)
            audio = model.generate_audio(
                voice_state,
                sentence,
                frames_after_eos=frames_after_eos if frames_after_eos > 0 else None,
                copy_state=True,
            )
            produced_any = True
            
            # Convert tensor to numpy
            audio_np = audio.cpu().numpy() if hasattr(audio, 'cpu') else audio
            
            if not initial_buffer_yielded:
                # Accumulate in buffer until we have enough audio
                audio_buffer.append(audio_np)
                buffer_samples += len(audio_np)
                
                # Check if we have enough or this is the last sentence
                if buffer_samples >= min_initial_samples or idx == len(sentences) - 1:
                    # Yield the accumulated buffer
                    combined = np.concatenate(audio_buffer, axis=0)
                    audio_int16 = _audio_np_to_int16(combined)
                    yield _wav_bytes_from_int16(audio_int16, sample_rate)
                    audio_buffer = []
                    buffer_samples = 0
                    initial_buffer_yielded = True
            else:
                # After initial buffer, yield each sentence immediately
                audio_int16 = _audio_np_to_int16(audio_np)
                yield _wav_bytes_from_int16(audio_int16, sample_rate)

    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Error during speech generation: {str(e)[:200]}...")

    if not produced_any:
        raise gr.Error("No audio was generated.")


def generate_tts(
    text: str,
    voice: str,
    custom_audio,
    temperature: float,
    lsd_decode_steps: int,
    noise_clamp: float,
    eos_threshold: float,
    frames_after_eos: int,
):
    """Main streaming dispatcher for Pocket TTS."""
    yield from pocket_tts_stream(
        text,
        voice,
        custom_audio,
        temperature,
        lsd_decode_steps,
        noise_clamp,
        eos_threshold,
        frames_after_eos,
    )


# --- Gradio UI ---
with gr.Blocks() as demo:
    gr.HTML(
        "<h1 style='text-align: center;'>Pocket-TTS</h1>"
    )
    device_info = gr.Markdown(
        "<p style='text-align: center;'>Powered by kyutai/pocket-tts | Running on CPU | Voices cloned from Kokoro-82M</p>"
    )
    
    def update_device_info():
        device = "CUDA" if torch.cuda.is_available() else "CPU"
        return f"<p style='text-align: center;'>Powered by kyutai/pocket-tts | Running on {device} | Voices cloned from Kokoro-82M</p>"
    
    demo.load(update_device_info, outputs=device_info)

    with gr.Row():
        with gr.Column():
            # Text input
            text_input = gr.Textbox(
                label="Input Text",
                placeholder="Enter the text you want to convert to speech here...",
                lines=5,
                value="The quick brown fox jumps over the lazy dog. I am already far north of London, and as I walk in the streets of Petersburgh, I feel a cold northern breeze play upon my cheeks, which braces my nerves and fills me with delight. Do you understand this feeling? This breeze, which has traveled from the regions towards which I am advancing, gives me a foretaste of those icy climes. Inspirited by this wind of promise, my daydreams become more fervent and vivid.",
            )

            # Voice selection
            with gr.Group():
                gr.Markdown("### Voice Selection")
                gr.Markdown("Select a preset voice OR upload your own WAV file for voice cloning.")

                voice_dropdown = gr.Dropdown(
                    choices=list(PRESET_VOICES.keys()),
                    label="Preset Voice",
                    value=list(PRESET_VOICES.keys())[0] if PRESET_VOICES else None,
                    info="Select a pre-loaded voice. Ignored if custom audio is uploaded.",
                )

                gr.Markdown("--- OR ---")

                ref_audio_input = gr.Audio(
                    label="Custom Voice (WAV)",
                    type="filepath",
                    sources=["upload", "microphone"],
                )

            generate_btn = gr.Button(
                "Generate Speech",
                variant="primary",
            )

        with gr.Column():
            audio_output = gr.Audio(
                label="Generated Speech",
                streaming=True,
                autoplay=True,
            )

            with gr.Accordion("Advanced Options", open=False):
                temp_slider = gr.Slider(
                    minimum=0.1,
                    maximum=1.5,
                    value=0.7,
                    step=0.05,
                    label="Temperature",
                    info="Controls randomness. Higher = more varied, lower = more consistent.",
                )
                lsd_steps_slider = gr.Slider(
                    minimum=1,
                    maximum=10,
                    value=1,
                    step=1,
                    label="LSD Decode Steps",
                    info="Number of generation steps. Higher = potentially better quality but slower.",
                )
                noise_clamp_slider = gr.Slider(
                    minimum=0.0,
                    maximum=5.0,
                    value=0.0,
                    step=0.1,
                    label="Noise Clamp",
                    info="Maximum value for noise sampling. 0 = disabled.",
                )
                eos_threshold_slider = gr.Slider(
                    minimum=-10.0,
                    maximum=0.0,
                    value=-4.0,
                    step=0.5,
                    label="EOS Threshold",
                    info="Threshold for end-of-sequence detection. More negative = longer audio.",
                )
                frames_after_eos_slider = gr.Slider(
                    minimum=0,
                    maximum=10,
                    value=2,
                    step=1,
                    label="Frames After EOS",
                    info="Additional frames to generate after EOS detection.",
                )

    # Connect inputs
    generate_inputs = [
        text_input,
        voice_dropdown,
        ref_audio_input,
        temp_slider,
        lsd_steps_slider,
        noise_clamp_slider,
        eos_threshold_slider,
        frames_after_eos_slider,
    ]

    generate_btn.click(
        fn=generate_tts,
        inputs=generate_inputs,
        outputs=audio_output,
        api_name="generate_speech",
    )

    text_input.submit(
        fn=generate_tts,
        inputs=generate_inputs,
        outputs=audio_output,
        api_name="generate_speech_enter",
    )


if __name__ == "__main__":
    demo.queue().launch(debug=True, theme="Nymbo/Nymbo_Theme")