import os
import traceback
from functools import lru_cache

import gradio as gr
import numpy as np
from datasets import load_dataset


DEFAULT_DATASET = "telecomadm1145/ASMR-dataviewer-fixed"
DEFAULT_SPLIT = "train"
TOKENIZER_ID = os.getenv("TOKENIZER_ID", "Qwen/Qwen3-TTS-Tokenizer-12Hz")


def _short_text(x, n=120):
    if x is None:
        return ""
    x = str(x).replace("\n", " ")
    return x if len(x) <= n else x[:n] + "..."


def _token_shape(tokens):
    if tokens is None:
        return "None"

    try:
        outer = len(tokens)
    except Exception:
        return "unknown"

    if outer == 0:
        return "0"

    try:
        first = tokens[0]
        inner = len(first)
        return f"{outer} x {inner}"
    except Exception:
        return f"{outer}"


def _token_head(tokens, n=12):
    if tokens is None:
        return ""

    try:
        if len(tokens) == 0:
            return "[]"

        first = tokens[0]

        if isinstance(first, (list, tuple)):
            return str(first[:n])
        else:
            return str(tokens[:n])
    except Exception:
        return ""


def make_stream_state(repo_id, split):
    ds = load_dataset(repo_id, split=split, streaming=True)
    return {
        "repo_id": repo_id,
        "split": split,
        "iterator": iter(ds),
        "page_rows": {},
        "next_index": 0,
        "finished": False,
    }


def row_to_preview(abs_idx, row):
    tokens = row.get("audio_tokens", None)

    return {
        "row_id": abs_idx,
        "text": _short_text(row.get("text", "")),
        "num_frames": row.get("num_frames", None),
        "num_quantizers": row.get("num_quantizers", None),
        "audio_tokens_shape": _token_shape(tokens),
        "tokens_head": _token_head(tokens),
    }


def load_next_page(repo_id, split, page_size, state):
    repo_id = repo_id.strip()
    split = split.strip()
    page_size = int(page_size)

    if page_size <= 0:
        page_size = 10

    if (
        state is None
        or state.get("repo_id") != repo_id
        or state.get("split") != split
    ):
        state = make_stream_state(repo_id, split)

    if state.get("finished"):
        previews = []
        choices = []
        return (
            state,
            previews,
            gr.update(choices=choices, value=None),
            "已经到达数据流结尾。"
        )

    previews = []
    page_rows = {}

    loaded = 0

    try:
        for _ in range(page_size):
            try:
                row = next(state["iterator"])
            except StopIteration:
                state["finished"] = True
                break

            abs_idx = state["next_index"]
            state["next_index"] += 1

            page_rows[str(abs_idx)] = row
            previews.append(row_to_preview(abs_idx, row))
            loaded += 1

    except Exception as e:
        return (
            state,
            previews,
            gr.update(choices=[], value=None),
            f"加载失败：{repr(e)}\n\n{traceback.format_exc()}"
        )

    state["page_rows"] = page_rows

    choices = []
    for item in previews:
        label = f'{item["row_id"]} | {_short_text(item["text"], 60)}'
        value = str(item["row_id"])
        choices.append((label, value))

    status = f"本页加载 {loaded} 条。当前流式位置：{state['next_index']}"

    if state.get("finished"):
        status += "，已到达结尾。"

    return (
        state,
        previews,
        gr.update(choices=choices, value=choices[0][1] if choices else None),
        status,
    )


def reset_stream(repo_id, split):
    repo_id = repo_id.strip()
    split = split.strip()

    try:
        state = make_stream_state(repo_id, split)
        return (
            state,
            [],
            gr.update(choices=[], value=None),
            f"已重置流：{repo_id} / {split}"
        )
    except Exception as e:
        return (
            None,
            [],
            gr.update(choices=[], value=None),
            f"重置失败：{repr(e)}\n\n{traceback.format_exc()}"
        )


