import os

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

import spaces  # noqa: E402  (must precede torch)

import json  # noqa: E402
import re  # noqa: E402
import time  # noqa: E402
from pathlib import Path  # noqa: E402

import gradio as gr  # noqa: E402
import librosa  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402
from transformers import (  # noqa: E402
    AutoModelForCausalLM,
    AutoTokenizer,
    LogitsProcessor,
    LogitsProcessorList,
    WhisperFeatureExtractor,
)

from hy_simulmt.model import build_llm_view, encode_audio, load_audio_tower, speech_token_count  # noqa: E402

MODEL_ID = "tencent/Hy-SimulMT-30B-A3B"
DTYPE = torch.bfloat16
SR = 16000
MAX_SECONDS = 60

# ----------------------------------------------------------------------------
# Load (module scope, eager .to("cuda") -- ZeroGPU streams weights in on first call)
# ----------------------------------------------------------------------------
root = Path(
    snapshot_download(MODEL_ID, allow_patterns=["*.safetensors", "*.json", "*.jinja"])
)
view = build_llm_view(root, Path("/tmp/hy_simulmt_llm_view"))
tokenizer = AutoTokenizer.from_pretrained(view)
feature_extractor = WhisperFeatureExtractor.from_pretrained(root)
llm = AutoModelForCausalLM.from_pretrained(view, dtype=DTYPE, attn_implementation="sdpa").eval()
encoder, adapter = load_audio_tower(root, DTYPE)
llm.to("cuda")
encoder.to("cuda")
adapter.to("cuda")
print("Hy-SimulMT loaded", flush=True)

# ----------------------------------------------------------------------------
# Protocol constants (SimulInfer/simulinfer/translation.py + templates/simulst.jinja)
# ----------------------------------------------------------------------------
HEADER = "<｜hy_begin▁of▁sentence｜><｜reasoning_mode｜>reasoning_effort:no_think"
USER, ASSISTANT, EOT = "<｜hy_User｜>", "<｜hy_Assistant｜><think></think>", "<｜hy_EOT｜>"
AUDIO_BEGIN = tokenizer.convert_tokens_to_ids("<｜hy_place▁holder▁no▁6｜>")
AUDIO_TOKEN = tokenizer.convert_tokens_to_ids("<｜hy_place▁holder▁no▁5｜>")
AUDIO_END = tokenizer.convert_tokens_to_ids("<｜hy_place▁holder▁no▁8｜>")
STOP_IDS = [
    tokenizer.convert_tokens_to_ids(EOT),
    tokenizer.convert_tokens_to_ids("<｜hy_end▁of▁sentence｜>"),
    tokenizer.eos_token_id,
]
INSTRUCTIONS = {
    "zh-en": "请完成下面的同声传译任务，将陆续收到的语音翻译成英文。请结合上下文边听边译，输出能够确定的译文，不必等待整句结束。",
    "en-zh": "请完成下面的同声传译任务，将陆续收到的语音翻译成中文。请结合上下文边听边译，输出能够确定的译文，不必等待整句结束。",
}
ST_INSTRUCTIONS = {
    "zh-en": "把这段语音翻译成英文，要全部输出，不要分块。",
    "en-zh": "把这段语音翻译成中文，要全部输出，不要分块。",
}
DIRECTIONS = {"Chinese → English": "zh-en", "English → Chinese": "en-zh"}
CHUNK = 1600 * 16  # first_chunk_ms = chunk_ms = 1600
MAX_TOKENS, ST_MAX_TOKENS = 64, 256
ST_WAIT_CHUNKS = 4
SILENCE_RMS = 0.003
TRAILING_SILENCE = int(2.4 * SR)


# Target-language mask for zh->en (SimulInfer target_language.py: block Han / partial-UTF8 tokens)
def _blocked_han_ids():
    han = re.compile("[⺀-⿟〇㇀-㇯㐀-䶿一-鿿豈-﫿\U00020000-\U0003347f]")
    values = list(range(33, 127)) + list(range(161, 173)) + list(range(174, 256))
    chars, extra = values[:], 0
    for v in range(256):
        if v not in values:
            values.append(v)
            chars.append(256 + extra)
            extra += 1
    decode = dict(zip(map(chr, chars), values))
    data = json.loads((root / "tokenizer.json").read_text())
    blocked = set()
    for tok, tid in data["model"]["vocab"].items():
        try:
            text = bytes(decode[c] for c in tok).decode("utf-8", errors="strict")
        except (KeyError, UnicodeDecodeError):
            blocked.add(tid)
            continue
        if han.search(text):
            blocked.add(tid)
    for item in data.get("added_tokens", []):
        if not item.get("special") and han.search(item["content"]):
            blocked.add(item["id"])
    return sorted(blocked)


HAN_IDS = torch.tensor(_blocked_han_ids(), dtype=torch.long)


