#!/usr/bin/env python3
"""Reference pipeline that turns a Vietnamese speech-recognition transcript into clean written text with
Qwen2.5-Instruct (ONNX): punctuation and capitals, misheard words fixed from context, fillers and stutters
removed, and lists written as "- " bullet points.

The Kotlin (kotlin/src/main/kotlin/QwenRefiner.kt) and Swift (swift/Sources/QwenRefiner/) ports follow this file
step by step, so keep the three in sync:

1. Tokenizer: Qwen2's byte-level BPE from vocab.txt and merges.txt. Text is NFC-normalized and split with Qwen2's
   pre-tokenizer regex, each piece is mapped to byte-level characters and merged by rank. Special tokens are never
   parsed out of the transcript; the prompt inserts their ids directly, so a transcript cannot inject a turn.
2. Prompt: ChatML with a fixed system message and few-shot turns. Everything up to the transcript is the same for
   every call, so its KV cache is computed once when the model loads and reused.
3. Decoding: greedy, stopping at <|im_end|> or <|endoftext|>, at most max_new_tokens(n) tokens.
4. Merge: small models reword and drop words, so the output is aligned word by word with the transcript (LCS)
   and only trusted edits are kept: punctuation, capitals, line breaks and bullets; dropped fillers, stutters and
   list markers ("thứ hai là") before a bullet; a word replaced by a similar-sounding word (a misheard word
   fixed); spoken numbers written as digits. Any other change is undone with the transcript's own words.
"""

import argparse
import json
import sys
import unicodedata
from pathlib import Path

import numpy as np
import onnxruntime as ort
import regex

IM_START = "<|im_start|>"
IM_END = "<|im_end|>"
END_OF_TEXT = "<|endoftext|>"

# Qwen2's pre-tokenizer split (tokenizer.json, pre_tokenizer.pretokenizers[0].pattern.Regex).
PRETOKENIZE = regex.compile(
    r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
)

SYSTEM = """You correct Vietnamese speech-recognition transcripts. Copy the transcript word for word, and only:
- add punctuation and capital letters (sentence starts, names, places);
- replace a word the recognizer clearly misheard with the similar-sounding word the context needs;
- delete filler sounds (ờ, à, ừm, ừ) and words said twice by mistake;
- when the speaker lists several items, steps or points, put each one on its own line starting with "- ".
Never reword, shorten, translate, answer or add words. Keep names and foreign words you do not know as they are. Reply with the corrected transcript only."""

# (transcript, corrected) turns shown before every transcript.
EXAMPLES = [
    (
        "ờ hôm nay em cần làm ba việc thứ nhất là đi chợ mua rau thứ hai là đón con lúc năm giờ chiều "
        "và thứ ba là gọi điện cho cho mẹ",
        "Hôm nay em cần làm ba việc:\n- Đi chợ mua rau.\n- Đón con lúc 5 giờ chiều.\n- Gọi điện cho mẹ.",
    ),
    (
        "bạn có biết bệnh viện bạch mai ở đâu không mình cần đến đó trước tám giờ sáng mai",
        "Bạn có biết Bệnh viện Bạch Mai ở đâu không? Mình cần đến đó trước 8 giờ sáng mai.",
    ),
    (
        "chúng tôi đã ký hợp đồng với công ty vào thán trước và ừm sẽ giao hàng trong tuần lày",
        "Chúng tôi đã ký hợp đồng với công ty vào tháng trước và sẽ giao hàng trong tuần này.",
    ),
    (
        "nhà nghiên cứu jan rossby ở viện max planck cho biết hệ số kờ rô càng lớn thì sao càng sáng",
        "Nhà nghiên cứu Jan Rossby ở Viện Max Planck cho biết hệ số kờ rô càng lớn thì sao càng sáng.",
    ),
]

# Filler sounds the merge lets the model delete.
FILLERS = frozenset("ờ ơ à ừ ừm ừ hừm ưm ừa um uh ah ừhm".split())

# Words of a spoken list marker ("thứ hai là", "và ba là") the merge lets the model replace with a bullet.
LIST_MARKERS = frozenset("thứ nhất hai ba bốn năm sáu bảy tám chín mười một là và".split())