@lru_cache(maxsize=1)
def get_tokenizer():
    """
    懒加载 tokenizer，避免 Gradio 启动时就加载模型。
    """
    from qwen_tts import Qwen3TTSTokenizer

    device_map = "cpu"

    try:
        import torch

        if torch.cuda.is_available():
            device_map = "cuda:0"
    except Exception:
        pass

    try:
        return Qwen3TTSTokenizer.from_pretrained(
            TOKENIZER_ID,
            device_map=device_map,
        )
    except TypeError:
        # 兼容旧版本 qwen-tts
        return Qwen3TTSTokenizer.from_pretrained(TOKENIZER_ID)


def tokens_to_array(tokens):
    """
    将 HF datasets 读出来的 list[list[int]] 转成 int64 numpy array。
    """
    if tokens is None:
        raise ValueError("audio_tokens is None")

    if len(tokens) == 0:
        raise ValueError("audio_tokens is empty")

    # 处理 list[list[int]]
    if isinstance(tokens[0], (list, tuple)):
        arr = np.array(
            [
                [0 if v is None else int(v) for v in inner]
                for inner in tokens
            ],
            dtype=np.int64,
        )
    else:
        arr = np.array(tokens, dtype=np.int64)

    if arr.ndim != 2:
        raise ValueError(f"audio_tokens should be 2-D, got shape={arr.shape}")

    return arr


def candidate_code_matrices(arr, num_frames=None, num_quantizers=None):
    """
    Qwen3-TTS 这类 codec token 常见形状可能是：

    1. [num_quantizers, num_frames]
    2. [num_frames, num_quantizers]

    数据集里有 num_frames / num_quantizers 字段，所以这里做自动判断。
    """
    candidates = []

    nq = None
    nf = None

    try:
        if num_quantizers is not None:
            nq = int(num_quantizers)
    except Exception:
        pass

    try:
        if num_frames is not None:
            nf = int(num_frames)
    except Exception:
        pass

    # 优先放 codebook-first: [num_quantizers, num_frames]
    if nq is not None:
        if arr.shape[0] == nq:
            candidates.append(arr)
        if arr.shape[1] == nq:
            candidates.append(arr.T)

    if nf is not None:
        if arr.shape[1] == nf:
            candidates.append(arr)
        if arr.shape[0] == nf:
            candidates.append(arr.T)

    # 兜底
    candidates.append(arr)
    candidates.append(arr.T)

    # 去重
    unique = []
    seen = set()

    for x in candidates:
        key = (x.shape[0], x.shape[1], x.strides)
        if key not in seen:
            seen.add(key)
            unique.append(np.ascontiguousarray(x))

    return unique


def parse_decode_output(out):
    """
    兼容 tokenizer.decode 返回：
    - wavs, sr
    - wavs
    """
    if isinstance(out, tuple) and len(out) == 2:
        wavs, sr = out
    else:
        wavs = out
        sr = 24000

    if isinstance(wavs, (list, tuple)):
        wav = wavs[0]
    else:
        wav = wavs

    try:
        import torch

        if isinstance(wav, torch.Tensor):
            wav = wav.detach().cpu().numpy()
    except Exception:
        pass

    wav = np.asarray(wav)

    # squeeze batch/channel
    wav = np.squeeze(wav)

    # Gradio Audio 期望 mono: [samples]，stereo/multichannel: [samples, channels]
    if wav.ndim == 2:
        # 如果是 [channels, samples]，转成 [samples, channels]
        if wav.shape[0] <= 8 and wav.shape[1] > wav.shape[0]:
            wav = wav.T

    wav = wav.astype(np.float32)

    # 简单防爆音裁剪
    wav = np.clip(wav, -1.0, 1.0)

    return int(sr), wav


