"""Audio Interaction Model — Gradio demo.

Takes an audio clip (microphone or file) and an optional text instruction,
streams a text response from the AudioInteraction model.
"""

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / any CUDA-touching import
import torch

# Monkey-patch torch.load to allow loading the audio encoder checkpoint
# (which pickles numpy/object globals and fails under torch 2.6+ weights_only=True default)
_orig_torch_load = torch.load
def _patched_torch_load(*args, **kwargs):
    return _orig_torch_load(*args, **{**kwargs, "weights_only": kwargs.get("weights_only", False)})
torch.load = _patched_torch_load

import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import random
import numpy as np
import gradio as gr
import whisper
import tempfile
import time
import json
from pathlib import Path
from transformers import AutoConfig
from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
from huggingface_hub import snapshot_download
from safetensors.torch import load_file

from src.audiointeraction.dataset.TOKENS import (
    ASSISTANT, AUDIO_BEGIN, ENGLISH, KEEP_SILENCE, ONLINE, PAD,
    SYSTEM, TEXT_BEGIN, TEXT_END,
    HAPPY, SAD, ANGRY, SURPRISE, NORMAL, URGENT,
)
from src.audiointeraction.generate.base import (
    AUDIO_TOKENS_PER_CHUNK, encode_audio_chunks, encode_silence_chunks, sample
)
from src.audiointeraction.model import GPT, Config
from src.audiointeraction.tokenizer import Tokenizer

# ── Constants ──────────────────────────────────────────────

MODEL_ID = "zhifeixie/AudioInteraction"
SYSTEM_PROMPT = (
    "You are a helpful assistant. When there is no user text, if the audio contains a question, "
    "please answer it. If it is a sound effect, determine based on the sound whether help is needed."
)

EMOTION_EMOJI = {
    HAPPY:    "😊",
    SAD:      "😢",
    ANGRY:    "😠",
    SURPRISE: "😲",
    NORMAL:   "😐",
    URGENT:   "⚠️",
}


def set_seed(seed: int = 1337) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


# ── Model loading (module scope, eager) ────────────────────

def _resolve_checkpoint_paths(checkpoint_dir: str):
    ckpt = Path(checkpoint_dir)
    return (
        str(ckpt),
        str(ckpt),
        str(ckpt / "qwen25OmniConfig"),
        str(ckpt / "audiointeraction_ChunkwisedEncoder.pth"),
    )


def _load_audio_encoder(qwen_omni_ckpt, audio_tower_ckpt, device):
    cfg = AutoConfig.from_pretrained(qwen_omni_ckpt)
    audio_cfg = cfg.thinker_config.audio_config
    encoder = Qwen2_5OmniAudioEncoder._from_config(audio_cfg)
    state_dict = torch.load(audio_tower_ckpt, map_location=device)
    encoder.load_state_dict(state_dict)
    encoder.to(device).requires_grad_(False).eval()
    return encoder


# Download model weights at module scope
print("Downloading model weights from HuggingFace...")
ckpt_dir = snapshot_download(
    repo_id=MODEL_ID,
    repo_type="model",
    local_dir="./checkpoints",
    resume_download=True,
)
print(f"Model downloaded to {ckpt_dir}")

model_config_dir, trained_checkpoint, qwen_omni_ckpt, audio_tower_ckpt = \
    _resolve_checkpoint_paths(ckpt_dir)

set_seed(1337)
device = "cuda"

print("Loading language model...")
config = Config.from_file(Path(model_config_dir) / "model_config.yaml")
model = GPT(config)
model.max_seq_length = config.block_size

checkpoint_dir = Path(trained_checkpoint)
index_path = checkpoint_dir / "model.safetensors.index.json"
with open(index_path) as f:
    index = json.load(f)
shard_files = sorted(set(index["weight_map"].values()))
state_dict = {}
for shard in shard_files:
    state_dict.update(load_file(str(checkpoint_dir / shard), device="cpu"))
missing, unexpected = model.load_state_dict(state_dict, strict=True)
if missing or unexpected:
    print(f"[load_model] missing={missing[:3]}… unexpected={unexpected[:3]}…")