class BlockTokens(LogitsProcessor):
    def __init__(self, ids):
        self.ids = ids

    def __call__(self, input_ids, scores):
        scores[:, self.ids.to(scores.device)] = -float("inf")
        return scores


# ----------------------------------------------------------------------------
# Inference helpers
# ----------------------------------------------------------------------------
def pcm16(x: np.ndarray) -> np.ndarray:
    """SimulInfer sends model audio as PCM16 WAV; reproduce that quantisation."""
    return (np.clip(np.rint(np.clip(x, -1, 1) * 32768.0), -32768, 32767) / 32768.0).astype(np.float32)


def encode_part(samples: np.ndarray) -> torch.Tensor:
    f = feature_extractor([pcm16(samples)], sampling_rate=SR, return_attention_mask=True,
                          padding="longest", return_tensors="pt")
    n = int(f["attention_mask"].sum())
    emb = encode_audio(encoder, adapter, f["input_features"][0], n)
    assert emb.shape[0] == speech_token_count(n)
    return emb


def render(messages):
    """messages: list of ('user', [str | Tensor ...]) / ('assistant', str). Returns ids + audio embeds."""
    ids, audio = [], []

    def text(s):
        ids.extend(tokenizer.encode(s, add_special_tokens=False))

    text(HEADER)
    for role, content in messages:
        if role == "user":
            text(USER)
            for p in content:
                if isinstance(p, str):
                    text(p)
                else:
                    ids.append(AUDIO_BEGIN)
                    ids.extend([AUDIO_TOKEN] * p.shape[0])
                    ids.append(AUDIO_END)
                    audio.append(p)
        else:
            text(ASSISTANT + content + EOT)
    text(ASSISTANT)
    return ids, audio


@torch.inference_mode()
def generate(messages, max_new_tokens, direction, repetition_penalty=1.0):
    ids, audio = render(messages)
    input_ids = torch.tensor([ids], device="cuda")
    embeds = llm.get_input_embeddings()(input_ids)
    if audio:
        mask = input_ids[0] == AUDIO_TOKEN
        embeds[0, mask] = torch.cat(audio, 0).to(embeds.dtype)
    procs = LogitsProcessorList([BlockTokens(HAN_IDS)] if direction == "zh-en" else [])
    out = llm.generate(
        inputs_embeds=embeds,
        attention_mask=torch.ones_like(input_ids),
        max_new_tokens=max_new_tokens,
        do_sample=False,
        repetition_penalty=repetition_penalty,
        no_repeat_ngram_size=3,
        eos_token_id=STOP_IDS,
        pad_token_id=tokenizer.pad_token_id,
        logits_processor=procs,
    )[0].tolist()
    truncated = not (out and out[-1] in STOP_IDS)
    while out and out[-1] in STOP_IDS:
        out = out[:-1]
    return tokenizer.decode(out, skip_special_tokens=True), truncated


def load_audio(path):
    if path is None:
        raise gr.Error("Please upload or record an audio clip.")
    audio, _ = librosa.load(path, sr=SR, mono=True)
    if not len(audio):
        raise gr.Error("The audio clip is empty.")
    return audio[: MAX_SECONDS * SR].astype(np.float32)


def fmt_t(samples):
    return f"{samples / SR:5.1f}s"


def _duration(audio_path, direction="Chinese → English", mode="Simultaneous", *args, **kwargs):
    try:
        secs = min(MAX_SECONDS, librosa.get_duration(path=audio_path))
    except Exception:
        secs = MAX_SECONDS
    return int(8 + secs * (0.4 if mode == "Simultaneous" else 0.2))


