#!/usr/bin/env python3
import sys

normal_first_letters = {
    'ā': 'a',
    'á': 'a',
    'à': 'a',
    'ǎ': 'a',
    'ế': 'e',
    'é': 'e',
    'ê̌': 'e',
    'ě': 'e',
    'ề': 'e',
    'è': 'e',
    'ê̄': 'e',
    'ē': 'e',
    'ō': 'o',
    'ň': 'n',
    'ǹ': 'n', 'ḿ': 'm',
    'ń': 'n',
    'ó': 'o',
    'ê': 'e',
    'ò': 'o',
    'ǒ': 'o',
}


def normal_first_letter(c):
    return normal_first_letters.get(c, c)


class Record:
    def __init__(self, l):
        words = l.split(':', 2)
        self.ord = int(words[0][2:], 16)

        words = words[1].split('#', 2)
        self.pinyin = words[0].strip().split(',')
        self.ch = words[1].strip()

        first_letters = []
        for w in self.pinyin:
            # do not use hash/set, keeps order
            c = normal_first_letter(w[0])
            if c not in first_letters:
                first_letters.append(c)
        assert chr(
            self.ord) == self.ch, f'{chr(self.ord)} {self.ch} {self.ord}'
        self.first_letters = first_letters

    def __str__(self):
        return f'{hex(self.ord)} {self.pinyin} {self.ch} {self.first_letters}'

    def encoded_tones(self):
        bits = 0
        first = True
        for ch in list(reversed(self.first_letters[:3])):
            if not first:
                bits <<= 5
            first = False
            bits += ord(ch) - ord('a') + 1
        return bits


first_letter_set = set()
max_multi_tones = Record('U+3405: wǔ  # 㐅')
more_than_three_tones = []
counts = 0
recs = []
for line in open('pinyin.txt'):
    if line.startswith('#'):
        continue

    line = line.rstrip('\n')
    rec = Record(line)
    if rec.ord > 0xffff:
        # Ignore char > 0xffff, they are uncommon chars, and can not hold inside uint16
        continue

    recs.append(rec)
    counts += 1
    first_letter_set.update(rec.first_letters)
    if len(max_multi_tones.first_letters) < len(rec.first_letters):
        max_multi_tones = rec
    if len(rec.first_letters) > 3:
        more_than_three_tones.append(rec)

# Char multi tones more than 3 are very little, and are uncommon chars.
# Keep 3 tones good enough

# print()
# print(f'Total {counts}, Max multi tones: {max_multi_tones}')
# print(f'{sorted(first_letter_set)} {len(first_letter_set)}')
# print(f'{len(more_than_three_tones)}')

# for x in more_than_three_tones:
#     print(x)


def gen_letter():
    print(r'''// Code generated by gen.py, DO NOT EDIT.

    package pinyin

    var records = []record{''')

    for r in recs:
        print(
            f'    {{CodePoint: {hex(r.ord)}'
            f', Tones: {r.encoded_tones()}}}, // {r.ch}')

    print('}')


def build_tree():
    root = []

    def find_or_create(s, pinyin):
        ns = root
        for i, ch in enumerate(s):
            for n in ns:
                if n['ch'] == ch:
                    return ns
            else:
                isLast = i + 1 == len(s)
                n = {
                    'ch': ch,
                    'children': None if isLast else [],
                    'pinyin': pinyin if isLast else None,
                }
                ns.append(n)
            ns = n['children']
        return n

    for line in open('words.dat'):
        pair = line.split(',')
        if len(pair) != 2:
            continue
        find_or_create(pair[0], pair[1])

    return {
        'ch': '',
        'children': root,
    }


def breadth_first_walk(node):
    if node['children'] is None:
        yield node['ch'], node['pinyin']
    else:
        yield node['ch'], node['children']
        for n in node['children']:
            yield from breadth_first_walk(n)


def gen_words():
    tree = build_tree()

    for ch, item in breadth_first_walk(tree):
        if isinstance(item, list):
            print(f'{ch},{len(item)}')
        else:
            print(f'{ch},0,{item}')


if len(sys.argv) == 1 or (sys.argv[1] not in ['--letter', '--words']):
    print('Usage:', sys.argv[0], '[--letter|--words]')
elif sys.argv[1] == '--letter':
    gen_letter()
else:
    gen_words()
