#include <fmt/format.h>
#include <fstream>
#include <nlohmann/json.hpp>

using namespace std::string_literals;
using json = nlohmann::json;
using u8   = std::uint8_t;
using u16  = std::uint16_t;
using u32  = std::uint32_t;

const u32 RomStartAddress = 0x8000000;

namespace {

    struct LanguageTable {
        const char *name;
        u32 address;
    };

    enum Color {
        Color_White,
        Color_Red,
        Color_Green,
        Color_Blue,
        Color_Yellow,

        Color_Count,
    };

    constexpr const char *const ColorStrings[] = {
        [Color_White]  = "White",
        [Color_Red]    = "Red",
        [Color_Green]  = "Green",
        [Color_Blue]   = "Blue",
        [Color_Yellow] = "Yellow",
    };

    enum Input {
        Input_A,
        Input_B,
        Input_Left,
        Input_Right,
        Input_DUp,
        Input_DDown,
        Input_DLeft,
        Input_DRight,
        Input_Dpad,
        Input_Select,
        Input_Start,

        Input_Count,
    };

    constexpr const char *const InputStrings[] = {
        [Input_A]      = "A",
        [Input_B]      = "B",
        [Input_Left]   = "Left",
        [Input_Right]  = "Right",
        [Input_DUp]    = "DUp",
        [Input_DDown]  = "DDown",
        [Input_DLeft]  = "DLeft",
        [Input_DRight] = "DRight",
        [Input_Dpad]   = "Dpad",
        [Input_Select] = "Select",
        [Input_Start]  = "Start",
    };

