#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import sentencepiece as spm


def build_text_rows(token_ids: list[int], cfg: dict) -> list[list[int]]:
    row_width = int(cfg["n_vq"]) + 1
    rows = []
    for tid in token_ids:
        row = [int(cfg["audio_pad_token_id"])] * row_width
        row[0] = int(tid)
        rows.append(row)
    return rows


def build_audio_prefix_rows(prompt_audio_codes: list[list[int]], cfg: dict) -> list[list[int]]:
    row_width = int(cfg["n_vq"]) + 1
    rows = []
    for code_row in prompt_audio_codes:
        row = [int(cfg["audio_pad_token_id"])] * row_width
        row[0] = int(cfg["audio_user_slot_token_id"])
        for i in range(min(len(code_row), int(cfg["n_vq"]))):
            row[i + 1] = int(code_row[i])
        rows.append(row)
    return rows


def main() -> None:
    ap = argparse.ArgumentParser(description="Prepare C++ MOSS-TTS request rows with Python tokenizer.")
    ap.add_argument("--config-dir", default="configs")
    group = ap.add_mutually_exclusive_group(required=True)
    group.add_argument("--text")
    group.add_argument("--token-ids", help="comma/space separated SentencePiece token ids")
    ap.add_argument("--voice", default="Junhao")
    ap.add_argument("--output", default="cpp/cpp_request.txt")
    ap.add_argument("--seed", type=int, default=1234)
    ap.add_argument("--max-new-frames", type=int, default=128)
    ap.add_argument("--greedy-prefix-frames", type=int, default=4)
    ap.add_argument(
        "--assistant-random-u",
        type=float,
        default=None,
        help="fix the continue/end sampling value; 0.5 selects the more probable decision",
    )
    args = ap.parse_args()
    if args.max_new_frames <= 0:
        raise SystemExit("--max-new-frames must be positive")
    if args.greedy_prefix_frames < 0:
        raise SystemExit("--greedy-prefix-frames must be non-negative")
    if args.assistant_random_u is not None and not 0.0 <= args.assistant_random_u < 1.0:
        raise SystemExit("--assistant-random-u must be in [0, 1)")

    config_dir = Path(args.config_dir)
    manifest = json.loads((config_dir / "browser_poc_manifest.json").read_text("utf-8"))
    tts_meta = json.loads((config_dir / "tts_browser_onnx_meta.json").read_text("utf-8"))
    cfg = manifest["tts_config"]
    templates = manifest["prompt_templates"]

    if args.token_ids:
        token_ids = [int(x) for x in args.token_ids.replace(",", " ").split() if x.strip()]
    else:
        sp = spm.SentencePieceProcessor(model_file=str(config_dir / manifest["model_files"]["tokenizer_model"]))
        token_ids = [int(x) for x in sp.encode(str(args.text), out_type=int)]

    voice = None
    for item in manifest["builtin_voices"]:
        if item["voice"] == args.voice:
            voice = item
            break
    if voice is None:
        raise SystemExit(f"voice not found: {args.voice}")

    prefix_ids = [
        *templates["user_prompt_prefix_token_ids"],
        int(cfg["audio_start_token_id"]),
    ]
    suffix_ids = [
        int(cfg["audio_end_token_id"]),
        *templates["user_prompt_after_reference_token_ids"],
        *token_ids,
        *templates["assistant_prompt_prefix_token_ids"],
        int(cfg["audio_start_token_id"]),
    ]
    rows = [
        *build_text_rows(prefix_ids, cfg),
        *build_audio_prefix_rows(voice["prompt_audio_codes"], cfg),
        *build_text_rows(suffix_ids, cfg),
    ]
    actual_len = len(rows)
    static_len = int(tts_meta["static_shapes"]["prefill_seq"])
    row_width = int(cfg["n_vq"]) + 1
    if actual_len > static_len:
        raise SystemExit(f"request too long: actual={actual_len}, static={static_len}")
    pad_row = [int(cfg["audio_pad_token_id"])] * row_width
    pad_row[0] = int(cfg["pad_token_id"])
    rows = rows + [pad_row[:] for _ in range(static_len - actual_len)]
    mask = [1] * actual_len + [0] * (static_len - actual_len)

    # Serialize NumPy PCG64 draws so C++ and Python feed exactly the same
    # four-input local-fixed sampler.  Prefix frames consume their draws first,
    # then replace audio u with zero (top-1), matching the Python runtime.
    sampler_steps = int(args.max_new_frames)
    rng = np.random.default_rng(int(args.seed))
    assistant_random_u: list[np.float32] = []
    audio_random_u: list[list[np.float32]] = []
    for step in range(sampler_steps):
        sampled_assistant_u = np.float32(rng.random())
        if args.assistant_random_u is not None:
            sampled_assistant_u = np.float32(args.assistant_random_u)
        assistant_random_u.append(sampled_assistant_u)
        audio_row = np.asarray(rng.random(int(cfg["n_vq"])), dtype=np.float32)
        if step < int(args.greedy_prefix_frames):
            audio_row.fill(0.0)
        audio_random_u.append([np.float32(value) for value in audio_row])

    out = Path(args.output)
    out.parent.mkdir(parents=True, exist_ok=True)
    with out.open("w", encoding="utf-8") as f:
        f.write(f"actual_len {actual_len}\n")
        f.write(f"prefill_seq {static_len}\n")
        f.write(f"row_width {row_width}\n")
        f.write("input_ids\n")
        for row in rows:
            f.write(" ".join(str(int(x)) for x in row) + "\n")
        f.write("attention_mask\n")
        f.write(" ".join(str(int(x)) for x in mask) + "\n")
        f.write(f"sampler_steps {sampler_steps}\n")
        f.write("assistant_random_u\n")
        f.write(" ".join(format(float(value), ".9g") for value in assistant_random_u) + "\n")
        f.write("audio_random_u\n")
        for row in audio_random_u:
            f.write(" ".join(format(float(value), ".9g") for value in row) + "\n")
    print(
        f"wrote {out} actual_len={actual_len} text_tokens={len(token_ids)} "
        f"sampler_steps={sampler_steps} greedy_prefix_frames={args.greedy_prefix_frames}"
    )


if __name__ == "__main__":
    main()