"""Download and build separately licensed, immutable dictionary shards. No model calls."""
import gzip
import hashlib
import json
import re
import shutil
import urllib.request
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
CACHE = ROOT / '.cache'
PUBLIC = ROOT / 'public'
CACHE.mkdir(exist_ok=True)


def download(name, url):
    folder = ROOT / 'data' / ('sources' if name in ['cedict.txt.gz', 'makemeahanzi.txt'] else 'licenses')
    folder.mkdir(parents=True, exist_ok=True)
    path = folder / name
    if not path.exists() and (CACHE / name).exists():
        shutil.copy2(CACHE / name, path)
    if not path.exists():
        print('Downloading', name, flush=True)
        with urllib.request.urlopen(url, timeout=90) as response:
            payload = response.read(35_000_000)
        tmp = path.with_suffix('.tmp')
        tmp.write_bytes(payload)
        tmp.replace(path)
    return path


cedict_url = 'https://www.mdbg.net/chinese/export/cedict/cedict_1_0_ts_utf-8_mdbg.txt.gz'
cedict = download('cedict.txt.gz', cedict_url)
lockpath = ROOT / 'sources.lock.json'
lock = json.loads(lockpath.read_text()) if lockpath.exists() else {}
revision = lock.get('makemeahanzi', {}).get('revision')
if not revision:
    req = urllib.request.Request('https://api.github.com/repos/skishore/makemeahanzi/commits/master', headers={'User-Agent': 'zh-dictionary-data-builder'})
    with urllib.request.urlopen(req, timeout=45) as response:
        revision = json.load(response)['sha']
mmh_url = f'https://raw.githubusercontent.com/skishore/makemeahanzi/{revision}/dictionary.txt'
mmh = download('makemeahanzi.txt', mmh_url)
for name, url in [('LGPL-3.0.txt', f'https://raw.githubusercontent.com/skishore/makemeahanzi/{revision}/LGPL'),
                  ('COPYING-makemeahanzi.txt', f'https://raw.githubusercontent.com/skishore/makemeahanzi/{revision}/COPYING'),
                  ('GPL-3.0.txt', 'https://www.gnu.org/licenses/gpl-3.0.txt'),
                  ('CC-BY-SA-4.0.html', 'https://raw.githubusercontent.com/creativecommons/cc-legal-tools-data/main/docs/licenses/by-sa/4.0/legalcode.en.html')]:
    download(name, url)

for name, path in [('cedict', cedict), ('makemeahanzi', mmh)]:
    digest = hashlib.sha256(path.read_bytes()).hexdigest()
    if name in lock and lock[name]['sha256'] != digest:
        raise RuntimeError(f'{name} checksum differs from sources.lock.json. Restore the original source or explicitly refresh the lock.')
    lock[name] = {'sha256': digest, 'url': cedict_url if name == 'cedict' else mmh_url, **({'revision': revision} if name == 'makemeahanzi' else {})}
lockpath.write_text(json.dumps(lock, ensure_ascii=False, indent=2) + '\n')

words = defaultdict(list)
count = 0
for line in gzip.decompress(cedict.read_bytes()).decode('utf-8-sig').splitlines():
    match = re.match(r'^(\S+) (\S+) \[([^\]]+)\] /(.+)/$', line)
    if not match:
        continue
    trad, simp, pinyin, definitions = match.groups()
    entry = {'simplified': simp, 'traditional': trad, 'pinyin': pinyin, 'definitions': definitions.split('/')}
    for form in sorted({trad, simp}):
        words[form].append(entry)
    count += 1
assert count > 100_000, 'Incomplete CC-CEDICT source'
characters = {row['character']: row for row in map(json.loads, mmh.read_text().splitlines())}
assert len(characters) > 8000, 'Incomplete character source'
related = defaultdict(list)
for word in sorted(words, key=lambda value: (len(value), value)):
    if len(word) > 1:
        for char in set(word):
            if '\u3400' <= char <= '\u9fff' or ord(char) >= 0x20000:
                related[char].append(word)

version = hashlib.sha256((lock['cedict']['sha256'] + lock['makemeahanzi']['sha256'] + 'format-2').encode()).hexdigest()[:12]
out = PUBLIC / 'data' / version
out.mkdir(parents=True, exist_ok=True)
for kind, values in [('words', words), ('characters', characters), ('related', related)]:
    buckets = [dict() for _ in range(256)]
    for key, value in values.items():
        buckets[ord(key[0]) % 256][key] = value
    for index, bucket in enumerate(buckets):
        (out / f'{kind}-{index:02x}.json').write_text(json.dumps(bucket, ensure_ascii=False, sort_keys=True, separators=(',', ':')))
manifest = {'version': version, 'format': 2, 'entries': count, 'headwords': len(words), 'characters': len(characters), 'sources': lock}
(PUBLIC / 'data' / 'manifest.json').write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + '\n')
sources = PUBLIC / 'sources'
sources.mkdir(exist_ok=True)
for filename in ['cedict.txt.gz', 'makemeahanzi.txt', 'LGPL-3.0.txt', 'GPL-3.0.txt', 'CC-BY-SA-4.0.html', 'COPYING-makemeahanzi.txt']:
    source = ROOT / 'data' / ('sources' if filename in ['cedict.txt.gz', 'makemeahanzi.txt'] else 'licenses') / filename
    shutil.copy2(source, sources / filename)
shutil.copy2(__file__, sources / 'build-data.py')
shutil.copy2(lockpath, sources / lockpath.name)
print(json.dumps({key: manifest[key] for key in ['version', 'entries', 'headwords', 'characters']}))