    const std::map<u8, std::string> CharConvertArray = {
        {0x0a, "\n"},
        {0x0d, "\r"},
        {0x20, " "},
        {0x21, "!"},
        {0x22, "\""},
        {0x23, "#"},
        {0x24, "$"},
        {0x25, "%"},
        {0x26, "&"},
        {0x27, "\'"},
        {0x28, "("},
        {0x29, ")"},
        {0x2a, "*"},
        {0x2b, "+"},
        {0x2c, ","},
        {0x2d, "-"},
        {0x2e, "."},
        {0x2f, "/"},
        {0x30, "0"},
        {0x31, "1"},
        {0x32, "2"},
        {0x33, "3"},
        {0x34, "4"},
        {0x35, "5"},
        {0x36, "6"},
        {0x37, "7"},
        {0x38, "8"},
        {0x39, "9"},
        {0x3a, ":"},
        {0x3b, ";"},
        {0x3c, "<"},
        {0x3d, "="},
        {0x3e, ">"},
        {0x3f, "?"},
        {0x40, "@"},
        {0x41, "A"},
        {0x42, "B"},
        {0x43, "C"},
        {0x44, "D"},
        {0x45, "E"},
        {0x46, "F"},
        {0x47, "G"},
        {0x48, "H"},
        {0x49, "I"},
        {0x4a, "J"},
        {0x4b, "K"},
        {0x4c, "L"},
        {0x4d, "M"},
        {0x4e, "N"},
        {0x4f, "O"},
        {0x50, "P"},
        {0x51, "Q"},
        {0x52, "R"},
        {0x53, "S"},
        {0x54, "T"},
        {0x55, "U"},
        {0x56, "V"},
        {0x57, "W"},
        {0x58, "X"},
        {0x59, "Y"},
        {0x5a, "Z"},
        {0x5b, "["},
        {0x5c, "\'"},
        {0x5d, "]"},
        {0x5e, "^"},
        {0x5f, "_"},
        {0x60, "`"},
        {0x61, "a"},
        {0x62, "b"},
        {0x63, "c"},
        {0x64, "d"},
        {0x65, "e"},
        {0x66, "f"},
        {0x67, "g"},
        {0x68, "h"},
        {0x69, "i"},
        {0x6a, "j"},
        {0x6b, "k"},
        {0x6c, "l"},
        {0x6d, "m"},
        {0x6e, "n"},
        {0x6f, "o"},
        {0x70, "p"},
        {0x71, "q"},
        {0x72, "r"},
        {0x73, "s"},
        {0x74, "t"},
        {0x75, "u"},
        {0x76, "v"},
        {0x77, "w"},
        {0x78, "x"},
        {0x79, "y"},
        {0x7a, "z"},
        {0x82, ","},
        {0x84, "„"},
        {0x85, "⋯"},
        {0x8A, "Š"},
        {0x8B, "‹"},
        {0x8C, "Œ"},
        {0x8E, "Ž"},
        {0x91, "‘"},
        {0x92, "’"},
        {0x93, "“"},
        {0x94, "”"},
        {0x95, "·"},
        {0x99, "™"},
        {0x9A, "š"},
        {0x9B, "›"},
        {0x9C, "œ"},
        {0x9E, "ž"},
        {0x9F, "Ÿ"},
        {0xA1, "¡"},
        {0xA3, "♪"},
        {0xAA, "ª"},
        {0xAB, "«"},
        {0xB0, "º"},
        {0xB4, "'"},
        {0xB7, "´"},
        {0xBA, "º"},
        {0xBB, "»"},
        {0xBF, "¿"},
        {0xC0, "À"},
        {0xC1, "Á"},
        {0xC2, "Â"},
        {0xC3, "Ã"},
        {0xC4, "Ä"},
        {0xC5, "Å"},
        {0xC6, "Æ"},
        {0xC7, "Ç"},
        {0xC8, "È"},
        {0xC9, "É"},
        {0xCA, "Ê"},
        {0xCB, "Ë"},
        {0xCC, "Ì"},
        {0xCD, "Í"},
        {0xCE, "Î"},
        {0xCF, "Ï"},
        {0xD0, "Đ"},
        {0xD1, "Ñ"},
        {0xD2, "Ò"},
        {0xD3, "Ó"},
        {0xD4, "Ô"},
        {0xD5, "Õ"},
        {0xD6, "Ö"},
        {0xD7, "×"},
        {0xD8, "Ø"},
        {0xD9, "Ù"},
        {0xDA, "Ú"},
        {0xDB, "Û"},
        {0xDC, "Ü"},
        {0xDD, "Ý"},
        {0xDE, "Þ"},
        {0xDF, "β"},
        {0xE0, "à"},
        {0xE1, "á"},
        {0xE2, "â"},
        {0xE3, "ã"},
        {0xE4, "ä"},
        {0xE5, "å"},
        {0xE6, "æ"},
        {0xE7, "ç"},
        {0xE8, "è"},
        {0xE9, "é"},
        {0xEA, "ê"},
        {0xEB, "ë"},
        {0xEC, "ì"},
        {0xED, "í"},
        {0xEE, "î"},
        {0xEF, "ï"},
        {0xF0, "ð"},
        {0xF1, "ñ"},
        {0xF2, "ò"},
        {0xF3, "ó"},
        {0xF4, "ô"},
        {0xF5, "õ"},
        {0xF6, "ö"},
        {0xF7, "÷"},
        {0xF8, "ø"},
        {0xF9, "ù"},
        {0xFA, "ú"},
        {0xFB, "û"},
        {0xFC, "ü"},
        {0xFD, "ý"},
        {0xFE, "þ"},
        {0xFF, "ÿ"},
    };

    using ConvertFunction = std::string (*const)(const char *&);

    std::string Unk1Handler(const char *&ptr) {
        u8 a = *ptr++;
        return fmt::format("{{01:{:02X}}}", a);
    }

    std::string ColorHandler(const char *&ptr) {
        u8 color = *ptr++;
        if (color >= Color_Count)
            throw std::runtime_error(ptr);
        return fmt::format("{{Color:{}}}", ColorStrings[color]);
    }

