// Copyright 2022, FOSS-VG Developers and Contributers
//
// This program is free software: you can redistribute it and/or modify it
// under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, version 3.
//
// 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 Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// version 3 along with this program.
// If not, see https://www.gnu.org/licenses/agpl-3.0.en.html

#include <iostream>
#include <fstream>
#include <vector>
#include <tinyutf8/tinyutf8.h>

#include "../lib/nbt.hpp"
#include "../lib/cli.hpp"

#define EXIT_SUCCESS 0
#define EXIT_RUNTIME 1
#define EXIT_USAGE 2
#define EXIT_UNIMPLEMENTED 3

void printTagTypeName(NBT::Tag::Generic* tag, uint64_t offsetBytes) {
    std::cout << "[" << offsetBytes << ": ";
    switch (tag->getTagType()) {
        case NBT::TagType::END:
            std::cout << "End";
            break;
        case NBT::TagType::INT8:
            std::cout << "8 Bit Integer";
            break;
        case NBT::TagType::INT16:
            std::cout << "16 Bit Integer";
            break;
        case NBT::TagType::INT32:
            std::cout << "32 Bit Integer";
            break;
        case NBT::TagType::INT64:
            std::cout << "64 Bit Integer";
            break;
        case NBT::TagType::FLOAT:
            std::cout << "Float";
            break;
        case NBT::TagType::DOUBLE:
            std::cout << "Double";
            break;
        case NBT::TagType::INT8_ARRAY:
            std::cout << "Array of 8 Bit Integers";
            break;
        case NBT::TagType::STRING:
            std::cout << "String";
            break;
        case NBT::TagType::LIST:
            std::cout << "List";
            break;
        case NBT::TagType::COMPOUND:
            std::cout << "Compound";
            break;
        case NBT::TagType::INT32_ARRAY:
            std::cout << "Array of 32 Bit Integers";
            break;
        case NBT::TagType::INT64_ARRAY:
            std::cout << "Array of 64 Bit Integers";
            break;
        default:
            // WTF? How'd you even get here?
            std::cout << "Unknown Type";
    }
    std::cout << "]";
    if (tag->name == "") {
        std::cout << ":";
    } else {
        std::cout << " " << tag->name << ":";
    }
}

void printNBytes(uint64_t bytes) {
    std::cout << bytes << (bytes==1? " byte":" bytes");
}

