import gradio as gr
import spaces
import torch
import numpy as np
import os
import tarfile
from pathlib import Path
from typing import Optional
from huggingface_hub import hf_hub_download
import sentencepiece

# PersonaPlex model imports - installed via: pip install git+https://github.com/NVIDIA/personaplex.git#subdirectory=moshi
from moshi.models import loaders, LMGen
from moshi.models.lm import load_audio, _iterate_audio, encode_from_sphn

# Configuration
HF_REPO = "nvidia/personaplex-7b-v1"
DEVICE = "cuda"
SAMPLE_RATE = 24000

# Available voices in PersonaPlex
ALL_VOICES = [
    "NATF0", "NATF1", "NATF2", "NATF3",  # Natural Female
    "NATM0", "NATM1", "NATM2", "NATM3",  # Natural Male
    "VARF0", "VARF1", "VARF2", "VARF3", "VARF4",  # Variety Female
    "VARM0", "VARM1", "VARM2", "VARM3", "VARM4",  # Variety Male
]

# Example persona prompts from PersonaPlex paper
EXAMPLE_PERSONAS = [
    "You are a wise and friendly teacher. Answer questions or provide advice in a clear and engaging way.",
    "You enjoy having a good conversation.",
    "You work for CitySan Services which is a waste management company and your name is Ayelen Lucero.",
    "You enjoy having a good conversation. Have a technical discussion about fixing a reactor core on a spaceship to Mars. You are an astronaut on a Mars mission. Your name is Alex.",
]

# Pre-download model weights at startup (cached by huggingface_hub)
print("Downloading model weights...")
MIMI_WEIGHT = hf_hub_download(HF_REPO, loaders.MIMI_NAME)
MOSHI_WEIGHT = hf_hub_download(HF_REPO, loaders.MOSHI_NAME)
TOKENIZER_PATH = hf_hub_download(HF_REPO, loaders.TEXT_TOKENIZER_NAME)
VOICES_TGZ = hf_hub_download(HF_REPO, "voices.tgz")

# Extract voices archive
VOICES_DIR = Path(VOICES_TGZ).parent / "voices"
if not VOICES_DIR.exists():
    print("Extracting voice embeddings...")
    with tarfile.open(VOICES_TGZ, "r:gz") as tar:
        tar.extractall(path=Path(VOICES_TGZ).parent)
print("Model weights ready.")

# Global model cache
_model_cache = {}


def get_models():
    """Lazy load models on first GPU call."""
    if "initialized" not in _model_cache:
        print("Loading models to GPU...")
        
        # Load Mimi encoder/decoder
        mimi = loaders.get_mimi(MIMI_WEIGHT, DEVICE)
        other_mimi = loaders.get_mimi(MIMI_WEIGHT, DEVICE)
        
        # Load Moshi LM
        lm = loaders.get_moshi_lm(MOSHI_WEIGHT, device=DEVICE)
        lm.eval()
        
        # Load text tokenizer
        text_tokenizer = sentencepiece.SentencePieceProcessor(TOKENIZER_PATH)
        
        # Create LMGen wrapper
        frame_size = int(mimi.sample_rate / mimi.frame_rate)
        lm_gen = LMGen(
            lm,
            audio_silence_frame_cnt=int(0.5 * mimi.frame_rate),
            sample_rate=mimi.sample_rate,
            device=DEVICE,
            frame_rate=mimi.frame_rate,
            temp=0.8,
            temp_text=0.7,
            top_k=250,
            top_k_text=25,
        )
        
        # Enable streaming mode
        mimi.streaming_forever(1)
        other_mimi.streaming_forever(1)
        lm_gen.streaming_forever(1)
        
        _model_cache.update({
            "mimi": mimi,
            "other_mimi": other_mimi,
            "lm_gen": lm_gen,
            "tokenizer": text_tokenizer,
            "frame_size": frame_size,
            "initialized": True,
        })
        print("Models loaded successfully.")
    
    return _model_cache


def wrap_with_system_tags(text: str) -> str:
    """Add system tags as PersonaPlex expects."""
    text = text.strip()
    if text.startswith("<system>") and text.endswith("<system>"):
        return text
    return f"<system> {text} <system>"


def decode_tokens_to_pcm(mimi, other_mimi, tokens: torch.Tensor) -> np.ndarray:
    """Decode audio tokens to PCM waveform."""
    # tokens shape: [B, num_codebooks, 1]
    # Agent audio is in codebooks 1:9
    agent_audio_tokens = tokens[:, 1:9, :]
    pcm = other_mimi.decode(agent_audio_tokens)
    return pcm[0, 0].cpu().numpy()


