#!/usr/bin/env python3
"""
Compile Plover JSON dictionaries into Javelin's JSC4 binary format.

Produces a StenoDictionaryCollection binary that can be embedded
in ZMK firmware at a known flash address.

Usage:
    python compile_dict.py main.json -o steno_dict.bin --base-addr 0
    python compile_dict.py main.json extra.json -o steno_dict.bin

The --base-addr flag sets the XIP base address for pointer resolution.
When embedded via .incbin in .rodata, set to 0 and the linker resolves
the _javelin_dict_start symbol. Javelin dereferences pointers relative
to the binary's load address.

IMPORTANT: This compiler uses position-independent pointers (offsets
from the start of the binary). The engine_init code adjusts pointers
at load time if needed, or the binary is placed at a fixed address.
"""

import argparse
import json
import struct
import sys
import zlib
from collections import defaultdict
from typing import Dict, List, Optional, Tuple


# --- Stroke parsing ---

STENO_KEYS = {
    'S-': 0x00000001, 'T-': 0x00000002, 'K-': 0x00000004, 'P-': 0x00000008,
    'W-': 0x00000010, 'H-': 0x00000020, 'R-': 0x00000040, 'A-': 0x00000080,
    'O-': 0x00000100, '*':  0x00000200, '-E': 0x00000400, '-U': 0x00000800,
    '-F': 0x00001000, '-R': 0x00002000, '-P': 0x00004000, '-B': 0x00008000,
    '-L': 0x00010000, '-G': 0x00020000, '-T': 0x00040000, '-S': 0x00080000,
    '-D': 0x00100000, '-Z': 0x00200000, '#':  0x00400000,
}

IMPLICIT_HYPHEN = set('AOEU*')
LEFT_BANK = 'STKPWHR'
RIGHT_BANK = 'FRPBLGTSDZ'
VOWELS = 'AOEU'


def parse_stroke(s: str) -> int:
    result = 0
    if '#' in s:
        result |= STENO_KEYS['#']
        s = s.replace('#', '')

    if '-' in s:
        left, right = s.split('-', 1)
        for c in left:
            if c == '*':
                result |= STENO_KEYS['*']
            elif c in 'STKPWHR':
                result |= STENO_KEYS[c + '-']
            elif c in VOWELS:
                result |= STENO_KEYS[{
                    'A': 'A-', 'O': 'O-', 'E': '-E', 'U': '-U'
                }[c]]
        for c in right:
            key = '-' + c
            if key in STENO_KEYS:
                result |= STENO_KEYS[key]
        return result

    # No explicit hyphen — determine split from steno order
    has_vowel_or_star = any(c in IMPLICIT_HYPHEN for c in s)
    if has_vowel_or_star:
        phase = 'left'
        for c in s:
            if phase == 'left':
                if c in IMPLICIT_HYPHEN:
                    phase = 'vowel'
                    if c == '*':
                        result |= STENO_KEYS['*']
                    else:
                        result |= STENO_KEYS[{'A': 'A-', 'O': 'O-', 'E': '-E', 'U': '-U'}[c]]
                elif c in LEFT_BANK:
                    result |= STENO_KEYS[c + '-']
            elif phase == 'vowel':
                if c in IMPLICIT_HYPHEN:
                    if c == '*':
                        result |= STENO_KEYS['*']
                    else:
                        result |= STENO_KEYS[{'A': 'A-', 'O': 'O-', 'E': '-E', 'U': '-U'}[c]]
                else:
                    phase = 'right'
                    result |= STENO_KEYS['-' + c]
            else:
                result |= STENO_KEYS['-' + c]
    else:
        for c in s:
            if c in LEFT_BANK:
                result |= STENO_KEYS[c + '-']
    return result


def parse_outline(outline: str) -> Tuple[int, ...]:
    return tuple(parse_stroke(s) for s in outline.split('/'))


# --- CRC32 (matching Javelin's implementation) ---

def crc32_hash(data: bytes) -> int:
    return zlib.crc32(data) & 0xFFFFFFFF


def stroke_hash(strokes: Tuple[int, ...]) -> int:
    data = b''.join(struct.pack('<I', s) for s in strokes)
    return crc32_hash(data)


# --- Binary builder ---

def uint24(v: int) -> bytes:
    return struct.pack('<I', v)[:3]


def round_up(v: int, align: int) -> int:
    return (v + align - 1) & ~(align - 1)


def popcount32(v: int) -> int:
    return bin(v & 0xFFFFFFFF).count('1')