void drawTree(NBT::Tag::Generic* tag, tiny_utf8::string prefix, uint64_t offsetBytes, bool listMode=false) {
    std::vector<uint8_t> serialized;
    uint64_t headerSize = 0;
    if (listMode) {
        tag->serializeWithoutHeader(&serialized);

        std::cout << prefix << "|–Payload: ";
        printNBytes(serialized.size());
        std::cout << std::endl;
    } else {
        tag->serialize(&serialized);
        if (tag->getTagType() == NBT::TagType::END) {
            headerSize = 1;
        } else {
            headerSize = (uint64_t) NBT::Helper::readInt16(serialized.data(), serialized.size(), 1).value+3;
        }

        std::cout << prefix << "|–Header: ";
        printNBytes(headerSize);
        std::cout << std::endl;
        std::cout << prefix << "|–Payload: ";
        printNBytes(serialized.size() - headerSize);
        std::cout << std::endl;
    }
    if (tag->getTagType() == NBT::TagType::END) {
        std::cout << prefix << "'–Total: ";
    } else {
        std::cout << prefix << "|–Total: ";
    }
    printNBytes(serialized.size());
    std::cout << std::endl;

    switch (tag->getTagType()) {
        case NBT::TagType::END:
            break;
        case NBT::TagType::INT8:
            std::cout << prefix << "'–Value: " << (int32_t) reinterpret_cast<NBT::Tag::Int8*>(tag)->getValue() << std::endl;
            break;
        case NBT::TagType::INT16:
            std::cout << prefix << "'–Value: " << (int32_t) reinterpret_cast<NBT::Tag::Int16*>(tag)->getValue() << std::endl;
            break;
        case NBT::TagType::INT32:
            std::cout << prefix << "'–Value: " << reinterpret_cast<NBT::Tag::Int32*>(tag)->getValue() << std::endl;
            break;
        case NBT::TagType::INT64:
            std::cout << prefix << "'–Value: " << reinterpret_cast<NBT::Tag::Int64*>(tag)->getValue() << std::endl;
            break;
        case NBT::TagType::FLOAT:
            std::cout << prefix << "'–Value: " << reinterpret_cast<NBT::Tag::Float*>(tag)->getValue() << std::endl;
            break;
        case NBT::TagType::DOUBLE:
            std::cout << prefix << "'–Value: " << reinterpret_cast<NBT::Tag::Double*>(tag)->getValue() << std::endl;
            break;
        case NBT::TagType::INT8_ARRAY: {
            NBT::Tag::Int8Array* array = reinterpret_cast<NBT::Tag::Int8Array*>(tag);
            std::cout << prefix << "|–Length: " << array->length() << std::endl;
            std::cout << prefix << "'–Values: " << std::endl;
            for (uint64_t i=0; i<array->length(); i++) {
                if (i == array->length()-1) {
                    std::cout << prefix << "  '–" << (int64_t) array->getValue(i).value << std::endl;
                } else {
                    std::cout << prefix << "  |–" << (int64_t) array->getValue(i).value << std::endl;
                }
            }
            break;
        }
        case NBT::TagType::STRING:
            std::cout << prefix << "'–Value: " << reinterpret_cast<NBT::Tag::String*>(tag)->getValue() << std::endl;
            break;
        case NBT::TagType::LIST: {
            NBT::Tag::List* list = reinterpret_cast<NBT::Tag::List*>(tag);

            std::cout << prefix << "|–Contained Type: ";
            switch (list->getContainedType()) {
                case NBT::TagType::END:
                    std::cout << "End";
                    break;
                case NBT::TagType::INT8:
                    std::cout << "8 Bit Integer";
                    break;
                case NBT::TagType::INT16:
                    std::cout << "16 Bit Integer";
                    break;
                case NBT::TagType::INT32:
                    std::cout << "32 Bit Integer";
                    break;
                case NBT::TagType::INT64:
                    std::cout << "64 Bit Integer";
                    break;
                case NBT::TagType::FLOAT:
                    std::cout << "Float";
                    break;
                case NBT::TagType::DOUBLE:
                    std::cout << "Double";
                    break;
                case NBT::TagType::INT8_ARRAY:
                    std::cout << "Array of 8 Bit Integers";
                    break;
                case NBT::TagType::STRING:
                    std::cout << "String";
                    break;
                case NBT::TagType::LIST:
                    std::cout << "List";
                    break;
                case NBT::TagType::COMPOUND:
                    std::cout << "Compound";
                    break;
                case NBT::TagType::INT32_ARRAY:
                    std::cout << "Array of 32 Bit Integers";
                    break;
                case NBT::TagType::INT64_ARRAY:
                    std::cout << "Array of 64 Bit Integers";
                    break;
                default:
                    // WTF? How'd you even get here?
                    std::cout << "Unknown Type";
            }
            std::cout << std::endl;

            if (list->length() > 0) {
                std::cout << prefix << "|–Length: " << list->length() << std::endl;
                std::cout << prefix << "|" << std::endl;
            } else {
                std::cout << prefix << "'–Length: " << list->length() << std::endl;
            }

            offsetBytes = offsetBytes + headerSize + 5;
            for (uint64_t i=0; i<list->length(); i++) {
                if (i == list->length()-1) {
                    std::cout << prefix << "'–";
                    printTagTypeName(list->getElementPointer(i).value, offsetBytes);
                    std::cout << std::endl;
                    drawTree(list->getElementPointer(i).value, prefix+"  ", offsetBytes, true);
                } else {
                    std::cout << prefix << "|–";
                    printTagTypeName(list->getElementPointer(i).value, offsetBytes);
                    std::cout << std::endl;
                    drawTree(list->getElementPointer(i).value, prefix+"| ", offsetBytes, true);
                }

                std::vector<uint8_t> serializedElement;
                list->getElementPointer(i).value->serializeWithoutHeader(&serializedElement);
                offsetBytes += serializedElement.size();
            }
            break;
            }
        case NBT::TagType::COMPOUND: {
            NBT::Tag::Compound* compound = reinterpret_cast<NBT::Tag::Compound*>(tag);

            std::cout << prefix << "|–Length: " << compound->length() << std::endl;
            std::cout << prefix << "|" << std::endl;

            offsetBytes = offsetBytes + headerSize;
            for (uint64_t i=0; i<compound->length(); i++) {
                if (i == compound->length()-1) {
                    std::cout << prefix << "'–";
                    printTagTypeName(compound->getElementPointer(i).value, offsetBytes);
                    std::cout << std::endl;
                    drawTree(compound->getElementPointer(i).value, prefix+"  ", offsetBytes);
                } else {
                    std::cout << prefix << "|–";
                    printTagTypeName(compound->getElementPointer(i).value, offsetBytes);
                    std::cout << std::endl;
                    drawTree(compound->getElementPointer(i).value, prefix+"| ", offsetBytes);
                }

                std::vector<uint8_t> serializedElement;
                compound->getElementPointer(i).value->serialize(&serializedElement);
                offsetBytes += serializedElement.size();
            }
            break;
        }
        case NBT::TagType::INT32_ARRAY: {
            NBT::Tag::Int32Array* array = reinterpret_cast<NBT::Tag::Int32Array*>(tag);
            std::cout << prefix << "|–Length: " << array->length() << std::endl;
            std::cout << prefix << "'–Values: " << std::endl;
            for (uint64_t i=0; i<array->length(); i++) {
                if (i == array->length()-1) {
                    std::cout << prefix << "  '–" << array->getValue(i).value << std::endl;
                } else {
                    std::cout << prefix << "  |–" << array->getValue(i).value << std::endl;
                }
            }
            break;
        }
        case NBT::TagType::INT64_ARRAY: {
            NBT::Tag::Int64Array* array = reinterpret_cast<NBT::Tag::Int64Array*>(tag);
            std::cout << prefix << "|–Length: " << array->length() << std::endl;
            std::cout << prefix << "'–Values: " << std::endl;
            for (uint64_t i=0; i<array->length(); i++) {
                if (i == array->length()-1) {
                    std::cout << prefix << "  '–" << array->getValue(i).value << std::endl;
                } else {
                    std::cout << prefix << "  |–" << array->getValue(i).value << std::endl;
                }
            }
            break;
        }
        default:
            // WTF? How'd you even get here?
            std::cout << prefix << "'–???" << std::endl;
    }

    std::cout << prefix << std::endl;
}

