#!/usr/bin/env python3

import argparse
import typing
import unicodedata
from pathlib import Path
from xml.etree import ElementTree as ET

import tqdm

import bunkai.constant


def get_sentence(node) -> typing.Iterator[str]:
    if node.tag == 'sentence':
        yield bunkai.constant.METACHAR_SENTENCE_BOUNDARY
    for luw in node:
        if luw.tag == 'webBr':
            yield bunkai.constant.METACHAR_LINE_BREAK
            continue
        elif luw.tag.lower() != 'luw':
            yield from get_sentence(luw)
            continue

        for fragment in luw:
            if fragment.tag in ['webBr', 'br']:
                yield bunkai.constant.METACHAR_LINE_BREAK
                continue
            elif fragment.tag.startswith('note'):
                continue
            elif fragment.tag in ['sampling', 'fraction', 'info']:
                continue
            elif fragment.tag.lower() != 'suw':
                raise NotImplementedError(fragment.tag)
            elif fragment.text is None:
                for unit in fragment:
                    if unit.tag == 'sampling':
                        pass
                    elif unit.tag == 'correction':
                        yield unit.get('originalText')
                    elif unit.tag not in ['ruby', 'quote', 'enclosedCharacter', 'delete', 'subScript', 'superScript']:
                        raise NotImplementedError(unit.tag)
                    else:
                        yield unit.text
                continue
            yield fragment.text


def operate_article(parent) -> typing.Iterator[str]:
    for node in parent:
        if node.tag in ['rejectedBlock', 'info', 'titleBlock', 'abstract', 'authorsData', 'figureBlock']:
            continue
        elif node.tag in ['webBr', 'br']:
            yield bunkai.constant.METACHAR_LINE_BREAK
        elif node.tag == 'sentence':
            yield from get_sentence(node)
        else:
            yield from operate_article(node)


def operation(data: str) -> typing.Iterator[str]:
    root = ET.fromstring(data)
    for child in root:
        if child.tag != 'article':
            continue
        for i, t in enumerate(operate_article(child)):
            if i == 0:
                yield t.lstrip(bunkai.constant.METACHAR_SENTENCE_BOUNDARY)
            else:
                yield t
        yield '\n'


def get_opts() -> argparse.Namespace:
    oparser = argparse.ArgumentParser()
    oparser.add_argument("--input", "-i", type=Path, required=True)
    oparser.add_argument("--output", "-o", type=Path, required=True)
    oparser.add_argument("--nonfkc", action='store_true')
    return oparser.parse_args()


def main() -> None:
    opts = get_opts()
    if opts.input.is_dir():
        opts.output.mkdir(exist_ok=True, parents=True)
        targets = [n for n in opts.input.iterdir()]
        for fname in tqdm.tqdm(targets, leave=False):
            with fname.open() as inf,\
                    opts.output.joinpath(f'{fname.stem}.txt').open('w') as outf:
                for string in operation(inf.read()):
                    if not opts.nonfkc:
                        string = unicodedata.normalize("NFKC", string)
                    outf.write(string)
    else:
        with opts.input.open() as inf,\
                opts.output.open('w') as outf:
            for string in operation(inf.read()):
                if not opts.nonfkc:
                    string = unicodedata.normalize("NFKC", string)
                outf.write(string)


if __name__ == '__main__':
    main()
