"""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 predictive_bucket(value):
    result=0
    for ch in value:result=(result*31+ord(ch))&0xffff
    return format(result%256,'02x')

def associations(entries,frequencies):
    """Derive bounded word continuations from licensed phrases; no user corpus."""
    result={}
    for word in {word for words in entries.values()for word in words}:
        frequency=frequencies.get(word,0)
        if frequency<20 or not re.fullmatch(r'[\u3400-\u9fff]{2,12}',word):continue
        for cut in range(1,min(4,len(word)-1)+1):
            prefix,suffix=word[:cut],word[cut:]
            result.setdefault(prefix,{})[suffix]=round(frequency/(1+.25*len(suffix)))
    # Small, original conversational defaults supplement dictionary word splits.
    seeds=json.loads((APP/'src/common/association-seeds.json').read_text(encoding='utf-8'))
    for prefix,words in seeds.items():
        for rank,word in enumerate(words):result.setdefault(prefix,{})[word]=200000-rank*10000
    return {prefix:sorted(words.items(),key=lambda x:(-x[1],x[0]))[:8]for prefix,words in sorted(result.items())}

def completions(entries,frequencies):
    result={}
    for key,words in entries.items():
        for mode,code in [('py',key),('t9',''.join(T9[ch]for ch in key)),('14',''.join(KEY14[ch]for ch in key))]:
            for length in range(2,min(3 if mode!='14'else 2,len(code)-1)+1):
                if len(code)-length>4:continue
                target=result.setdefault(mode+':'+code[:length],{})
                for ordinal,word in enumerate(words):
                    frequency=frequencies.get(word,0)
                    if frequency<20:continue
                    value=[word,key,frequency,ordinal]
                    if word not in target or target[word][2]<frequency:target[word]=value
    return {prefix:sorted(words.values(),key=lambda x:(-x[2],len(x[1]),x[3],x[0]))[:12]for prefix,words in sorted(result.items())}

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

def presence(names,alphabet,lengths,prefix):
    masks=[]
    for length in lengths:
        mask=bytearray((len(alphabet)**length+7)//8)
        for name in names:
            code=name[len(prefix):] if name.startswith(prefix) else ''
            if len(code)!=length or any(ch not in alphabet for ch in code):continue
            index=0
            for ch in code:index=index*len(alphabet)+alphabet.index(ch)
            mask[index//8]|=1<<(index%8)
        masks.append(mask.hex())
    return masks

def convert(source, target):
    target.parent.mkdir(parents=True, exist_ok=True)
    assert target.resolve().is_relative_to((APP/'src/assets').resolve())
    marker=target.parent/('.'+target.name+'.generating')
    with marker.open('x',encoding='ascii') as handle:handle.write('generating')
    try:
        _convert(source,target)
    finally:
        marker.unlink()

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'],*previous.get('association_shards',[]),*previous.get('completion_shards',[])]:
            assert re.fullmatch(r'(?:short|s?[0-9]+|ds?[0-9a-d]+|[ac][0-9a-f]{2})',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)
    association_shards={};completion_shards={}
    for prefix,items in associations(entries,frequencies).items():association_shards.setdefault('a'+predictive_bucket(prefix),{})[prefix]=items
    for prefix,items in completions(entries,frequencies).items():completion_shards.setdefault('c'+predictive_bucket(prefix),{})[prefix]=items
    for name,content in {**association_shards,**completion_shards}.items():
        (target/(name+'.json')).write_text(json.dumps(content,ensure_ascii=False,separators=(',',':')),encoding='utf-8',newline='\n')
    # 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(association_seed_sha256=hashlib.sha256((APP/'src/common/association-seeds.json').read_bytes()).hexdigest(),
                  generator_sha256=hashlib.sha256(pathlib.Path(__file__).read_bytes()).hexdigest(),
                  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),
                  association_shards=sorted(association_shards),completion_shards=sorted(completion_shards),
                  predictive_largest_shard_bytes=max((target/(name+'.json')).stat().st_size for name in [*association_shards,*completion_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()}
    # Runtime needs only presence bits, not thousands of filenames and audit data.
    runtime={'version':1,'short':presence(shards,'23456789',range(1,4),'s'),
             'long':presence([name for name in shards if re.fullmatch('[2-9]{4}',name)],'23456789',[4],''),
             'dualShort':presence(shards,'0123456789abcd',range(1,3),'ds'),
             'dualLong':presence([name for name in shards if re.fullmatch('d[0-9a-d]{3}',name)],'0123456789abcd',[3],'d'),
             'associations':presence(association_shards,'0123456789abcdef',[2],'a'),
             'completions':presence(completion_shards,'0123456789abcdef',[2],'c')}
    output=target/'runtime.json';data=(json.dumps(runtime,separators=(',',':'))+'\n').encode('utf-8')
    if not output.exists() or output.read_bytes()!=data:output.write_bytes(data)
    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)