import csv
import errno
import io
import json
import os
import traceback
from contextlib import suppress

from functools import partial
from glob import glob
from multiprocessing.pool import Pool


def translate_dir(dir_in, tsv_out_dir, json_out_dir, is_global):
    files = [
        f
        for f in glob(os.path.join(dir_in, "**/Scenario/*.tsv"), recursive=True)
        if f"{os.path.sep}Diva{os.path.sep}" not in f
    ]

    if len(files) == 0:
        raise FileNotFoundError("No valid files found in directory: " + dir_in)

    ptrans = partial(
        translate_file,
        tsv_out_dir=tsv_out_dir,
        json_out_dir=json_out_dir,
        is_global=is_global,
    )
    with Pool() as p:
        p.map(ptrans, files)


def translate_file(file_in, tsv_out_dir, json_out_dir, is_global):
    try:
        if not file_in.endswith(".tsv"):
            raise ValueError("Invalid File Type for {}".format(file_in))

        # we need to get the event folder
        # ie, Utage/>>>main01<<</Scenario/whatever.tsv
        try:
            event_folder = os.path.normpath(file_in).split(os.sep)[-3]
        except Exception:
            traceback.print_exc()
            print("Error processing tsv: Please reference from Utage root")
            raise

        id_num = os.path.splitext(os.path.basename(file_in))[0]

        tsv_output_path = os.path.join(
            tsv_out_dir, event_folder, "Scenario", "{}_t.tsv".format(id_num)
        )
        languages = ["jpn"]
        if is_global:
            languages = ["enm", "zho", "kor"]
        json_output_paths = {
            lang: os.path.join(
                json_out_dir,
                event_folder,
                "{}_translations_{}.json".format(id_num, lang),
            )
            for lang in languages
        }

        # need to create output paths and avoid races when threading
        if not os.path.exists(os.path.dirname(tsv_output_path)):
            try:
                os.makedirs(os.path.dirname(tsv_output_path))
            except OSError as e:
                if e.errno != errno.EEXIST:
                    raise

        if not os.path.exists(os.path.dirname(json_output_paths[languages[0]])):
            try:
                os.makedirs(os.path.dirname(json_output_paths[languages[0]]))
            except OSError as e:
                if e.errno != errno.EEXIST:
                    raise

        with open(file_in, "r", encoding="utf-8") as tsv_file:
            tsv = csv.DictReader(tsv_file, delimiter="\t", quoting=csv.QUOTE_NONE)

            t_fieldnames = tsv.fieldnames
            if is_global:
                t_fieldnames = tsv.fieldnames[: tsv.fieldnames.index("English") + 1]

            tsv_keyed, json_str = process_tsv(tsv, id_num, is_global)

            # csv handles newlines, don't set it in io.open
            with io.open(tsv_output_path, "w", newline="", encoding="utf-8") as tsv_out:
                writer = csv.DictWriter(
                    tsv_out,
                    delimiter="\t",
                    quotechar='"',
                    fieldnames=t_fieldnames,
                    extrasaction="ignore",
                )
                writer.writeheader()
                for row in tsv_keyed:
                    writer.writerow(row)

        for lang in languages:
            if not json_str[lang]:
                continue
            with io.open(json_output_paths[lang], "w", encoding="utf-8") as json_out:
                # we don't want to sort these because they're _1, ..., _10, etc
                json.dump(
                    json_str[lang],
                    json_out,
                    ensure_ascii=False,
                    indent="\t",
                    sort_keys=False,
                )

    except Exception:
        traceback.print_exc()
        print("Error processing file: " + file_in)


def process_tsv(tsv, id_num, is_global):
    i = 0
    tsv_keyed = []
    key_dict = {"jpn": {}}
    if is_global:
        key_dict["enm"] = {}
        key_dict["zho"] = {}
        key_dict["kor"] = {}

    for row in tsv:
        try:
            if row["Command"].startswith("//"):
                tsv_keyed.append(row)
                continue
            if row["Text"] and row["Text"].strip():
                key = "{}_{}".format(id_num, i)
                i += 1
                key_dict["jpn"][key] = row["Text"]
                if is_global:
                    with suppress(KeyError):
                        if row["English"]:
                            key_dict["enm"][key] = row["English"]
                        key_dict["zho"][key] = row["Chinese"]
                        key_dict["kor"][key] = row["Korean"]
                row["English"] = key
            tsv_keyed.append(row)
        except Exception:
            traceback.print_exc()
            print("Error processing tsv: " + id_num)
            raise

    return tsv_keyed, key_dict