@spaces.GPU(duration=_duration, size="xlarge")
def translate(audio_path: str, direction: str = "Chinese → English", mode: str = "Simultaneous"):
    """Translate Chinese/English speech with Hy-SimulMT.

    Args:
        audio_path: path or URL of a speech recording (Chinese or English, up to 60 s).
        direction: "Chinese → English" or "English → Chinese".
        mode: "Simultaneous" replays the clip in 1.6 s chunks like a live interpreter
              (incremental output); "Offline" translates the whole clip at once.

    Returns:
        The translated text and a per-chunk timeline (markdown).
    """
    d = DIRECTIONS.get(direction, direction)
    if d not in INSTRUCTIONS:
        raise gr.Error("Unknown direction")
    audio = load_audio(audio_path)
    t0 = time.perf_counter()

    if mode != "Simultaneous":
        parts = [ST_INSTRUCTIONS[d]] + [encode_part(audio[i:i + 480000]) for i in range(0, len(audio), 480000)]
        text, _ = generate([("user", parts)], ST_MAX_TOKENS, d, repetition_penalty=1.05)
        yield text, f"Offline translation of {len(audio) / SR:.1f}s audio in {time.perf_counter() - t0:.1f}s."
        return

    messages, pending, rows = [], [], []
    full = ""
    voiced_waits = continuous_silence = 0
    st_count = 0

    def st_flush(start_label):
        nonlocal full, voiced_waits, st_count
        if not any(not b[2] for b in pending):
            return ""
        samples = np.concatenate([b[0][: b[1]] for b in pending])
        parts = [ST_INSTRUCTIONS[d]] + [encode_part(samples[i:i + 480000]) for i in range(0, len(samples), 480000)]
        text, _ = generate([("user", parts)], ST_MAX_TOKENS, d, repetition_penalty=1.05)
        st_count += 1
        if not text.strip():
            return ""
        full += text
        if messages and messages[-1][0] == "assistant":
            messages[-1] = ("assistant", text)
        pending.clear()
        voiced_waits = 0
        return text

    def table():
        head = "| Audio | Action | Output |\n|---|---|---|\n"
        return head + "\n".join(rows[-200:])

    cursor = 0
    n_chunks = (len(audio) + CHUNK - 1) // CHUNK
    for idx in range(n_chunks):
        chunk = audio[idx * CHUNK:(idx + 1) * CHUNK]
        valid = len(chunk)
        final = idx == n_chunks - 1
        if valid < CHUNK:
            chunk = np.pad(chunk, (0, CHUNK - valid))
        begin, cursor = cursor, cursor + valid
        label = f"{fmt_t(begin)}–{fmt_t(cursor)}"
        silent = float(np.sqrt(np.mean(chunk[:valid] ** 2))) < SILENCE_RMS
        continuous_silence = continuous_silence + valid if silent else 0
        if silent and not messages:  # skip_initial_silence
            rows.append(f"| {label} | skip (silence) | |")
            yield full, table()
            continue
        turn = ("user", ([INSTRUCTIONS[d]] if not messages else []) + [encode_part(chunk)])
        text, truncated = generate(messages + [turn], MAX_TOKENS, d)
        messages += [turn, ("assistant", text)]
        pending.append((chunk, valid, silent))
        if text.strip():
            full += text
            pending.clear()
            voiced_waits = 0
            rows.append(f"| {label} | EMIT | {text.strip()} |")
        else:
            if not silent:
                voiced_waits += 1
            rows.append(f"| {label} | WAIT | |")
            periodic = voiced_waits >= ST_WAIT_CHUNKS
            silence_end = voiced_waits and silent and continuous_silence >= TRAILING_SILENCE
            if periodic or silence_end or final:
                st = st_flush(label)
                if st:
                    rows.append(f"| {label} | ST recovery | {st.strip()} |")
        if truncated and not final:  # max_tokens reached -> cut the run
            if any(not b[2] for b in pending):
                st = st_flush(label)
                if st:
                    rows.append(f"| {label} | ST recovery | {st.strip()} |")
            messages = []
            pending.clear()
            voiced_waits = 0
        yield full, table()

    elapsed = time.perf_counter() - t0
    rows.append(f"| done | {n_chunks} chunks, {st_count} ST recoveries | processed {len(audio) / SR:.1f}s audio in {elapsed:.1f}s |")
    yield full, table()


# ----------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Hy-SimulMT 30B-A3B\n"
            "Tencent Hunyuan's simultaneous speech translation model (Chinese ⇄ English). "
            "In **Simultaneous** mode the clip is replayed in 1.6 s chunks, like a live interpreter: "
            "for each chunk the model either emits translated text or waits for more context.\n\n"
            "[Model](https://huggingface.co/tencent/Hy-SimulMT-30B-A3B) · "
            "[GitHub](https://github.com/Tencent-Hunyuan/Hy-SimulMT) · "
            "Text output only; the official speech output and voice cloning need a separate TTS service."
        )
        with gr.Row():
            with gr.Column():
                audio_in = gr.Audio(label="Source speech (max 60 s)", sources=["upload", "microphone"], type="filepath")
                direction = gr.Radio(list(DIRECTIONS), value="Chinese → English", label="Direction")
                mode = gr.Radio(["Simultaneous", "Offline"], value="Simultaneous", label="Mode",
                                info="Offline translates the whole clip in one pass")
                run = gr.Button("Translate", variant="primary")
            with gr.Column():
                out_text = gr.Textbox(label="Translation", lines=8)
                with gr.Accordion("Chunk timeline", open=True):
                    timeline = gr.Markdown()
        gr.Examples(
            examples=[
                ["examples/zh_lunar_south_pole.wav", "Chinese → English"],
                ["examples/en_reasonable_measures.wav", "English → Chinese"],
                ["examples/zh_womens_basketball.wav", "Chinese → English"],
                ["examples/en_learning_culture.wav", "English → Chinese"],
            ],
            inputs=[audio_in, direction],
            outputs=[out_text, timeline],
            fn=translate,
            cache_examples=True,
            cache_mode="lazy",
        )
    run.click(translate, inputs=[audio_in, direction, mode], outputs=[out_text, timeline], api_name="translate")

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)