    std::string SoundHandler(const char *&ptr) {
        u8 a = *ptr++;
        u8 b = *ptr++;
        return fmt::format("{{Sound:{:02X}:{:02X}}}", a, b);
    }

    std::string Unk4Handler(const char *&ptr) {
        u8 a = *ptr++;
        if (a == 0x10) {
            u8 b = *ptr++;
            return fmt::format("{{04:{:02X}:{:02X}}}", a, b);
        } else {
            return fmt::format("{{04:{:02X}}}", a);
        }
    }

    std::string ChoiceHandler(const char *&ptr) {
        u8 category = *ptr++;
        if (category == 0xff) {
            return "{Choice:FF}";
        } else {
            u8 action = *ptr++;
            return fmt::format("{{Choice:{:02X}:{:02X}}}", category, action);
        }
    }

    std::string VariableHandler(const char *&ptr) {
        u8 idx = *ptr++;
        if (idx == 0) {
            return "{Player}";
        } else {
            return fmt::format("{{Var:{:X}}}", idx);
        }
    }

    std::string Unk7Handler(const char *&ptr) {
        u8 a = *ptr++;
        u8 b = *ptr++;
        return fmt::format("{{07:{:02X}:{:02X}}}", a, b);
    }

    std::string Unk8Handler(const char *&ptr) {
        u8 a = *ptr++;
        if (a != 0xff)
            throw std::runtime_error("unmatched unk8: "s + std::to_string(a));
        return "{08:FF}";
    }

    std::string Unk9Handler(const char *&ptr) {
        u8 a = *ptr++;
        if (a != 0x78 && a != 0x00)
            throw std::runtime_error("unmatched unk8: "s + std::to_string(a));
        return fmt::format("{{09:{:02X}}}", a);
    }

    std::string InputHandler(const char *&ptr) {
        u8 key = *ptr++;
        if (key > 8)
            throw std::runtime_error("unmatched key: "s + std::to_string(key));
        return fmt::format("{{Key:{}}}", InputStrings[key]);
    }

    std::string SymbolHandler(const char *&ptr) {
        u8 a = *ptr++;
        return fmt::format("{{Symbol:{:02X}}}", a);
    }

    const std::map<u8, ConvertFunction> FuncConvertArray = {
        {0x01, Unk1Handler},
        {0x02, ColorHandler},
        {0x03, SoundHandler},
        {0x04, Unk4Handler},
        {0x05, ChoiceHandler},
        {0x06, VariableHandler},
        {0x07, Unk7Handler},
        {0x08, Unk8Handler},
        {0x09, Unk9Handler},
        {0x0c, InputHandler},
        {0x0f, SymbolHandler},
    };

    std::string ParseTMCString(const char *ptr) {
        std::string ret;

        while (*ptr) {
            u8 c = *ptr++;

            /* Convert character. */
            {
                const auto it = CharConvertArray.find(c);

                if (it != std::end(CharConvertArray)) {
                    ret += it->second;
                    continue;
                }
            }

            /* Convert function. */
            {
                const auto it = FuncConvertArray.find(c);

                if (it != std::end(FuncConvertArray)) {
                    ret += it->second(ptr);
                    continue;
                }
            }

            throw std::runtime_error(fmt::format("Unknown characters: {}", ptr));
        }

        return ret;
    }

    using RevertFunction = void (*)(const char *&, char *&);

    void ColorRevert(const char *&src, char *&dst) {
        *dst++ = 0x02;
        src++; // ':'
        for (u32 i = 0; i < Color_Count; i++) {
            const char *const color     = ColorStrings[i];
            const std::size_t color_len = std::strlen(color);
            if (std::strncmp(src, color, color_len) == 0) {
                *dst++ = i;
                src += color_len;
                return;
            }
        }

        throw std::runtime_error(fmt::format("Color not found: {}", src));
    }

