#!/usr/bin/env python3
"""Decode tern_tc token ids to text with the published 8K tokenizer.

Pure-python GPT-2 byte-level BPE decoder over results/data8k/tokenizer.json
(no dependencies beyond the standard library) -- the same file the C engine
and the harness use.

Usage:
    python3 scripts/decode_ids.py 1 3 204 276 405 659 85 1516
    python3 scripts/decode_ids.py --specials 1 3 204 276 405 659 85 1516
"""

import json
import os
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


def _byte_decoder():
    cs = list(range(ord("!"), ord("~") + 1)) + list(range(ord("\xa1"), ord("\xac") + 1)) \
        + list(range(ord("\xae"), ord("\xff") + 1))
    n = 0
    table = {}
    for b in range(2 ** 8):
        if b in cs:
            table[chr(b)] = b
        else:
            table[chr(256 + n)] = b
            n += 1
    return table


BYTE_DECODER = _byte_decoder()


def decode(ids, tokenizer, skip_specials=True):
    vocab = tokenizer["model"]["vocab"]
    id2tok = {v: k for k, v in vocab.items()}
    special = {a["id"]: a["content"] for a in tokenizer.get("added_tokens", [])}
    parts = []
    for i in ids:
        if i in special:
            if not skip_specials:
                parts.append(special[i])
            continue
        parts.append(id2tok[i])
    joined = "".join(parts)
    return bytes(BYTE_DECODER[c] for c in joined).decode("utf-8", errors="replace")


def main():
    args = sys.argv[1:]
    skip = True
    if args and args[0] == "--specials":
        skip = False
        args = args[1:]
    if not args:
        print(__doc__)
        sys.exit(1)
    ids = [int(a) for a in args]
    tok = json.load(open(os.path.join(ROOT, "results/data8k/tokenizer.json")))
    print(decode(ids, tok, skip_specials=skip))


if __name__ == "__main__":
    main()