"""Losslessly build the Android SQLite asset from the licensed original V5 trie."""
import argparse,hashlib,importlib.util,json,pathlib,sqlite3
APP=pathlib.Path(__file__).resolve().parents[1]

def build(rtos):
    source=rtos/'third_party/crabKeyboard/pinyinDictV5_20260416'
    files={p.name:p.read_bytes()for p in source.glob('*.dat')};entries={}
    assert len(files)==41,'Expected the complete upstream V5 dictionary'
    def number(text):
        result=0
        for ch in text:assert 33<=ord(ch)<=126;result=result*94+ord(ch)-33
        return result
    def walk(key,depth,position,length):
        for item in files[f'key{depth}.dat'][position:position+length].decode('ascii').split('\x1e'):
            if not item:continue
            offset,size=map(number,item[1:].split('\x1f'))
            if item[0]=='_':entries[key]=[word for word in files[f'value{depth}.dat'][offset:offset+size].decode('utf-8').split(' ')if word]
            else:walk(key+item[0],depth+1,offset,size)
    walk('',0,0,len(files['key0.dat']))
    frequency_path=rtos/'third_party/jieba/dict.txt'
    frequencies={}
    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])
    spec=importlib.util.spec_from_file_location('crab_dictionary',rtos/'scripts/convert-dictionary.py')
    converter=importlib.util.module_from_spec(spec);spec.loader.exec_module(converter)
    associations=converter.associations(entries,frequencies)
    groups9=['abc','def','ghi','jkl','mno','pqrs','tuv','wxyz'];groups14=['qw','er','ty','ui','op','as','df','gh','jk','l','zx','cv','bn','m']
    map9={ch:str(i+2)for i,group in enumerate(groups9)for ch in group}
    map14={ch:'0123456789abcd'[i]for i,group in enumerate(groups14)for ch in group}
    target=APP/'keyboard/src/main/assets/dictionary.db';target.parent.mkdir(parents=True,exist_ok=True)
    temporary=target.with_suffix('.tmp')
    if temporary.exists():temporary.unlink()
    db=sqlite3.connect(temporary)
    db.executescript('PRAGMA page_size=4096; PRAGMA journal_mode=OFF; CREATE TABLE entries(id INTEGER PRIMARY KEY,pinyin TEXT NOT NULL UNIQUE,code9 TEXT NOT NULL,code14 TEXT NOT NULL); CREATE TABLE words(entry INTEGER NOT NULL,ordinal INTEGER NOT NULL,word TEXT NOT NULL,frequency INTEGER NOT NULL,PRIMARY KEY(entry,ordinal)) WITHOUT ROWID;')
    for index,(key,words)in enumerate(sorted(entries.items()),1):
        db.execute('INSERT INTO entries VALUES(?,?,?,?)',(index,key,''.join(map9[ch]for ch in key),''.join(map14[ch]for ch in key)))
        db.executemany('INSERT INTO words VALUES(?,?,?,?)',[(index,ordinal,word,frequencies.get(word,0))for ordinal,word in enumerate(words)])
    db.execute('CREATE TABLE associations(prefix TEXT NOT NULL,word TEXT NOT NULL,frequency INTEGER NOT NULL,PRIMARY KEY(prefix,word)) WITHOUT ROWID')
    db.executemany('INSERT INTO associations VALUES(?,?,?)',[(prefix,word,frequency)for prefix,rows in associations.items()for word,frequency in rows])
    db.executescript('CREATE INDEX entries_code9 ON entries(code9);CREATE INDEX entries_code14 ON entries(code14);ANALYZE;PRAGMA user_version=2;')
    db.commit()
    restored={}
    for key,word in db.execute('SELECT e.pinyin,w.word FROM entries e JOIN words w ON w.entry=e.id ORDER BY e.pinyin,w.ordinal'):restored.setdefault(key,[]).append(word)
    assert restored==entries,'Full dictionary round-trip mismatch'
    assert db.execute('PRAGMA integrity_check').fetchone()[0]=='ok';db.close();temporary.replace(target)
    report={'schema_version':2,'association_prefixes':len(associations),'association_candidates':sum(map(len,associations.values())),'keys':len(entries),'candidates':sum(map(len,entries.values())),'longest_word':max(len(word)for words in entries.values()for word in words),'sha256':hashlib.sha256(target.read_bytes()).hexdigest(),'bytes':target.stat().st_size,'round_trip':True,'source_files':{name:hashlib.sha256(data).hexdigest()for name,data in sorted(files.items())},'jieba_sha256':hashlib.sha256(frequency_path.read_bytes()).hexdigest()}
    (APP/'keyboard/src/main/assets/dictionary-manifest.json').write_text(json.dumps(report,ensure_ascii=False,indent=2)+'\n',encoding='utf-8')
    print(json.dumps({k:v for k,v in report.items()if k!='source_files'},ensure_ascii=False))

if __name__=='__main__':
    parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--rtos-source',type=pathlib.Path)
    args=parser.parse_args();rtos=args.rtos_source
    if rtos is None:
        rtos=next((p for p in [APP.parent/'com.sydxky.rtos.inputmethod',APP.parent/'RTOS/com.sydxky.rtos.inputmethod']if p.exists()),None)
    assert rtos is not None,'Provide the licensed RTOS source with --rtos-source';build(rtos)