    void SoundRevert(const char *&src, char *&dst) {
        *dst++ = 0x03;
        *dst++ = std::strtoul(++src, nullptr, 0x10);
        src += 2;
        *dst++ = std::strtoul(++src, nullptr, 0x10);
        src += 2;
    }

    void ChoiceRevert(const char *&src, char *&dst) {
        *dst++    = 0x05;
        u8 choice = std::strtoul(++src, nullptr, 0x10);

        *dst++ = choice;
        src += 2;

        if (choice == 0xff)
            return;

        *dst++ = std::strtoul(++src, nullptr, 0x10);
        src += 2;
    }

    void PlayerRevert(const char *&src, char *&dst) {
        *dst++ = 0x06;
        *dst++ = 0x00;

        (void)src;
    }

    void VariableRevert(const char *&src, char *&dst) {
        *dst++ = 0x06;
        src++; // ':'

        *dst++ = std::strtoul(src, nullptr, 0x10);
        src++;
    }

    void KeyRevert(const char *&src, char *&dst) {
        *dst++ = 0x0c;
        src++; // ':'
        for (u32 i = 0; i < Input_Count; i++) {
            const char *const input     = InputStrings[i];
            const std::size_t input_len = std::strlen(input);
            if (std::strncmp(src, input, input_len) == 0) {
                *dst++ = i;
                src += input_len;
                return;
            }
        }

        throw std::runtime_error(fmt::format("Input not found: {}", src));
    }

    void SymbolRevert(const char *&src, char *&dst) {
        *dst++ = 0x0f;
        src++; // ':'

        *dst++ = std::strtoul(src, nullptr, 0x10);
        src += 2;
    }

    const std::pair<std::string, RevertFunction> FuncRevertArray[] = {
        {"Color", ColorRevert},
        {"Sound", SoundRevert},
        {"Choice", ChoiceRevert},
        {"Player", PlayerRevert},
        {"Var", VariableRevert},
        {"Key", KeyRevert},
        {"Symbol", SymbolRevert},
    };

    void WriteTMCString(char *&dst, const std::string &src) {
        const char *ptr = src.data();

        while (*ptr) {
            /* Parse special */
            {
                if (*ptr == '{') {
                    ptr++;
                    const auto it = std::find_if(std::begin(FuncRevertArray), std::end(FuncRevertArray), [&](const auto &data) {
                        return std::strncmp(ptr, data.first.c_str(), data.first.size()) == 0;
                    });

                    if (it != std::end(FuncRevertArray)) {
                        ptr += it->first.size();
                        it->second(ptr, dst);
                    } else {
                        do {
                            *dst++ = std::strtoul(ptr, nullptr, 0x10);
                            ptr += 2;
                        } while (*ptr == ':' && (ptr++, true));
                    }

                    if (*ptr != '}')
                        throw std::runtime_error(fmt::format("unmatched characters: \"{}\"\n", ptr));

                    ptr++;
                    continue;
                }
            }

            /* Convert character. */
            {
                const auto it = std::find_if(std::begin(CharConvertArray), std::end(CharConvertArray), [ptr](const auto &data) {
                    return std::strncmp(ptr, data.second.c_str(), data.second.length()) == 0;
                });

                if (it != std::end(CharConvertArray)) {
                    ptr += it->second.size();
                    *dst++ = it->first;
                    continue;
                }
            }

            throw std::runtime_error(fmt::format("unmatched characters: \"{}\"\n", ptr));
        }
        *dst++ = '\0';
    }

}

