#!/usr/bin/env python3
"""Pack the unchanged Reader source for a 4 KiB sliding-window nRF52 decoder."""
import argparse
import hashlib
from functools import lru_cache
import shutil
import subprocess
import tempfile
import json
from pathlib import Path
import re
import zlib

import pack_reader as blocks

HEADER = blocks.ROOT / "src/helpers/bible/ReaderStreamData.generated.h"
COMPACT_HEADER = blocks.ROOT / "src/helpers/bible/ReaderCompactStreamData.generated.h"
TOKENS = Path(__file__).with_name("reader-stream-tokens.json")
COMPRESSOR = Path(__file__).with_name("zopfli4k")
STREAM_LIMIT = 17000
COMPACT_STREAM_LIMIT = 60000


@lru_cache(maxsize=1)
def host_compressor():
    compiler = shutil.which("cc") or shutil.which("gcc") or shutil.which("clang")
    if not compiler:
        raise RuntimeError("Regenerating the nRF52 Reader asset requires a host C compiler")
    directory = tempfile.TemporaryDirectory(prefix="mesh-reader-zopfli4k-")
    binary = Path(directory.name) / "pack4k"
    subprocess.run([compiler, "-std=c11", "-O2", "-DNDEBUG", "-I" + str(COMPRESSOR),
                    *map(str, sorted(COMPRESSOR.glob("*.c"))), "-lm", "-o", str(binary)], check=True)
    return directory, binary


@lru_cache(maxsize=16)
def compress_4k(data):
    directory, binary = host_compressor()
    source = Path(directory.name) / "input"
    target = Path(directory.name) / "output"
    source.write_bytes(data)
    subprocess.run([str(binary), str(source), str(target), "1000"], check=True, capture_output=True)
    packed = target.read_bytes()
    assert zlib.decompress(packed, wbits=-12) == data
    return packed


def token_dictionary(verses):
    # The vocabulary is an independently tunable compression asset. Tokens
    # only substitute exact source bytes, including matches inside a word.
    selected = json.loads(TOKENS.read_text(encoding="ascii"))
    if len(selected) > 128 or len(set(selected)) != len(selected):
        raise ValueError("Reader vocabulary must contain at most 128 unique byte strings")
    words = [word.encode("ascii") for word in selected]
    if any(not word or any(byte < 32 or byte >= 127 for byte in word) for word in words):
        raise ValueError("Reader tokens must be nonempty printable ASCII")
    return [word for word in words if any(word in verse for verse in verses)]


def encode(document, *, compact=False):
    verses = blocks.source_verses(document)
    words = token_dictionary(verses)
    tokens = {word: bytes([128 + index]) for index, word in enumerate(words)}
    # Tokens are exact byte strings, including any leading space. Matching
    # inside another word is lossless too: expansion restores those same bytes.
    pattern = re.compile(b"|".join(re.escape(word) for word in sorted(words, key=lambda x: (-len(x), x)))) if words else None
    encoded_verses = []
    dictionary_size = sum(map(len, words)) + len(words)
    for verse in verses:
        encoded = pattern.sub(lambda match: tokens[match.group()], verse) if pattern else verse
        if encoded != verse and len(encoded) + len(verse) + dictionary_size + 2 > blocks.BLOCK_SIZE:
            encoded = verse  # Literal output needs no dictionary workspace.
        encoded_verses.append(encoded)
    records = [verse + b"\0" for verse in encoded_verses]
    chunks, first, current = [], 0, b""
    for index, record in enumerate(records):
        if current and len(current) + len(record) > (COMPACT_STREAM_LIMIT if compact else STREAM_LIMIT):
            chunks.append((first, current))
            first, current = index, b""
        current += record
    chunks.append((first, current))
    # A single stream needs no framing. Larger corpora retain independently
    # bounded 12-bit histories and six-byte headers, without another buffer.
    if compact and len(chunks) == 1:
        return compress_4k(chunks[0][1]), b"".join(records), words, 0
    packed = bytearray()
    for first, plain in chunks:
        compressed = compress_4k(plain)
        packed += (first.to_bytes(2, "little") + len(plain).to_bytes(2, "little")
                   + len(compressed).to_bytes(2, "little") + compressed)
    return bytes(packed), b"".join(records), words, len(chunks)


def render(document, *, compact=False):
    packed, plain, words, chunk_count = encode(document, compact=compact)
    dictionary_plain = b"\0".join(words) + b"\0" if words else b""
    dictionary = compress_4k(dictionary_plain) if words else b""
    digest = hashlib.sha256(json.dumps(document, ensure_ascii=False, sort_keys=True,
                                     separators=(",", ":")).encode()).hexdigest()
    command = "tools/bible/pack_reader_stream.py" + (" --compact" if compact else "")
    lines = [f"// Generated by {command}; do not edit.",
             "// Offline Reader corpus. Public domain; see tools/bible/README.md.",
             "// Google Zopfli, 1000 iterations, 4096-byte history; exact ASCII text retained.",
             f"// Canonical source SHA-256: {digest}",
             "#pragma once", '#include "ReaderStreamLookup.h"',
             "namespace mesh { namespace bible { namespace generated {",
             "static_assert(kBlockSize == 4096, \"Reader history must remain 4096 bytes\");",
             "static const uint8_t readerData[] = {"]
    for start in range(0, len(packed), 20):
        lines.append("  " + ",".join(f"0x{v:02x}" for v in packed[start:start + 20]) + ",")
    lines += ["};", "static const uint8_t readerWords[] = {"]
    for start in range(0, len(dictionary), 20):
        lines.append("  " + ",".join(f"0x{v:02x}" for v in dictionary[start:start + 20]) + ",")
    if not dictionary:
        lines.append("  0, // Empty vocabulary: literal records remain supported.")
    lines += ["};", "static const StreamTokenCorpus readerCorpus = {",
              f"  {json.dumps(document['translation'])},",
              f"  {json.dumps(document['attribution'])},",
              f"  readerData, sizeof(readerData), {len(plain)},",
              f"  readerWords, {len(dictionary)}, {len(dictionary_plain)},",
              f"  {len(words)}, kVerseCount, {chunk_count}", "};", "}}}", ""]
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, default=blocks.SOURCE)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--compact", action="store_true", help="RAK4631 Full: smaller asset with slower cold lookup")
    parser.add_argument("--check", action="store_true")
    args = parser.parse_args()
    document = json.loads(args.source.read_text(encoding="utf-8"))
    output = args.output or (COMPACT_HEADER if args.compact else HEADER)
    header = render(document, compact=args.compact)
    if args.check:
        if output.read_text(encoding="ascii") != header:
            raise SystemExit("Reader stream header differs; regenerate with tools/bible/pack_reader_stream.py")
    else:
        output.write_text(header, encoding="ascii", newline="\n")
    packed, plain, words, chunk_count = encode(document, compact=args.compact)
    dictionary_plain = b"\0".join(words) + b"\0" if words else b""
    dictionary_size = len(compress_4k(dictionary_plain)) if words else 0
    print(f"Reader stream: {len(blocks.references())} unchanged verses, "
          f"{len(plain)} encoded bytes -> {len(packed)} DEFLATE + "
          f"{dictionary_size} dictionary bytes; {chunk_count or 1} independent streams, 4096-byte history")


if __name__ == "__main__":
    main()