"""Losslessly convert CrabKeyboard's V5 trie into bounded, digit-indexed shards."""
import argparse, hashlib, json, pathlib, re

APP = pathlib.Path(__file__).resolve().parents[1]
GROUPS = ['abc', 'def', 'ghi', 'jkl', 'mno', 'pqrs', 'tuv', 'wxyz']
T9 = {letter: str(i + 2) for i, group in enumerate(GROUPS) for letter in group}
GROUPS14 = ['qw','er','ty','ui','op','as','df','gh','jk','l','zx','cv','bn','m']
KEY14 = {letter: '0123456789abcd'[i] for i, group in enumerate(GROUPS14) for letter in group}

def from94(value):
    number = 0
    for ch in value:
        assert 33 <= ord(ch) <= 126
        number = number * 94 + ord(ch) - 33
    return number

def convert(source, target):
    files = {p.name: p.read_bytes() for p in source.glob('*.dat')}
    frequency_path=APP/'third_party/jieba/dict.txt'
    frequencies={}
    if frequency_path.exists():
        for line in frequency_path.read_text(encoding='utf-8').splitlines():
            parts=line.split(' ')
            if len(parts)>=2 and parts[1].isdigit():frequencies[parts[0]]=int(parts[1])
    entries = {}
    def walk(key, depth, position, length):
        data = files['key%d.dat' % depth][position:position+length].decode('ascii')
        for item in data.split('\x1e'):
            if not item: continue
            char = item[0]
            off, size = map(from94, item[1:].split('\x1f'))
            if char == '_':
                words = files['value%d.dat' % depth][off:off+size].decode('utf-8').split(' ')
                entries[key] = [word for word in words if word]
            else:
                walk(key+char, depth+1, off, size)
    walk('', 0, 0, len(files['key0.dat']))
    shards = {}
    for key, words in sorted(entries.items()):
        digits = ''.join(T9[ch] for ch in key)
        bucket = digits[:4] if len(digits) >= 4 else 's'+digits
        shards.setdefault(bucket, {})[key] = words
        code14 = ''.join(KEY14[ch] for ch in key)
        bucket14 = 'd'+code14[:3] if len(code14)>=3 else 'ds'+code14
        shards.setdefault(bucket14,{})[key]=words
    target.mkdir(parents=True, exist_ok=True)
    assert target.resolve().is_relative_to((APP/'src/assets').resolve())
    # Only replace files listed in the previous generated manifest.
    if (target/'meta.json').exists():
        previous=json.loads((target/'meta.json').read_text(encoding='utf-8'))
        for name in previous['shards']:
            assert re.fullmatch(r'(?:short|s?[0-9]+|ds?[0-9a-d]+)',name)
            old=target/(name+'.json')
            if name not in shards and old.exists():old.unlink()
    for name, content in shards.items():
        weights={word:frequencies[word] for words in content.values() for word in words if word in frequencies}
        output=target/(name+'.json')
        data=json.dumps({'entries':content,'frequencies':weights}, ensure_ascii=False, separators=(',', ':')).encode('utf-8')
        if not output.exists() or output.read_bytes()!=data:output.write_bytes(data)
    # Every original candidate remains accessible, in original per-key order.
    restored = {};restored14={}
    for name in shards:
        (restored14 if name.startswith('d')else restored).update(json.loads((target/(name+'.json')).read_text(encoding='utf-8'))['entries'])
    assert restored == entries and restored14 == entries
    report = dict(upstream='https://github.com/crabPangXie/crabKeyboard',
                  commit='93d36b5002dee80a3d253b1ec7ec9e29800a104c',
                  variant='pinyinDictV5_20260416', keys=len(entries),
                  candidates=sum(map(len, entries.values())), shards=sorted(shards),
                  t9_shards=sum(not name.startswith('d')for name in shards),
                  key14_shards=sum(name.startswith('d')for name in shards),key14_groups=GROUPS14,key14_prefix_length=3,
                  largest_shard_bytes=max((target/(name+'.json')).stat().st_size for name in shards),
                  source_files={name:hashlib.sha256(data).hexdigest() for name,data in sorted(files.items())})
    if frequency_path.exists():report['t9_frequency_source']={'project':'jieba 0.42.1','license':'MIT','sha256':hashlib.sha256(frequency_path.read_bytes()).hexdigest()}
    output=target/'meta.json';data=(json.dumps(report, ensure_ascii=False, indent=2)+'\n').encode('utf-8')
    if not output.exists() or output.read_bytes()!=data:output.write_bytes(data)
    print(json.dumps({**{k:v for k,v in report.items() if k not in ['source_files','shards']},'shards':len(shards)}, ensure_ascii=False))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--source', type=pathlib.Path, default=APP/'third_party/crabKeyboard/pinyinDictV5_20260416')
    parser.add_argument('--target', type=pathlib.Path, default=APP/'src/assets/ime-dict')
    args = parser.parse_args(); convert(args.source, args.target)