del state_dict  # free memory

model = model.to(device).to(torch.bfloat16)
model.eval()

print("Loading audio encoder...")
audio_encoder = _load_audio_encoder(qwen_omni_ckpt, audio_tower_ckpt, device)
# Keep audio encoder in float32 — the original code does not convert it to bf16

tokenizer = Tokenizer(model_config_dir)

system_ids = tokenizer.encode(SYSTEM_PROMPT).cpu().tolist()
prefix_ids = torch.LongTensor(
    [ONLINE, ENGLISH, SYSTEM, TEXT_BEGIN] + system_ids + [TEXT_END]
).to(device)

print("Model loaded successfully!")


# ── Inference ──────────────────────────────────────────────

@spaces.GPU(duration=180)
def interact(audio_file, text_instruction=""):
    """Process an audio clip and optional text instruction, return the model's text response.

    Args:
        audio_file: Path to an audio file (wav, mp3, m4a, etc.).
        text_instruction: Optional text instruction to guide the model's response.

    Returns:
        A tuple of (full_response_text, status_text) where full_response_text
        is the model's text response and status_text provides processing info.
    """
    if audio_file is None:
        return "Please provide an audio file.", "No audio input"

    t0 = time.perf_counter()

    # Set up KV cache on the correct device
    model.set_kv_cache(batch_size=1, device=device, dtype=torch.bfloat16)
    model.max_seq_length = model.config.block_size
    # Ensure cos/sin rope buffers are on the right device
    model.cos = model.cos.to(device)
    model.sin = model.sin.to(device)

    # Build prefix
    token = prefix_ids.clone()
    input_pos = torch.arange(0, prefix_ids.size(0), device=device, dtype=torch.int64)
    input_pos_maxp1 = torch.tensor(prefix_ids.size(0), device=device)

    # If text instruction provided, append it
    if text_instruction.strip():
        text_ids = tokenizer.encode(text_instruction).cpu().tolist()
        text_tokens = torch.LongTensor(
            [TEXT_BEGIN] + text_ids + [TEXT_END]
        ).to(device)
        token = torch.cat([token, text_tokens])
        input_pos = torch.cat([input_pos, torch.arange(
            prefix_ids.size(0), prefix_ids.size(0) + len(text_tokens),
            device=device, dtype=torch.int64
        )])

    # Encode audio
    audio_chunks = encode_audio_chunks(audio_file, audio_encoder, device)
    # Follow the original inference pattern: audio then 1s silence
    # The model often responds during the silence round after audio
    silence_chunks = encode_silence_chunks(1.0, audio_encoder, device)
    all_chunks = audio_chunks + silence_chunks
    n_chunks = len(all_chunks)

    # Run streaming inference
    response_parts = []
    current_text = []
    listening = True
    audio_idx = -1
    emotion_tag = ""
    max_tokens = 4096

    try:
        with torch.inference_mode():
            for _ in range(max_tokens - input_pos.numel()):
                if listening:
                    audio_idx += 1
                    if audio_idx >= n_chunks:
                        break
                    # Append [AUDIO_BEGIN, PAD*N, ASSISTANT]
                    new_tokens = torch.LongTensor(
                        [AUDIO_BEGIN] + [PAD] * AUDIO_TOKENS_PER_CHUNK + [ASSISTANT]
                    ).to(device)
                    new_positions = input_pos[-1] + torch.arange(1, len(new_tokens) + 1, device=device)
                    token = torch.cat([token, new_tokens])
                    input_pos = torch.cat([input_pos, new_positions])
                    input_pos_maxp1 = input_pos_maxp1 + len(new_tokens)

                    logits = model(
                        token.view(1, -1), None, 1,
                        all_chunks[audio_idx].to(device),
                        input_pos,
                        input_pos_maxp1=input_pos_maxp1,
                        audio_tokens_per_chunk=AUDIO_TOKENS_PER_CHUNK,
                    )
                else:
                    logits = model(
                        token.view(1, -1), None, 1,
                        None,
                        input_pos,
                        input_pos_maxp1=input_pos_maxp1,
                        audio_tokens_per_chunk=AUDIO_TOKENS_PER_CHUNK,
                    )

                token = sample(logits, temperature=0.0, top_p=0.0).to(torch.int64)
                int_token = token.item()
                input_pos = input_pos[-1].unsqueeze(0).add_(1)
                input_pos_maxp1 = input_pos_maxp1 + 1

                if listening:
                    if int_token == TEXT_BEGIN:
                        listening = False
                        current_text = [int_token]
                    elif int_token == KEEP_SILENCE:
                        pass  # model stays silent
                    else:
                        break
                else:
                    current_text.append(int_token)
                    if int_token == TEXT_END:
                        # Decode the turn (skip TEXT_BEGIN and optional emotion)
                        text_tokens = current_text[1:-1]  # strip TEXT_BEGIN and TEXT_END
                        if text_tokens:
                            # Check if first token is an emotion tag
                            if text_tokens[0] in EMOTION_EMOJI:
                                emotion_tag = EMOTION_EMOJI[text_tokens[0]]
                                text_tokens = text_tokens[1:]
                            decoded = tokenizer.decode(torch.tensor(text_tokens))
                            if decoded:
                                response_parts.append((emotion_tag, decoded))
                        current_text = []
                        listening = True
                    elif len(current_text) >= 2:
                        # Could be streaming text, but we collect at TEXT_END
                        pass

        # Handle any incomplete turn
        if not listening and len(current_text) > 1:
            text_tokens = current_text[1:]  # skip TEXT_BEGIN
            if text_tokens and text_tokens[0] in EMOTION_EMOJI:
                emotion_tag = EMOTION_EMOJI[text_tokens[0]]
                text_tokens = text_tokens[1:]
            decoded = tokenizer.decode(torch.tensor(text_tokens))
            if decoded:
                response_parts.append((emotion_tag, decoded))

    finally:
        model.clear_kv_cache()

    elapsed = time.perf_counter() - t0

    if not response_parts:
        return "The model listened to the audio but chose to stay silent. (This is expected for non-speech sounds when no response is warranted.)", \
               f"Processed {n_chunks} audio chunks in {elapsed:.1f}s — model stayed silent"

    # Format response
    formatted = []
    for emoji, text in response_parts:
        if emoji:
            formatted.append(f"{emoji} {text}")
        else:
            formatted.append(text)
    full_response = "\n".join(formatted)

    status = f"Processed {n_chunks} audio chunks in {elapsed:.1f}s — {len(response_parts)} response(s)"

    return full_response, status