@spaces.GPU(duration=120)
def generate_response(audio_input, persona: str, voice: str):
    """Process audio input and generate PersonaPlex response."""
    if audio_input is None:
        return None, "Please record audio first."
    
    models = get_models()
    mimi = models["mimi"]
    other_mimi = models["other_mimi"]
    lm_gen = models["lm_gen"]
    tokenizer = models["tokenizer"]
    frame_size = models["frame_size"]
    
    # Process input audio
    sr, audio = audio_input
    audio = audio.astype(np.float32)
    
    # Convert to mono if stereo
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    
    # Normalize to [-1, 1]
    if audio.max() > 1.0 or audio.min() < -1.0:
        audio = audio / 32768.0 if audio.dtype == np.int16 else audio / np.abs(audio).max()
    
    # Resample to model's sample rate if needed
    if sr != mimi.sample_rate:
        import sphn
        audio = sphn.resample(audio, sr, mimi.sample_rate)
    
    # Add channel dimension: (T,) -> (1, T)
    if audio.ndim == 1:
        audio = audio[None, :]
    
    # Load voice prompt
    voice_path = str(VOICES_DIR / f"{voice}.pt")
    if not os.path.exists(voice_path):
        return None, f"Voice '{voice}' not found."
    lm_gen.load_voice_prompt_embeddings(voice_path)
    
    # Set text prompt
    if persona.strip():
        lm_gen.text_prompt_tokens = tokenizer.encode(wrap_with_system_tags(persona))
    else:
        lm_gen.text_prompt_tokens = None
    
    # Reset streaming state
    mimi.reset_streaming()
    other_mimi.reset_streaming()
    lm_gen.reset_streaming()
    
    # Run system prompts (voice + text conditioning)
    with lm_gen.streaming(1):
        lm_gen.step_system_prompts(mimi)
        mimi.reset_streaming()
        
        # Process user audio frames
        generated_frames = []
        generated_text = []
        
        for user_encoded in encode_from_sphn(
            mimi,
            _iterate_audio(audio, sample_interval_size=frame_size, pad=True),
            max_batch=1,
        ):
            for c in range(user_encoded.shape[-1]):
                step_in = user_encoded[:, :, c:c+1]
                tokens = lm_gen.step(step_in)
                
                if tokens is None:
                    continue
                
                # Decode agent audio
                pcm = decode_tokens_to_pcm(mimi, other_mimi, tokens)
                generated_frames.append(pcm)
                
                # Decode text token
                text_token = tokens[0, 0, 0].item()
                if text_token not in (0, 3):  # Skip special tokens
                    text_piece = tokenizer.id_to_piece(text_token).replace("▁", " ")
                    generated_text.append(text_piece)
    
    if not generated_frames:
        return None, "No audio generated. Try speaking more clearly."
    
    # Concatenate output audio
    output_audio = np.concatenate(generated_frames, axis=-1)
    output_text = "".join(generated_text).strip()
    
    return (mimi.sample_rate, output_audio), output_text


# Build Gradio interface
with gr.Blocks(title="PersonaPlex Demo", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎭 PersonaPlex
        **Voice and Role Control for Full Duplex Conversational Speech Models**
        
        [Paper](https://arxiv.org/abs/2503.04721) | [GitHub](https://github.com/NVIDIA/personaplex) | [Model](https://huggingface.co/nvidia/personaplex-7b-v1)
        
        ---
        
        Record your message, and PersonaPlex will respond with the configured persona and voice.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            persona = gr.Textbox(
                label="Persona Description",
                placeholder="Describe the assistant's persona...",
                value=EXAMPLE_PERSONAS[0],
                lines=4,
            )
            voice = gr.Dropdown(
                choices=ALL_VOICES,
                value="NATF2",
                label="Voice"
            )
            gr.Examples(
                examples=[[p] for p in EXAMPLE_PERSONAS],
                inputs=[persona],
                label="Example Personas"
            )
        
        with gr.Column(scale=2):
            audio_input = gr.Audio(
                label="🎤 Record your message",
                sources=["microphone"],
                type="numpy",
            )
            generate_btn = gr.Button("Generate Response", variant="primary", size="lg")
            
            audio_output = gr.Audio(
                label="🔊 PersonaPlex Response",
                type="numpy",
                autoplay=True,
            )
            text_output = gr.Textbox(
                label="📝 Response Text",
                interactive=False,
            )
    
    generate_btn.click(
        fn=generate_response,
        inputs=[audio_input, persona, voice],
        outputs=[audio_output, text_output],
    )


if __name__ == "__main__":
    demo.launch()