void ExtractStringTable(std::string &rom_path, const std::vector<LanguageTable> &tables) {
    const std::vector<char> rom = [&]() {
        std::vector<char> rom;

        /* Open ROM. */
        std::ifstream rom_file(rom_path, std::ios::in);

        if (!rom_file.good()) {
            fmt::print(stderr, "Couldn't open ROM {}\n", rom_path);
            exit(EXIT_FAILURE);
        }

        /* Get ROM size. */
        rom_file.seekg(0, std::ios::end);
        const std::size_t rom_size = rom_file.tellg();

        /* Read ROM to buffer. */
        rom.resize(rom_size);
        rom_file.seekg(0, std::ios::beg);
        rom_file.read(&rom[0], rom_size);

        return rom;
    }();

    auto ReadAbsolute = [&](u32 address) -> u32 {
        return *(u32 *)&rom[address - RomStartAddress];
    };

    for (auto &language : tables) {
        /* Get category table start. */
        const std::size_t table_start = language.address;

        /* Get category count. */
        const std::size_t category_count = ReadAbsolute(table_start) / sizeof(u32);

        /* Read category offsets. */
        std::vector<u32> category_table(category_count);
        std::memcpy(&category_table[0], &rom[table_start - RomStartAddress], category_count * sizeof(u32));

        json j = json::array();

        for (auto &category_offset : category_table) {
            /* Get string table start. */
            const std::size_t category_start = table_start + category_offset;

            /* Get string count. */
            const std::size_t string_count = ReadAbsolute(category_start) / sizeof(u32);

            /* Read string offsets. */
            std::vector<u32> string_table(string_count);
            std::memcpy(&string_table[0], &rom[category_start - RomStartAddress], string_count * sizeof(u32));

            auto &category = j.emplace_back();

            for (std::size_t l = 0; l < string_count; l++) {
                /* Get string start. */
                const std::size_t string_start = category_start + string_table[l];

                /* Parse TMC. */
                category[l] = ParseTMCString(&rom[string_start - RomStartAddress]);
            }
        }

        const std::string out_path = language.name + ".json"s;
        std::ofstream ofs(out_path);
        ofs << j.dump(4);
    }
}

void PackStringTable(const std::string &src_path, const std::string &dst_path, const std::size_t out_size) {
    const json j = [&]() -> json {
        std::ifstream ifs(src_path);

        if (!ifs.good()) {
            fmt::print(stderr, "Couldn't open JSON {}\n", src_path);
            exit(EXIT_FAILURE);
        }

        return json::parse(ifs);
    }();

    std::vector<char> buffer(0x100000);
    std::uintptr_t root_start = (std::uintptr_t)buffer.data();

    char *root_ptr = buffer.data();
    char *table    = buffer.data() + j.size() * sizeof(u32);

    for (auto &category : j) {
        char *table_ptr = table;
        char *str_start = table_ptr + category.size() * sizeof(u32);
        char *str_ptr   = str_start;
        for (auto &str_j : category) {
            /* Write string offset to table. */
            *(u32 *)table_ptr = (std::uintptr_t)str_ptr - (std::uintptr_t)table;
            table_ptr += sizeof(u32);

            auto str = str_j.get<std::string>();

            /* Copy string to table. */
            WriteTMCString(str_ptr, str);
        }

        /* Align string table size by 16 bytes. */
        while (static_cast<std::uintptr_t>(str_ptr - str_start) & 0xf) {
            *str_ptr = 0xff;
            str_ptr++;
        }

        /* Write category offset to root table. */
        *(u32 *)root_ptr = (std::uintptr_t)table - root_start;
        root_ptr += sizeof(u32);

        table = str_ptr;
    }

    /* Align table end to 0x10 bytes. */
    while ((std::uintptr_t)table & 0xf) {
        *table++ = 0xff;
    }

    std::size_t table_size = (uintptr_t)table - root_start;

    if (out_size) {
        /* Abort if string table is too big. */
        if (table_size > out_size) {
            fmt::print(stderr, "Table is too big. Should be {} but was {}", out_size, table_size);
            exit(EXIT_FAILURE);
        }

        /* Pad table if it's too small. */
        while (table_size < out_size) {
            *table++ = 0xff;
            table_size++;
        }
    }

    std::ofstream ofs(dst_path);
    ofs.write(buffer.data(), table_size);
}

