import os.path
import jsonpickle
import ujson as json
import pandas as pd
from pathlib import Path
from sys import argv
from tqdm import tqdm
from argparse import ArgumentParser, RawDescriptionHelpFormatter

from eis1600.repositories.repo import get_ready_and_double_checked_files, TEXT_REPO, JSON_REPO, TSV_REPO, COLUMNS, \
    SEP, SEP2
from eis1600.helper.CheckFileEndingActions import CheckFileEndingEIS1600JsonAction


ALL_LABELS = ("SECTIONS", "TOKENS", "TAGS_LISTS", "NER_LABELS", "LEMMAS", "POS_TAGS", "ROOTS", "TOPONYM_LABELS",
              "NER_TAGS", "DATE_TAGS", "MONTH_TAGS", "ONOM_TAGS", "ONOMASTIC_TAGS")


def dump_file(fpath: str, label_list: tuple[str] = ALL_LABELS):

    if not fpath.endswith(".json"):
        fpath = fpath.replace(TEXT_REPO, JSON_REPO)
        fpath = fpath.replace('.EIS1600', '.json')

    structural_data, content_data = [], []
    with open(fpath, "r", encoding="utf-8") as fp:
        data = json.load(fp)

    for miu in data:
        header = miu["yml"]
        uid = header["UID"]

        # get structural data
        for entity in header.keys():
            if entity == "UID":
                continue
            value = header[entity]
            if type(value) in (str, int, bool):
                structural_data.append((uid, entity, value))
            elif type(value) == list:
                for sub_value in value:
                    if type(sub_value) == dict:
                        parsed_sub_value = SEP.join(f"{k}{SEP2}{v}" for k, v in sub_value.items())
                        structural_data.append((uid, entity, parsed_sub_value))
                    elif type(sub_value) == list:
                        parsed_sub_value = SEP.join(sub_value)
                        structural_data.append((uid, entity, parsed_sub_value))
                    else:
                        structural_data.append((uid, entity, sub_value))
            elif type(value) == dict:
                for sub_entity, sub_value in value.items():
                    if type(sub_value) == list:
                        for val in sub_value:
                            structural_data.append((uid, entity, f"{sub_entity}{SEP}{val}"))
                    elif type(sub_value) == dict:
                        # e.g. onomastics elements
                        for sub_sub_ent, sub_sub_val in sub_value.items():
                            structural_data.append((uid, entity, f"{sub_entity}{SEP}{sub_sub_ent}{SEP2}"))
                    else:
                        structural_data.append((uid, entity, f"{sub_entity}{SEP}{sub_value}"))
            else:
                raise ValueError(f'Fatal error dumping data of "{fpath}".')

        # get content data
        miu_df = jsonpickle.decode(miu["df"])
        for entity in label_list:
            if entity not in miu_df:
                continue
            for j, (_, value) in enumerate(miu_df[entity].items(), 1):
                if type(value) == list:
                    value = SEP2.join(value)
                if value:
                    content_data.append((uid, entity, f"{j}{SEP}{value}"))

    fbase, _ = os.path.splitext(fpath)

    fbase = fbase.replace(JSON_REPO, TSV_REPO)
    dir_path, _ = os.path.split(fbase)
    Path(dir_path).mkdir(parents=True, exist_ok=True)

    content_df = pd.DataFrame(content_data, columns=COLUMNS)
    content_df.to_csv(f"{fbase}_df.tsv", sep="\t", index=False)

    struct_df = pd.DataFrame(structural_data, columns=COLUMNS)
    struct_df.to_csv(f"{fbase}_yml.tsv", sep="\t", index=False)


def main():
    arg_parser = ArgumentParser(
            prog=argv[0], formatter_class=RawDescriptionHelpFormatter,
            description="Script to dump all eis1600 corpus into a tsv with the structure and "
                        "another tsv file with the enriched textual information."
    )
    arg_parser.add_argument(
            'infile', type=str, nargs='?',
            help='json file to process',
            action=CheckFileEndingEIS1600JsonAction
    )
    arg_parser.add_argument(
        "--label_list",
        nargs="*",
        default=ALL_LABELS,
        help="entities from content data to add to output. The default is all entities: "
             "SECTIONS, TOKENS, TAGS_LISTS, NER_LABELS, LEMMAS, POS_TAGS, ROOTS, TOPONYM_LABELS, "
             "NER_TAGS, DATE_TAGS, MONTH_TAGS, ONOM_TAGS, ONOMASTIC_TAGS"
    )
    args = arg_parser.parse_args()

    if args.infile:
        dump_file(args.infile, args.label_list)

    else:
        files_ready, files_double_checked = get_ready_and_double_checked_files()
        files = files_ready + files_double_checked

        for fpath in tqdm(files):

            dump_file(fpath, args.label_list)

        print(f"Processed {len(files)} files")
        print(f"For each json file in {JSON_REPO} directory, "
              f"a tsv with the yml data and a tsv with the df data have been generated.")