int main(int argc, char* argv[]) {
    std::vector<CLI::Flag> flags;
    flags.push_back(CLI::Flag('h', "help", "print help and exit"));
    flags.push_back(CLI::Flag('l', "license", "print license information and exit"));
    std::vector<CLI::Option> options;
    std::vector<CLI::Argument> arguments;
    arguments.push_back(CLI::Argument("FILE", "path of the file to dump"));

    CLI::ArgumentsParser cliParser = CLI::ArgumentsParser(argc, argv, flags, options, arguments, "Present NBT in human or machine readable formats");

    if (cliParser.getFlag("help").value){
        std::cout << cliParser.getUsage() << std::endl;
        return EXIT_SUCCESS;
    }
    if (cliParser.getFlag("license").value){
        std::cout
            << "Copyright 2022, FOSS-VG Developers and Contributers\n"
            << "\n"
            << "DumpNBT is part of the FOSS-VG development tool suite.\n"
            << "\n"
            << "This program is free software: you can redistribute it and/or modify it\n"
            << "under the terms of the GNU Affero General Public License as published\n"
            << "by the Free Software Foundation, version 3.\n"
            << "\n"
            << "This program is distributed in the hope that it will be useful,\n"
            << "but WITHOUT ANY WARRANTY; without even the implied\n"
            << "warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.\n"
            << "See the GNU Affero General Public License for more details.\n"
            << "\n"
            << "You should have received a copy of the GNU Affero General Public License\n"
            << "version 3 along with this program.\n"
            << "If not, see https://www.gnu.org/licenses/agpl-3.0.en.html"
            << std::endl;
        return EXIT_SUCCESS;
    }

    if (cliParser.wrongUsage) {
        std::cerr << cliParser.getUsage() << std::endl;
        return EXIT_USAGE;
    }

    std::ifstream fileStream;
    fileStream.open(cliParser.getArgument(0).value, std::ios::in | std::ios::binary | std::ios::ate);
    if (!fileStream.is_open()) {
        std::cerr << argv[0] << ": Could not open file: " << cliParser.getArgument(0).value << std::endl;
        return EXIT_RUNTIME;
    }

    uint64_t fileSize = fileStream.tellg();
    fileStream.seekg(0, std::ios::beg);

    uint8_t data[fileSize];
    uint8_t* nextByte = new uint8_t;

    for (uint64_t i=0; i<fileSize; i++) {
        fileStream.read(reinterpret_cast<char*>(nextByte), 1);
        data[i] = *nextByte;
    }
    fileStream.close();

    ErrorOr<std::vector<NBT::Tag::Generic*>> tags = NBT::deserialize(data, fileSize);
    if (tags.isError) {
        std::cerr << "Invalid data." << std::endl;
        return EXIT_RUNTIME;
    }

    uint64_t offsetBytes = 0;
    for (uint64_t i=0; i<tags.value.size(); i++) {
        printTagTypeName(tags.value[i], offsetBytes);
        std::cout << std::endl;
        drawTree(tags.value[i], "", offsetBytes);

        std::vector<uint8_t> serialized;
        tags.value[i]->serialize(&serialized);
        offsetBytes += serialized.size();
    }

    return EXIT_SUCCESS;
}