// clang-format off
const std::vector<LanguageTable> LanguageTableUS = {
    {"USA",      0x89B1D90},
};

const std::vector<LanguageTable> LanguageTableEU = {
    {"English",  0x89AEB60},
    {"French",   0x89F7420},
    {"German",   0x8A3EEB0},
    {"Spanish",  0x8A81E70},
    {"Italian",  0x8AC37A0},
};
// clang-format on

#include <getopt.h>

const char *progname;

void usage() {
    fmt::print(stderr,
               "tmc_strings (c) Luis Scheurenbrand\n"
               "Built: " __TIME__ " " __DATE__ "\n"
               "\n"
               "Usage: {} [options...]\n"
               "Options:\n"
               "  -x, --extract     Extract string table from ROM and store it in json format. (Default)\n"
               "  -p, --pack        Pack a string table from json format.\n"
               "  --region          Specify ROM region. [USA, EU]\n"
               "  --source          Specify source (-x: ROM, -p: JSON)\n"
               "  --dest            Specify string table destination.\n"
               "  --size            Specify string table size.\n",
               progname);
}

// clang-format off
constexpr const struct option long_options[] = {
    {"extract", no_argument,        nullptr, 'x'},
    {"pack",    no_argument,        nullptr, 'p'},
    {"region",  required_argument,  nullptr, 0},
    {"source",  required_argument,  nullptr, 1},
    {"dest",    required_argument,  nullptr, 2},
    {"size",    required_argument,  nullptr, 3},
    {},
};
// clang-format on

int main(int argc, char **argv) {
    std::string src_path;
    std::string dst_path;
    std::size_t max_size = 0;

    progname = (argc < 1) ? "tmc_strings" : argv[0];

    enum {
        Mode_Extract,
        Mode_Pack,
    } mode = Mode_Extract;

    enum {
        Region_USA,
        Region_JAPAN,
        Region_EU,
    } region = Region_USA;

    while (true) {
        int opt_index;
        int c = getopt_long(argc, argv, "xp", long_options, &opt_index);
        if (c == -1)
            break;

        switch (c) {
            case 'x':
                mode = Mode_Extract;
                break;
            case 'p':
                mode = Mode_Pack;
                break;
            case 0:
                if (strcmp(optarg, "USA") == 0) {
                    region = Region_USA;
                } else if (strcmp(optarg, "EU") == 0) {
                    region = Region_EU;
                } else if (strcmp(optarg, "JAPAN") == 0) {
                    region = Region_JAPAN;
                    fmt::print("Region unsupported\n\n");
                    usage();
                    exit(EXIT_FAILURE);
                }
                break;
            case 1:
                src_path = optarg;
                break;
            case 2:
                dst_path = optarg;
                break;
            case 3:
                max_size = std::strtoul(optarg, nullptr, 0);
                break;
            default:
                usage();
                exit(EXIT_FAILURE);
        }
    }

    switch (mode) {
        case Mode_Extract:
            if (src_path == "") {
                fmt::print("ROM source path not supplied.\n\n");
                usage();
                exit(EXIT_FAILURE);
            }
            ExtractStringTable(src_path, [region]() {
                switch (region) {
                    case Region_USA:
                        return LanguageTableUS;
                    case Region_EU:
                        return LanguageTableEU;
                    default:
                        fmt::print("Region unsupported\n\n");
                        usage();
                        exit(EXIT_FAILURE);
                }
            }());
            break;
        case Mode_Pack:
            if (src_path == "" || dst_path == "") {
                fmt::print("Source or destination path not supplied\n\n");
                usage();
                exit(EXIT_FAILURE);
            }
            if (max_size == 0) {
                fmt::print("Max size not supplied. Assuming shiftable.\n");
            }
            PackStringTable(src_path, dst_path, max_size);
    }

    return EXIT_SUCCESS;
}