class CompactMapBuilder:
    """Build a compact map dictionary for a set of entries with the same stroke count."""

    def __init__(self, stroke_length: int, entries: Dict[Tuple[int, ...], int]):
        """entries maps stroke tuples to text block offsets."""
        self.stroke_length = stroke_length
        self.entries = entries

    def build(self) -> Tuple[bytes, bytes, int]:
        """Returns (data_block, hash_blocks, hashMapMask)."""
        if not self.entries:
            return b'', b'', 0

        # Determine hash map size: next power of 2 * 128, at least 2x entries
        n = len(self.entries)
        num_blocks = max(1, (n * 2 + 127) // 128)
        # Round up to power of 2
        p = 1
        while p < num_blocks:
            p *= 2
        num_blocks = p
        total_slots = num_blocks * 128
        hash_map_mask = total_slots - 1

        # Place entries using open addressing (linear probing)
        slots = [None] * total_slots
        for strokes, text_offset in self.entries.items():
            h = stroke_hash(strokes) & hash_map_mask
            while slots[h] is not None:
                h = (h + 1) & hash_map_mask
            slots[h] = (strokes, text_offset)

        # Build data block (entries in slot order) and hash map blocks.
        # Data entry = textOffset(3 bytes) + strokes(3 bytes each)
        entry_size = 3 + 3 * self.stroke_length
        data_block = bytearray()
        hash_blocks = bytearray()

        for block_idx in range(num_blocks):
            masks = [0, 0, 0, 0]
            block_entries = []

            for bit in range(128):
                slot_idx = block_idx * 128 + bit
                if slots[slot_idx] is not None:
                    masks[bit // 32] |= 1 << (bit % 32)
                    block_entries.append(slots[slot_idx])

            # baseOffset: count of all entries in previous blocks,
            # MINUS popcount of current block's masks (Javelin convention).
            # From compact_map_dictionary.cc GetOffset():
            #   result = PopCount(mask << (31 - bitIndex)) + baseOffset
            #   + PopCount of masks[0..maskIndex-1]
            # The builder sets baseOffset = running_total - PopCount(all masks in block)
            running_count = len(data_block) // entry_size
            block_popcount = sum(popcount32(m) for m in masks)
            base_offset = running_count - block_popcount

            for m in masks:
                hash_blocks.extend(struct.pack('<I', m))
            hash_blocks.extend(struct.pack('<I', base_offset & 0xFFFFFFFF))

            for strokes, text_offset in block_entries:
                data_block.extend(uint24(text_offset))
                for s in strokes:
                    data_block.extend(uint24(s))

        return bytes(data_block), bytes(hash_blocks), hash_map_mask


class DictionaryCollectionBuilder:
    """Build a complete JSC4 StenoDictionaryCollection binary.

    Memory layout (all pointers are offsets from binary start):
      [StenoDictionaryCollection header]
      [StenoDictionaryDefinition pointers (one per dict)]
      [StenoCompactMapDictionaryDefinition structs]
      [StenoCompactMapDictionaryStrokesDefinition arrays]
      [Hash map blocks]
      [Data blocks]
      [Text block]
      [Timestamp (4 bytes)]
    """

    def __init__(self, base_addr: int = 0):
        self.base_addr = base_addr

    def build(self, dict_entries: List[Tuple[str, Dict[str, str]]]) -> bytes:
        # Parse all entries, group by stroke count per dictionary
        parsed_dicts = []
        for dict_name, raw_dict in dict_entries:
            by_length = defaultdict(dict)
            max_outline_len = 0
            for outline, translation in raw_dict.items():
                try:
                    strokes = parse_outline(outline)
                except (KeyError, ValueError):
                    continue
                stroke_count = len(strokes)
                by_length[stroke_count][strokes] = translation
                max_outline_len = max(max_outline_len, stroke_count)
            parsed_dicts.append((dict_name, by_length, max_outline_len))

        # Build shared text block (deduplicated)
        all_texts = set()
        for _, by_length, _ in parsed_dicts:
            for entries in by_length.values():
                all_texts.update(entries.values())

        text_list = sorted(all_texts)
        text_to_offset = {}
        text_block = bytearray()
        for text in text_list:
            text_to_offset[text] = len(text_block)
            encoded = text.encode('utf-8')
            text_block.extend(encoded)
            text_block.append(0)

        # Replace translation strings with text offsets
        for _, by_length, _ in parsed_dicts:
            for stroke_len in by_length:
                entries = by_length[stroke_len]
                by_length[stroke_len] = {
                    strokes: text_to_offset[text]
                    for strokes, text in entries.items()
                }

        # Build compact map data for each (dict, stroke_length)
        # Each dict has a StrokesDefinition array indexed by stroke_length
        dict_stroke_data = []
        for dict_name, by_length, max_outline_len in parsed_dicts:
            stroke_defs = []
            for length in range(1, max_outline_len + 1):
                entries = by_length.get(length, {})
                builder = CompactMapBuilder(length, entries)
                data_block, hash_blocks, hash_map_mask = builder.build()
                stroke_defs.append((length, data_block, hash_blocks, hash_map_mask))
            dict_stroke_data.append((dict_name, stroke_defs, max_outline_len))

        # Now lay out the binary. We do two passes:
        # 1. Calculate sizes to determine offsets
        # 2. Write the actual data with correct pointers

        timestamp = 0x12345678

        # Header: StenoDictionaryCollection
        #   uint32_t magic
        #   uint16_t dictionaryCount
        #   bool hasReverseLookup
        #   uint8_t _padding
        #   SizedList<uint8_t> textBlock (count + ptr = 4 + ptr_size)
        #   SizedList<const uint8_t*> prefixes
        #   SizedList<const uint8_t*> suffixes
        #   uint32_t timestamp
        #   XipPointer<StenoDictionaryDefinition> dictionaries[]

        # On ARM32 (nRF52840): pointers are 4 bytes
        PTR = 4

        # SizedList<T> = { size_t count; T* data; } = 4 + PTR
        SIZED_LIST = 4 + PTR

        header_size = (
            4 +       # magic
            2 +       # dictionaryCount
            1 +       # hasReverseLookup
            1 +       # padding
            SIZED_LIST +  # textBlock
            SIZED_LIST +  # prefixes
            SIZED_LIST +  # suffixes
            4         # timestamp
        )
        dict_ptrs_size = len(parsed_dicts) * PTR

        # StenoCompactMapDictionaryDefinition:
        #   uint8_t defaultEnabled, maximumOutlineLength, type, options (= 4)
        #   XipPointer<char> name (= PTR)
        #   const uint8_t* textBlock (= PTR)
        #   const StrokesDefinition* strokes (= PTR)
        COMPACT_DEF_SIZE = 4 + PTR + PTR + PTR

        # StenoCompactMapDictionaryStrokesDefinition:
        #   size_t hashMapMask (= 4)
        #   const uint8_t* data (= PTR)
        #   const Block* offsets (= PTR)
        STROKES_DEF_SIZE = 4 + PTR + PTR

        # Layout plan
        offset = header_size + dict_ptrs_size

        # Align
        offset = round_up(offset, 4)

        # Dictionary definitions
        dict_def_offsets = []
        for i in range(len(parsed_dicts)):
            dict_def_offsets.append(offset)
            offset += COMPACT_DEF_SIZE

        offset = round_up(offset, 4)

        # Dictionary names (null-terminated strings)
        dict_name_offsets = []
        for dict_name, _, _ in dict_stroke_data:
            dict_name_offsets.append(offset)
            offset += len(dict_name.encode('utf-8')) + 1

        offset = round_up(offset, 4)

        # Strokes definitions arrays
        strokes_def_offsets = []  # [(dict_idx, [(length, offset)])]
        for di, (_, stroke_defs, max_len) in enumerate(dict_stroke_data):
            arr_offset = offset
            arr = []
            for si in range(len(stroke_defs)):
                arr.append(offset)
                offset += STROKES_DEF_SIZE
            strokes_def_offsets.append((arr_offset, arr))

        offset = round_up(offset, 4)

        # Data blocks (per dict, per stroke length)
        data_offsets = []  # [[(length, offset)]]
        for di, (_, stroke_defs, _) in enumerate(dict_stroke_data):
            doffs = []
            for si, (length, data_block, hash_blocks, _) in enumerate(stroke_defs):
                doffs.append(offset)
                offset += len(data_block)
            data_offsets.append(doffs)

        offset = round_up(offset, 4)

        # Hash map blocks (per dict, per stroke length)
        hash_offsets = []
        for di, (_, stroke_defs, _) in enumerate(dict_stroke_data):
            hoffs = []
            for si, (length, data_block, hash_blocks, _) in enumerate(stroke_defs):
                hoffs.append(offset)
                offset += len(hash_blocks)
            hash_offsets.append(hoffs)

        offset = round_up(offset, 4)

        # Text block
        text_block_offset = offset
        offset += len(text_block)

        # Timestamp at end of text block
        timestamp_end_offset = offset
        offset += 4

        total_size = offset

        # --- Pass 2: write binary ---

        buf = bytearray(total_size)
        base = self.base_addr

        def write_u8(off, v):
            buf[off] = v & 0xFF

        def write_u16(off, v):
            struct.pack_into('<H', buf, off, v)

        def write_u32(off, v):
            struct.pack_into('<I', buf, off, v & 0xFFFFFFFF)

        def write_ptr(off, v):
            struct.pack_into('<I', buf, off, (base + v) & 0xFFFFFFFF)

        def write_bytes(off, data):
            buf[off:off+len(data)] = data

        def write_sized_list_u8(off, data_off, count):
            write_u32(off, count)
            write_ptr(off + 4, data_off)

        def write_sized_list_ptr(off, data_off, count):
            write_u32(off, count)
            write_ptr(off + 4, data_off)

        # Header
        p = 0
        write_u32(p, 0x3443534A); p += 4  # magic 'JSC4'
        write_u16(p, len(parsed_dicts)); p += 2  # dictionaryCount
        write_u8(p, 0); p += 1  # hasReverseLookup = false
        write_u8(p, 0); p += 1  # padding

        # textBlock SizedList
        write_sized_list_u8(p, text_block_offset, len(text_block)); p += SIZED_LIST

        # prefixes SizedList (empty)
        write_u32(p, 0); write_u32(p + 4, 0); p += SIZED_LIST

        # suffixes SizedList (empty)
        write_u32(p, 0); write_u32(p + 4, 0); p += SIZED_LIST

        # timestamp
        write_u32(p, timestamp); p += 4

        # Dictionary definition pointers
        for di in range(len(parsed_dicts)):
            write_ptr(p, dict_def_offsets[di]); p += PTR

        # Dictionary definitions
        for di, (_, stroke_defs, max_len) in enumerate(dict_stroke_data):
            off = dict_def_offsets[di]
            write_u8(off + 0, 1)  # defaultEnabled = true
            write_u8(off + 1, max_len)  # maximumOutlineLength
            write_u8(off + 2, 0)  # type = COMPACT_MAP
            write_u8(off + 3, 0)  # options
            write_ptr(off + 4, dict_name_offsets[di])  # name
            write_ptr(off + 4 + PTR, text_block_offset)  # textBlock
            # strokes pointer: points to array[0], but Javelin indexes from [1]
            # so we point to (array_start - STROKES_DEF_SIZE)
            arr_start = strokes_def_offsets[di][0]
            write_ptr(off + 4 + PTR + PTR, arr_start - STROKES_DEF_SIZE)

        # Dictionary names
        for di, (dict_name, _, _) in enumerate(dict_stroke_data):
            write_bytes(dict_name_offsets[di], dict_name.encode('utf-8') + b'\x00')

        # Strokes definitions
        for di, (_, stroke_defs, _) in enumerate(dict_stroke_data):
            for si, (length, data_block, hash_blocks, hash_map_mask) in enumerate(stroke_defs):
                off = strokes_def_offsets[di][1][si]
                write_u32(off, hash_map_mask)
                write_ptr(off + 4, data_offsets[di][si])
                write_ptr(off + 4 + PTR, hash_offsets[di][si])

        # Data blocks
        for di, (_, stroke_defs, _) in enumerate(dict_stroke_data):
            for si, (length, data_block, hash_blocks, _) in enumerate(stroke_defs):
                write_bytes(data_offsets[di][si], data_block)

        # Hash map blocks
        for di, (_, stroke_defs, _) in enumerate(dict_stroke_data):
            for si, (length, data_block, hash_blocks, _) in enumerate(stroke_defs):
                write_bytes(hash_offsets[di][si], hash_blocks)

        # Text block
        write_bytes(text_block_offset, text_block)

        # Timestamp at end of text block
        write_u32(timestamp_end_offset, timestamp)

        return bytes(buf)


def main():
    parser = argparse.ArgumentParser(
        description='Compile Plover JSON dictionaries to Javelin JSC4 binary')
    parser.add_argument('inputs', nargs='+', help='Input JSON dictionary files')
    parser.add_argument('-o', '--output', required=True, help='Output binary file')
    parser.add_argument('-n', '--name', action='append',
                        help='Dictionary name (one per input, defaults to filename)')
    parser.add_argument('--base-addr', type=lambda x: int(x, 0), default=0,
                        help='XIP base address for pointer resolution (default: 0)')
    parser.add_argument('--max-entries', type=int, default=0,
                        help='Limit entries per dict (0 = no limit, for testing)')
    args = parser.parse_args()

    dict_entries = []
    for i, input_path in enumerate(args.inputs):
        with open(input_path, 'r') as f:
            raw = json.load(f)
        if args.max_entries > 0:
            items = list(raw.items())[:args.max_entries]
            raw = dict(items)
        name = args.name[i] if args.name and i < len(args.name) else input_path.rsplit('/', 1)[-1].rsplit('.', 1)[0]
        dict_entries.append((name, raw))
        print(f"Loaded {input_path}: {len(raw)} entries as '{name}'", file=sys.stderr)

    builder = DictionaryCollectionBuilder(base_addr=args.base_addr)
    result = builder.build(dict_entries)

    with open(args.output, 'wb') as f:
        f.write(result)

    print(f"Wrote {len(result)} bytes ({len(result)/1024:.1f} KB) to {args.output}", file=sys.stderr)


if __name__ == '__main__':
    main()