# Copyright (c) 2020-2021 The reone project contributors

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

"""Script to generate a filelist for training a Flowtron model of a particular NPC."""

import glob
import json
import os
import random
import sys

extract_dir = r"D:\OpenKotOR\Extract\KotOR"

if not os.path.exists(extract_dir):
    raise RuntimeError("Extraction directory does not exist")


def index_or_negative_one(string, substr, beg=0):
    try:
        return string.index(substr, beg)
    except ValueError:
        return -1


def erase_special_chars(string, left, right):
    left_idx = index_or_negative_one(string, left)
    if left_idx == -1:
        return (string, False)

    right_idx = index_or_negative_one(string, right, left_idx + len(left))
    if right_idx == -1:
        return (string, False)

    return (string[:left_idx] + string[(right_idx + len(right)):], True)


def clear_text(text):
    while True:
        text, found = erase_special_chars(text, "[", "]")
        if not found:
            text, found = erase_special_chars(text, "{", "}")
            if not found:
                text, found = erase_special_chars(text, "::", "::")
                if not found:
                    break

    return text.strip()


def match_entry(entry_speaker, entry_voresref, speaker, voresref):
    if entry_speaker:
        return speaker in entry_speaker

    return voresref in entry_voresref if voresref else False


def get_lines_from_dlg(obj, speaker, voresref, tlk_strings, uniq_sounds):
    global extract_dir

    wav_dir = os.path.join(extract_dir, "voices")
    if not os.path.exists(wav_dir):
        return []

    lines = []
    if "EntryList" in obj:
        for entry in obj["EntryList"]:
            if "VO_ResRef" in entry:
                textstrref = int(entry["Text"].split("|")[0])
                entry_speaker = entry["Speaker"].lower()
                entry_voresref = entry["VO_ResRef"].lower()
                if textstrref != -1 and match_entry(entry_speaker, entry_voresref, speaker, voresref) and not entry_voresref in uniq_sounds:
                    text = clear_text(tlk_strings[textstrref])
                    if text:
                        wav_filename = os.path.join(wav_dir, entry_voresref + ".wav")
                        if os.path.exists(wav_filename):
                            lines.append("{}|{}|0\n".format(wav_filename, text))
                            uniq_sounds.add(entry_voresref)

    return lines


def generate_filelist(speaker, voresref):
    """
    :param speaker: substring to search for in Speaker field of dialog entries
    :param voresref: substring to search for in VO_ResRef field of dialog entries
    """
    global extract_dir

    tlk_strings = dict()
    uniq_sounds = set()

    # Read strings from dialog.tlk into dictionary
    tlk_path = os.path.join(extract_dir, "dialog.tlk.json")
    if os.path.exists(tlk_path):
        with open(tlk_path, "r") as fp:
            obj = json.load(fp)
            if "strings" in obj:
                for string in obj["strings"]:
                    strref = int(string["_index"])
                    text = string["text"]
                    tlk_strings[strref] = text

    # Extract lines from DLG files
    lines = []
    for f in glob.glob("{}/**".format(extract_dir), recursive=True):
        if f.endswith(".dlg.json"):
            with open(f, "r") as fp:
                obj = json.load(fp)
                lines.extend(get_lines_from_dlg(obj, speaker, voresref, tlk_strings, uniq_sounds))

    # Split lines into training and validation filelists
    random.shuffle(lines)
    num_val = int(5 * len(lines) / 100)
    lines_train = lines[num_val:]
    lines_val = lines[:num_val]

    with open(speaker + "_train_filelist.txt", "w") as fp:
        fp.writelines(lines_train)
    with open(speaker + "_val_filelist.txt", "w") as fp:
        fp.writelines(lines_val)


if len(sys.argv) > 1:
    # Interpet first argument as Speaker substring, and second argument as VO_ResRef substring
    speaker = sys.argv[1]
    voresref = sys.argv[2] if len(sys.argv) > 2 else None
    generate_filelist(speaker, voresref)
else:
    print("Usage: python filelist.py SPEAKER [VORESREF]")