def decode_selected(row_id, state):
    if state is None:
        return None, "请先加载数据。"

    if not row_id:
        return None, "请先选择一行。"

    page_rows = state.get("page_rows", {})
    row = page_rows.get(str(row_id))

    if row is None:
        return None, "当前页找不到该 row_id，请重新选择。"

    tokens = row.get("audio_tokens", None)

    try:
        arr = tokens_to_array(tokens)

        num_frames = row.get("num_frames", None)
        num_quantizers = row.get("num_quantizers", None)

        matrices = candidate_code_matrices(
            arr,
            num_frames=num_frames,
            num_quantizers=num_quantizers,
        )

        tokenizer = get_tokenizer()

        errors = []

        for mat in matrices:
            codes = mat.tolist()

            # 不同版本 qwen-tts 的 decode 输入格式可能略有差异。
            # 这里按常见格式依次尝试。
            decode_inputs = [
                #[codes],  # batch of one, shape: [1, Q, T]
                #codes,    # shape: [Q, T]
                {
                    "audio_codes": [codes],
                    #"num_frames": [int(num_frames)] if num_frames is not None else None,
                    #"num_quantizers": [int(num_quantizers)] if num_quantizers is not None else None,
                },
                #{
                #    "codes": [codes],
                #    "num_frames": [int(num_frames)] if num_frames is not None else None,
                #    "num_quantizers": [int(num_quantizers)] if num_quantizers is not None else None,
                #},
            ]

            for decode_input in decode_inputs:
                try:
                    out = tokenizer.decode(decode_input)
                    sr, wav = parse_decode_output(out)

                    msg = (
                        f"解码成功。\n"
                        f"row_id={row_id}\n"
                        f"token_matrix_shape={mat.shape}\n"
                        f"sample_rate={sr}"
                    )

                    return (sr, wav), msg

                except Exception as e:
                    errors.append(
                        f"shape={mat.shape}, input_type={type(decode_input)}: {repr(e)}"
                    )

        return (
            None,
            "所有解码尝试均失败。\n\n"
            + "\n".join(errors[-10:])
        )

    except Exception as e:
        return (
            None,
            f"解码失败：{repr(e)}\n\n{traceback.format_exc()}"
        )


with gr.Blocks(title="ASMR Dataset Streaming Preview") as demo:
    gr.Markdown(
        """
# ASMR Dataset Streaming Preview

数据集：

```text
telecomadm1145/ASMR-dataviewer-fixed
```

该数据集字段主要包括：

- `text`
- `audio_tokens`
- `num_frames`
- `num_quantizers`

本页面使用 `datasets` 的 `streaming=True` 逐条流式读取，不会一次性下载整个数据集。
"""
    )

    state = gr.State(None)

    with gr.Row():
        repo_id = gr.Textbox(
            label="Dataset Repo",
            value=DEFAULT_DATASET,
            scale=3,
        )
        split = gr.Textbox(
            label="Split",
            value=DEFAULT_SPLIT,
            scale=1,
        )
        page_size = gr.Number(
            label="每次加载条数",
            value=10,
            precision=0,
            scale=1,
        )

    with gr.Row():
        load_btn = gr.Button("加载下一页", variant="primary")
        reset_btn = gr.Button("重置流")

    status = gr.Textbox(
        label="状态",
        value="等待加载。",
        lines=4,
    )

    table = gr.Dataframe(
        label="预览",
        headers=[
            "row_id",
            "text",
            "num_frames",
            "num_quantizers",
            "audio_tokens_shape",
            "tokens_head",
        ],
        datatype=[
            "number",
            "str",
            "number",
            "number",
            "str",
            "str",
        ],
        interactive=False,
        wrap=True,
    )

    gr.Markdown("## 解码试听")

    row_select = gr.Dropdown(
        label="选择当前页中的一行",
        choices=[],
        value=None,
    )

    decode_btn = gr.Button("解码 audio_tokens", variant="secondary")

    audio = gr.Audio(
        label="Decoded Audio",
        type="numpy",
    )

    decode_status = gr.Textbox(
        label="解码状态",
        lines=8,
    )

    load_btn.click(
        fn=load_next_page,
        inputs=[repo_id, split, page_size, state],
        outputs=[state, table, row_select, status],
    )

    reset_btn.click(
        fn=reset_stream,
        inputs=[repo_id, split],
        outputs=[state, table, row_select, status],
    )

    decode_btn.click(
        fn=decode_selected,
        inputs=[row_select, state],
        outputs=[audio, decode_status],
    )


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