# ── Gradio UI ──────────────────────────────────────────────

CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown("# 🎙️ Audio Interaction Model")
    gr.Markdown(
        "Upload or record an audio clip — the model listens and responds with text. "
        "This is a demo of [AudioInteraction](https://huggingface.co/zhifeixie/AudioInteraction), "
        "a streaming audio-language model that perceives audio and decides when to speak."
    )

    with gr.Column(elem_id="col-container"):
        with gr.Row():
            audio_input = gr.Audio(
                label="Audio Input",
                type="filepath",
                sources=["microphone", "upload"],
            )

        text_input = gr.Textbox(
            label="Text Instruction (optional)",
            placeholder="e.g., 'Translate this to English' or leave blank for auto-response",
            lines=2,
        )

        run_btn = gr.Button("Run", variant="primary")

        with gr.Row():
            response_output = gr.Textbox(
                label="Model Response",
                lines=8,
                interactive=False,
            )

        status_output = gr.Textbox(
            label="Status",
            interactive=False,
        )

    run_btn.click(
        fn=interact,
        inputs=[audio_input, text_input],
        outputs=[response_output, status_output],
    )

    gr.Examples(
        examples=[
            ["samples/sample01_count_bark.wav"],
            ["samples/sample03_cough_music.wav"],
            ["samples/sample02a_translate.wav"],
            ["samples/sample02b_translate.wav"],
        ],
        inputs=[audio_input],
        outputs=[response_output, status_output],
        fn=interact,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True)