#!/usr/bin/env python3
"""Build a llama.cpp static n-gram (lookup) cache from a navi export corpus.

llama.cpp dumps `common_ngram_cache` (an unordered_map) without header:

    repeat until EOF:
        int32 tokens[4]     # unused slots = -1 (LLAMA_TOKEN_NULL)
        int32 n_pairs       # > 0
        n_pairs x { int32 token, int32 count }   # token = the token after the n-gram

Only 2-grams are emitted: the static cache is looked up with
LLAMA_NGRAM_STATIC == 2, so 3- and 4-gram keys are dead weight.

Tokenization goes through a running llama-server (`POST /tokenize`,
add_special=false) so the 11-14 GiB model is loaded once instead of once per
message.

Usage:
    python3 ngram_build.py --input messages.jsonl.gz --out ngram-cache.bin \
        --server http://127.0.0.1:8199 --roles assistant --max-bytes 4000000
"""

from __future__ import annotations

import argparse
import gzip
import hashlib
import http.client
import json
import os
import struct
import sys
import time
from collections import Counter, defaultdict
from urllib.parse import urlsplit

LLAMA_TOKEN_NULL = -1
NGRAM = 4
STATIC_NGRAM = 2


class Tokenizer:
    """Keep-alive client for llama-server /tokenize."""

    def __init__(self, url: str):
        parts = urlsplit(url)
        self.host = parts.hostname or "127.0.0.1"
        self.port = parts.port or 80
        self.conn = http.client.HTTPConnection(self.host, self.port, timeout=120)

    def tokenize(self, text: str) -> list[int]:
        body = json.dumps({"content": text, "add_special": False})
        for attempt in (0, 1):
            try:
                self.conn.request(
                    "POST", "/tokenize", body=body,
                    headers={"Content-Type": "application/json"},
                )
                resp = self.conn.getresponse()
                raw = resp.read()
                if resp.status != 200:
                    raise RuntimeError(f"/tokenize HTTP {resp.status}: {raw[:300]!r}")
                data = json.loads(raw)
                tokens = data.get("tokens")
                if not isinstance(tokens, list):
                    raise RuntimeError(f"/tokenize: no 'tokens' in {raw[:300]!r}")
                return [int(t) for t in tokens]
            except (http.client.HTTPException, OSError):
                if attempt:
                    raise
                self.conn.close()
                self.conn = http.client.HTTPConnection(
                    self.host, self.port, timeout=120
                )
        raise AssertionError("unreachable")


def iter_messages(path: str):
    opener = gzip.open if path.endswith(".gz") else open
    with opener(path, "rt", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError as exc:
                print(f"skip malformed line: {exc}", file=sys.stderr)


def build(args) -> None:
    tok = Tokenizer(args.server)
    roles = set(args.roles.split(",")) if args.roles else None

    table: dict[tuple[int, ...], Counter] = defaultdict(Counter)
    n_msgs = 0
    n_tokens = 0
    content_bytes = 0
    t0 = time.time()

    for rec in iter_messages(args.input):
        if roles is not None and rec.get("role") not in roles:
            continue
        text = rec.get("content") or ""
        if not text.strip():
            continue
        if args.min_chars and len(text) < args.min_chars:
            continue
        if args.sample_percent < 100.0:
            digest = hashlib.sha256(
                f"{rec.get('session')}:{rec.get('seq')}".encode()
            ).digest()
            if int.from_bytes(digest[:4], "big") / 2**32 * 100 >= args.sample_percent:
                continue

        ids = tok.tokenize(text)
        n_msgs += 1
        n_tokens += len(ids)
        content_bytes += len(text.encode("utf-8"))

        # 2-grams inside one message only: message boundaries make no pairs.
        for i in range(len(ids) - STATIC_NGRAM):
            key = tuple(ids[i:i + STATIC_NGRAM]) + (LLAMA_TOKEN_NULL,) * (
                NGRAM - STATIC_NGRAM
            )
            table[key][ids[i + STATIC_NGRAM]] += 1

        if n_msgs % 1000 == 0:
            print(
                f"  {n_msgs} msgs, {n_tokens} tokens, "
                f"{content_bytes / 1e6:.2f} MB text, "
                f"{len(table)} keys, {time.time() - t0:.1f}s",
                file=sys.stderr,
            )

        if args.limit and n_msgs >= args.limit:
            break
        if args.max_bytes and content_bytes >= args.max_bytes:
            break

    print(
        f"tokenized: {n_msgs} messages, {n_tokens} tokens, "
        f"{content_bytes / 1e6:.2f} MB text, {time.time() - t0:.1f}s",
        file=sys.stderr,
    )

    # Truncate each distribution. Counts of the kept tokens stay exact; the
    # dropped tail only shrinks the denominator used by the percentage gates.
    n_pairs = 0
    n_dropped = 0
    n_lowconf = 0
    with open(args.out, "wb") as out:
        for key, counts in table.items():
            if len(counts) < args.min_count:
                continue
            total = sum(counts.values())
            top = max(counts.values())
            if total < args.min_total or top / total < args.min_top_share:
                n_lowconf += 1
                continue
            if args.top_k and len(counts) > args.top_k:
                kept = counts.most_common(args.top_k)
                n_dropped += len(counts) - args.top_k
            else:
                kept = list(counts.items())
            out.write(struct.pack(f"<{NGRAM}i", *key))
            out.write(struct.pack("<i", len(kept)))
            for token, count in kept:
                out.write(struct.pack("<ii", token, count))
            n_pairs += len(kept)

    print(
        f"wrote {args.out}: {len(table)} keys seen, {n_lowconf} dropped as "
        f"low-confidence, {n_pairs} pairs, {n_dropped} tail pairs dropped, "
        f"{os.path.getsize(args.out) / 1e6:.1f} MB",
        file=sys.stderr,
    )


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--input", required=True, help="corpus jsonl or jsonl.gz")
    ap.add_argument("--out", required=True, help="output cache file")
    ap.add_argument("--server", default="http://127.0.0.1:8199",
                    help="llama-server base URL used only for /tokenize")
    ap.add_argument("--roles", default="assistant",
                    help="comma separated roles, empty = all")
    ap.add_argument("--max-bytes", type=int, default=0,
                    help="stop once this many bytes of text are read (0 = all)")
    ap.add_argument("--limit", type=int, default=0,
                    help="stop after this many messages (0 = all)")
    ap.add_argument("--min-chars", type=int, default=0,
                    help="skip messages shorter than this")
    ap.add_argument("--sample-percent", type=float, default=100.0,
                    help="deterministic hash sampling of eligible messages, percent")
    ap.add_argument("--top-k", type=int, default=8,
                    help="keep at most K next-token counts per key (0 = keep all)")
    ap.add_argument("--min-count", type=int, default=1,
                    help="drop keys whose distribution has fewer than K tokens")
    ap.add_argument("--min-total", type=int, default=1,
                    help="drop keys seen fewer than K times in the corpus")
    ap.add_argument("--min-top-share", type=float, default=0.0,
                    help="drop keys whose most frequent next token has a share "
                         "below this fraction (0 = keep all)")
    args = ap.parse_args()
    build(args)
    return 0


if __name__ == "__main__":
    sys.exit(main())