# Spoken digits, and the words that scale them, for checking a number the model wrote in digits.
DIGITS = {
    "không": 0, "một": 1, "mốt": 1, "hai": 2, "ba": 3, "bốn": 4, "tư": 4, "năm": 5, "lăm": 5, "nhăm": 5,
    "sáu": 6, "bảy": 7, "tám": 8, "chín": 9,
}
SCALES = {"nghìn": 1_000, "ngàn": 1_000, "triệu": 1_000_000, "tỷ": 1_000_000_000, "tỉ": 1_000_000_000}

def bytes_to_unicode() -> dict[int, str]:
    """GPT-2's byte-level alphabet: printable bytes map to themselves, the rest to chars from U+0100 up."""
    printable = list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
    chars, extra = {}, 0
    for b in range(256):
        if b in printable:
            chars[b] = chr(b)
        else:
            chars[b] = chr(256 + extra)
            extra += 1
    return chars


class Tokenizer:
    def __init__(self, repo: Path):
        tokens = (repo / "vocab.txt").read_text(encoding="utf-8").split("\n")[:-1]
        self.id_of = {t: i for i, t in enumerate(tokens)}
        self.tokens = tokens
        merges = (repo / "merges.txt").read_text(encoding="utf-8").split("\n")[:-1]
        self.rank = {tuple(m.split(" ")): r for r, m in enumerate(merges)}
        self.byte_char = bytes_to_unicode()
        self.char_byte = {c: b for b, c in self.byte_char.items()}
        self.cache: dict[str, list[int]] = {}

    def special(self, token: str) -> int:
        return self.id_of[token]

    def encode(self, text: str) -> list[int]:
        ids = []
        for piece in PRETOKENIZE.findall(unicodedata.normalize("NFC", text)):
            ids += self.bpe("".join(self.byte_char[b] for b in piece.encode("utf-8")))
        return ids

    def bpe(self, word: str) -> list[int]:
        cached = self.cache.get(word)
        if cached is not None:
            return cached
        parts = list(word)
        while len(parts) > 1:
            best, best_rank = -1, None
            for i in range(len(parts) - 1):
                r = self.rank.get((parts[i], parts[i + 1]))
                if r is not None and (best_rank is None or r < best_rank):
                    best, best_rank = i, r
            if best < 0:
                break
            merged = parts[best] + parts[best + 1]
            # Merge every occurrence of the best pair, left to right, as the reference BPE does.
            out, i = [], 0
            while i < len(parts):
                if i < len(parts) - 1 and parts[i] == parts[best] and parts[i + 1] == parts[best + 1]:
                    out.append(merged)
                    i += 2
                else:
                    out.append(parts[i])
                    i += 1
            parts = out
        ids = [self.id_of[p] for p in parts]
        self.cache[word] = ids
        return ids

    def decode(self, ids: list[int]) -> str:
        data = bytes(self.char_byte[c] for i in ids for c in self.tokens[i])
        return data.decode("utf-8", errors="replace")


def prompt_prefix(tok: Tokenizer) -> list[int]:
    """System message and few-shot turns, ending right before the transcript."""
    start, end, nl = tok.special(IM_START), tok.special(IM_END), tok.encode("\n")
    ids = [start] + tok.encode("system\n" + SYSTEM) + [end] + nl
    for transcript, edited in EXAMPLES:
        ids += [start] + tok.encode("user\n" + transcript) + [end] + nl
        ids += [start] + tok.encode("assistant\n" + edited) + [end] + nl
    return ids + [start] + tok.encode("user\n")


def prompt_suffix(tok: Tokenizer, transcript: str) -> list[int]:
    """The transcript and the opening of the assistant turn."""
    return tok.encode(transcript) + [tok.special(IM_END)] + tok.encode("\n") + [tok.special(IM_START)] + tok.encode("assistant\n")


def max_new_tokens(input_tokens: int) -> int:
    """Room for punctuation, bullet markers and line breaks on top of the input's length."""
    return min(input_tokens * 3 // 2 + 32, 1024)


def prepare(transcript: str) -> str:
    """NFC, whitespace collapsed. Whisper's punctuation is kept: it helps the model."""
    return " ".join(unicodedata.normalize("NFC", transcript).split())


class Unit:
    """A word of the text with what is written around it: punctuation before and after it, and whether it starts
    a line and a bullet. key is the word in lower case, for matching."""

    __slots__ = ("lead", "word", "trail", "newline", "bullet")

    def __init__(self, word: str, lead: str = "", trail: str = "", newline: bool = False, bullet: bool = False):
        self.word, self.lead, self.trail, self.newline, self.bullet = word, lead, trail, newline, bullet

    @property
    def key(self) -> str:
        return self.word.lower()


# A word: from its first to its last letter or digit; what is around it is its punctuation.
WORD_EDGES = regex.compile(r"^([^\p{L}\p{N}]*)(.*?)([^\p{L}\p{N}]*)$", regex.DOTALL)
BULLETS = ("-", "*", "•", "–")


def units(text: str) -> list[Unit]:
    """Splits text into Units. Line breaks and bullet markers go to the next word, punctuation-only pieces to the
    next word's lead (or the last word's trail at the end)."""
    out: list[Unit] = []
    newline = bullet = False
    lead = ""
    for piece in regex.findall(r"\n|[^\s]+", text):
        if piece == "\n":
            if out:
                newline = True
            continue
        before, word, after = WORD_EDGES.match(piece).groups()
        if not word:
            if piece in BULLETS and not lead:
                bullet = True
            elif out and not newline and not bullet:
                out[-1].trail += piece
            else:
                lead += piece
            continue
        if before[:1] in BULLETS and len(before) == 1 and not lead:
            bullet, before = True, ""
        out.append(Unit(word, lead + before, after, newline, bullet))
        newline = bullet = False
        lead = ""
    return out


def render(items: list[Unit]) -> str:
    parts = []
    for i, u in enumerate(items):
        if i > 0:
            parts.append("\n" if u.newline else " ")
        parts.append(("- " if u.bullet else "") + u.lead + u.word + u.trail)
    return "".join(parts)


def plain(word: str) -> str:
    """Lower case without tone and vowel marks (đ is d), so similar-sounding words compare close."""
    decomposed = unicodedata.normalize("NFD", word.lower().replace("đ", "d"))
    return "".join(c for c in decomposed if not unicodedata.combining(c))


def edit_distance(a: str, b: str) -> int:
    row = list(range(len(b) + 1))
    for i in range(1, len(a) + 1):
        prev, row[0] = row[0], i
        for j in range(1, len(b) + 1):
            cur = min(row[j] + 1, row[j - 1] + 1, prev + (a[i - 1] != b[j - 1]))
            prev, row[j] = row[j], cur
    return row[len(b)]


def sounds_alike(a: str, b: str) -> bool:
    """One edit apart without marks for short words, two for longer ones: lội/nội, thán/tháng, tuần/tuần."""
    pa, pb = plain(a), plain(b)
    return edit_distance(pa, pb) <= (1 if max(len(pa), len(pb)) <= 4 else 2)


def spoken_number(words: list[str]) -> int | None:
    """The value of a spoken Vietnamese number ("hai nghìn không trăm hai mươi sáu" is 2026), or None when the
    words are not one number."""
    total = group = 0  # group: the part below the current thousand
    digit: int | None = None
    for w in words:
        if w in DIGITS:
            if digit is not None:
                return None  # two digits in a row: a list of digits, not one number
            digit = DIGITS[w]
        elif w == "mười":
            if digit is not None:
                return None
            group += 10
        elif w == "mươi" and digit is not None:
            group, digit = group + digit * 10, None
        elif w == "trăm" and digit is not None:
            group, digit = group + digit * 100, None
        elif w in ("linh", "lẻ"):
            continue
        elif w in SCALES:
            total, group, digit = total + (group + (digit or 0)) * SCALES[w], 0, None
        else:
            return None
    if not words:
        return None
    return total + group + (digit or 0)


def written_number(word: str) -> int | None:
    """The value of digits the model wrote, with "." or "," as thousands separators: 15, 2026, 1.000."""
    if not regex.fullmatch(r"\d{1,3}([.,]\d{3})*|\d+", word):
        return None
    return int(word.replace(".", "").replace(",", ""))


def accept_replacement(old: list[Unit], new: list[Unit]) -> bool:
    """A misheard word fixed one for one, or a spoken number written as digits with the same value."""
    if len(old) == len(new) and all(sounds_alike(o.key, n.key) for o, n in zip(old, new)):
        return True
    if len(new) != 1:
        return False
    value = written_number(new[0].word)
    return value is not None and value == spoken_number([o.key for o in old])


def align(a: list[str], b: list[str]) -> list[tuple[str, int, int, int, int]]:
    """Longest common subsequence of a and b as difflib-style opcodes (tag, i1, i2, j1, j2); a match is taken as
    early as possible, and an unmatched run of a before one of b."""
    n, m = len(a), len(b)
    lcs = [[0] * (m + 1) for _ in range(n + 1)]
    for i in range(n - 1, -1, -1):
        for j in range(m - 1, -1, -1):
            lcs[i][j] = lcs[i + 1][j + 1] + 1 if a[i] == b[j] else max(lcs[i + 1][j], lcs[i][j + 1])
    ops, i, j = [], 0, 0
    while i < n or j < m:
        if i < n and j < m and a[i] == b[j]:
            i2, j2 = i, j
            while i2 < n and j2 < m and a[i2] == b[j2]:
                i2, j2 = i2 + 1, j2 + 1
            ops.append(("equal", i, i2, j, j2))
            i, j = i2, j2
            continue
        i2, j2 = i, j
        while i2 < n or j2 < m:
            if i2 < n and j2 < m and a[i2] == b[j2]:
                break
            if j2 >= m or (i2 < n and lcs[i2 + 1][j2] >= lcs[i2][j2 + 1]):
                i2 += 1
            else:
                j2 += 1
        tag = "replace" if i2 > i and j2 > j else ("delete" if i2 > i else "insert")
        ops.append((tag, i, i2, j, j2))
        i, j = i2, j2
    return ops


def may_delete(old: list[Unit], start: int, end: int, next_unit: Unit | None) -> list[bool]:
    """For each transcript word in old[start:end] the model left out: whether leaving it out is trusted."""
    run = old[start:end]
    if next_unit is not None and next_unit.bullet and all(u.key in LIST_MARKERS for u in run):
        return [True] * len(run)
    keep = []
    for k, u in enumerate(run):
        i = start + k
        repeated = (i > 0 and old[i - 1].key == u.key) or (i + 1 < len(old) and old[i + 1].key == u.key)
        keep.append(u.key in FILLERS or repeated)
    return keep


def merge(transcript: str, edited: str) -> str:
    """The model's text, with every edit that is not trusted undone (see the module docstring)."""
    old, new = units(transcript), units(edited)
    out: list[Unit] = []
    carry: Unit | None = None  # layout of dropped words, for the next word written

    def emit(u: Unit) -> None:
        nonlocal carry
        if carry is not None:
            u.newline, u.bullet = u.newline or carry.newline, u.bullet or carry.bullet
            u.lead = carry.lead + u.lead
            carry = None
        out.append(u)

    def drop(dropped: list[Unit]) -> None:
        """Keeps the layout and the closing punctuation of words that are left out."""
        nonlocal carry
        for u in dropped:
            if u.newline or u.bullet or u.lead:
                carry = carry or Unit("")
                carry.newline, carry.bullet = carry.newline or u.newline, carry.bullet or u.bullet
                carry.lead += u.lead
            if u.trail and out and not out[-1].trail:
                out[-1].trail = u.trail

    def restore(words: list[Unit], like: list[Unit]) -> list[Unit]:
        """The transcript's words in place of the model's [like]: its layout, its first capital, its last mark."""
        back = [Unit(u.word) for u in words]
        if like and back:
            back[0].newline, back[0].bullet, back[0].lead = like[0].newline, like[0].bullet, like[0].lead
            back[-1].trail = like[-1].trail
            if like[0].word[:1].isupper():
                back[0].word = back[0].word[:1].upper() + back[0].word[1:]
        return back

    for tag, i1, i2, j1, j2 in align([u.key for u in old], [u.key for u in new]):
        if tag == "equal":
            for u in new[j1:j2]:
                emit(u)
        elif tag == "insert":
            drop(new[j1:j2])
        elif tag == "replace" and accept_replacement(old[i1:i2], new[j1:j2]):
            for u in new[j1:j2]:
                emit(u)
        elif tag == "replace":
            for u in restore(old[i1:i2], new[j1:j2]):
                emit(u)
        else:  # delete
            next_unit = new[j1] if j1 < len(new) else None
            for u, ok in zip(old[i1:i2], may_delete(old, i1, i2, next_unit)):
                if not ok:
                    restored = Unit(u.word)
                    # The word belongs before the mark that ends the previous word's sentence or clause.
                    if out and out[-1].trail:
                        restored.trail, out[-1].trail = out[-1].trail, ""
                    emit(restored)
    if carry is not None and out and carry.lead:
        out[-1].trail += carry.lead
    capitalize_sentence_starts(out)
    return render(out)


def capitalize_sentence_starts(items: list[Unit]) -> None:
    """Restored words keep the transcript's lower case; one that now starts a sentence, line or bullet gets a capital."""
    for i, u in enumerate(items):
        starts = i == 0 or u.newline or u.bullet or items[i - 1].trail.rstrip("\"')»”").endswith((".", "?", "!"))
        if starts and u.word[:1].islower() and not u.lead:
            u.word = u.word[:1].upper() + u.word[1:]


class Refiner:
    def __init__(self, repo: Path, threads: int = 4):
        self.tok = Tokenizer(repo)
        options = ort.SessionOptions()
        options.intra_op_num_threads = threads
        self.session = ort.InferenceSession(str(repo / "onnx" / "model.onnx"), options, providers=["CPUExecutionProvider"])
        names = [i.name for i in self.session.get_inputs()]
        self.past_names = [n for n in names if n.startswith("past_key_values.")]
        self.uses_position_ids = "position_ids" in names
        config = json.loads((repo / "onnx" / "genai_config.json").read_text(encoding="utf-8"))["model"]["decoder"]
        self.kv_heads, self.head_size = config["num_key_value_heads"], config["head_size"]
        self.stop = {self.tok.special(IM_END), self.tok.special(END_OF_TEXT)}
        self.prefix = prompt_prefix(self.tok)
        empty = np.zeros((1, self.kv_heads, 0, self.head_size), np.float32)
        _, self.prefix_cache = self.forward(self.prefix, {n: empty for n in self.past_names}, 0)

    def forward(self, ids: list[int], cache: dict[str, np.ndarray], past_len: int) -> tuple[np.ndarray, dict[str, np.ndarray]]:
        """Runs ids after past_len cached tokens; returns the last position's logits and the grown cache."""
        total = past_len + len(ids)
        feed = {
            "input_ids": np.array([ids], np.int64),
            "attention_mask": np.ones((1, total), np.int64),
            **cache,
        }
        if self.uses_position_ids:
            feed["position_ids"] = np.arange(past_len, total, dtype=np.int64)[None, :]
        outputs = self.session.run(None, feed)
        out_names = [o.name for o in self.session.get_outputs()]
        result = dict(zip(out_names, outputs))
        new_cache = {n: result[n.replace("past_key_values.", "present.")] for n in self.past_names}
        return result["logits"][0, -1], new_cache

    def generate(self, transcript: str) -> list[int]:
        suffix = prompt_suffix(self.tok, transcript)
        logits, cache = self.forward(suffix, self.prefix_cache, len(self.prefix))
        length = len(self.prefix) + len(suffix)
        out: list[int] = []
        for _ in range(max_new_tokens(len(suffix))):
            nxt = int(np.argmax(logits))
            if nxt in self.stop:
                break
            out.append(nxt)
            logits, cache = self.forward([nxt], cache, length)
            length += 1
        return out

    def refine(self, transcript: str) -> str:
        text = prepare(transcript)
        if not text:
            return transcript
        edited = self.tok.decode(self.generate(text)).strip()
        merged = merge(text, edited) if edited else ""
        return merged if merged.strip() else text


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("repo", type=Path)
    parser.add_argument("text", nargs="+")
    args = parser.parse_args()
    refiner = Refiner(args.repo)
    for text in args.text:
        print(refiner.refine(text))
        sys.stdout.flush()


if __name__ == "__main__":
    main()