#include <cstdint>
#include <array>

#include <fmt/format.h>

#include "os_utils.h"
#include "utf_convert.h"

constexpr uint32_t kInvalidPartition = 0xFFFFFFFFu;

struct LeapPartitionItem {
    uint32_t fixed = 1;
    uint32_t start = 0;
    uint32_t end = 0;
    uint32_t size = 0;
};

struct MinatoPartitionItem {
    uint32_t start = 0;
    uint32_t end = 0;
    uint32_t size = 0;
    uint32_t flags = 0;
};

struct XeltekPartitionItem {
    uint32_t start = 0;
    uint32_t end = 0;
    uint32_t size = 0;
    uint32_t unknown = kInvalidPartition;
};


enum class PartitionType {
    kNone = 0,
    kLeap,
    kMinato,
    kXeltek
};

static_assert(sizeof(MinatoPartitionItem) == 16 &&
              sizeof(XeltekPartitionItem) == 16 &&
              sizeof(LeapPartitionItem) == 16);

struct SimplePartitionItem {
    uint32_t start = 0;
    uint32_t end = 0;
    uint32_t size = 0;
};

PartitionType ReadPartitionTable(const std::vector<uint8_t>& data, std::vector<SimplePartitionItem>& out) {
    if (data.size() % 16 != 0) {
        fmt::print("Invalid partition table size\n");
        return PartitionType::kNone;
    }
    if (data.empty()) {
        fmt::print("Empty partition table\n");
        return PartitionType::kNone;
    }
    out.clear();
    out.reserve(data.size() / 16);
    bool isLeap = false;
    const std::array<char, 16> kLeapHeader{"GROUP DEFINE2\0\0"};
    if (memcmp(data.data(), kLeapHeader.data(), 16) == 0) {
        isLeap = true;
    }
    if (isLeap) {
        const auto* items = reinterpret_cast<const LeapPartitionItem*>(data.data());
        // start from 1
        for (size_t i = 1; i < data.size() / 16; i++) {
            const auto& item = items[i];
            if (item.start == kInvalidPartition) {
                break;
            }
            if (item.fixed != 1) {
                fmt::print("Invalid partition table, got fixed value {}\n", item.fixed);
                return PartitionType::kNone;
            }
            SimplePartitionItem simpleItem;
            simpleItem.start = item.start;
            simpleItem.end = item.end;
            simpleItem.size = item.size;
            out.push_back(simpleItem);
        }
        return PartitionType::kLeap;
    } else {
        struct UnknownPartitionItem {
            uint32_t start = 0;
            uint32_t end = 0;
            uint32_t size = 0;
            uint32_t unknown = 0;
        };
        PartitionType type = PartitionType::kNone;
        const auto* items = reinterpret_cast<const UnknownPartitionItem*>(data.data());
        if (items[0].unknown == 0 || items[0].unknown == 1) {
            type = PartitionType::kMinato;
        } else if (items[0].unknown == kInvalidPartition) {
            type = PartitionType::kXeltek;
        } else {
            fmt::print("Invalid partition table, unknown value {}\n", items[0].unknown);
            return PartitionType::kNone;
        }
        for (size_t i = 0; i < data.size() / 16; i++) {
            const auto& item = items[i];
            if (item.start == kInvalidPartition) {
                break;
            }
            // check unknown
            if (item.unknown != kInvalidPartition && item.unknown != 0 && item.unknown != 1) {
                fmt::print("Invalid partition table, got unknown value {}\n", item.unknown);
                return PartitionType::kNone;
            }
            SimplePartitionItem simpleItem;
            simpleItem.start = item.start;
            simpleItem.end = item.end;
            simpleItem.size = item.size;
            out.push_back(simpleItem);
        }
        return type;
    }
}

std::vector<uint8_t> GenerateLeapPartitionTable(const std::vector<SimplePartitionItem>& items) {
    std::vector<uint8_t> data;
    data.resize((items.size() + 2) * 16);
    std::array<char, 16> header{"GROUP DEFINE2\0\0"};
    memcpy(data.data(), header.data(), 16);
    auto* outItems = reinterpret_cast<LeapPartitionItem*>(data.data());
    for (size_t i = 0; i < items.size(); i++) {
        const auto& item = items[i];
        outItems[i + 1].fixed = 1;
        outItems[i + 1].start = item.start;
        outItems[i + 1].end = item.end;
        outItems[i + 1].size = item.size;
    }
    // add end
    outItems[items.size() + 1].fixed = kInvalidPartition;
    outItems[items.size() + 1].start = kInvalidPartition;
    outItems[items.size() + 1].end = kInvalidPartition;
    outItems[items.size() + 1].size = kInvalidPartition;
    return data;
}

std::vector<uint8_t> GenerateXeltekMinatoPartitionTable(const std::vector<SimplePartitionItem>& items,
                                                        PartitionType type) {
    std::vector<uint8_t> data;
    data.resize(items.size() * 16);
    auto* outItems = reinterpret_cast<MinatoPartitionItem*>(data.data());
    for (size_t i = 0; i < items.size(); i++) {
        const auto& item = items[i];
        outItems[i].start = item.start;
        outItems[i].end = item.end;
        outItems[i].size = item.size;
        if (type == PartitionType::kMinato) {
            outItems[i].flags = 0;
        } else if (type == PartitionType::kXeltek) {
            outItems[i].flags = kInvalidPartition;
        }
    }
    return data;
}

std::vector<uint8_t> GeneratePartitionTable(const std::vector<SimplePartitionItem>& items, PartitionType type) {
    if (type == PartitionType::kLeap) {
        return GenerateLeapPartitionTable(items);
    } else {
        return GenerateXeltekMinatoPartitionTable(items, type);
    }
}

const char* GetPartionTypeLongName(PartitionType type) {
    switch (type) {
        case PartitionType::kLeap:
            return "Leap Electronic NAND Flash programmer";
        case PartitionType::kMinato:
            return "Minato NAND Flash programmer, Data I/O NAND Flash programmer";
        case PartitionType::kXeltek:
            return "Xeltek Universal IC Chip Programmer";
        default:
            return "Unknown";
    }
}

int MainEntry(int argc, char** argv) {
    PartitionType inputType = PartitionType::kNone;
    PartitionType outputType = PartitionType::kNone;
    std::string inputFile;
    std::string outputFile;
    bool quiet = false;
    if (argc == 1 || (argc == 2 && (strcmp(argv[1], "-h") == 0 || strcmp(argv[1], "--help") == 0))) {
        fmt::print("Usage: {} -i <input> [-q] [-L/M/X output]\n", argv[0]);
        fmt::print("  -i <input>  : input file\n");
        fmt::print("  -q          : quiet mode\n");
        fmt::print("  -L <output> : output file for Leap\n");
        fmt::print("  -M <output> : output file for Minato\n");
        fmt::print("  -X <output> : output file for Xeltek\n");
        return 1;
    }
    // parse arguments
    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "-i") == 0) {
            if (i + 1 < argc) {
                inputFile = argv[i + 1];
                i++;
            }
        } else if (strcmp(argv[i], "-q") == 0) {
            quiet = true;
        } else if (strcmp(argv[i], "-L") == 0) {
            if (i + 1 < argc) {
                outputFile = argv[i + 1];
                outputType = PartitionType::kLeap;
                i++;
            }
        } else if (strcmp(argv[i], "-M") == 0) {
            if (i + 1 < argc) {
                outputFile = argv[i + 1];
                outputType = PartitionType::kMinato;
                i++;
            }
        } else if (strcmp(argv[i], "-X") == 0) {
            if (i + 1 < argc) {
                outputFile = argv[i + 1];
                outputType = PartitionType::kXeltek;
                i++;
            }
        }
    }
    if (inputFile.empty()) {
        fmt::print("Input file is required\n");
        return 1;
    }
    std::vector<uint8_t> inputData;
    if (int rc; !ReadFile(inputFile, inputData, 512, rc)) {
        fmt::print("Unable to open '{}', rc={}, {}", inputFile, rc, GetErrorMsg(-rc));
        return 1;
    }
    std::vector<SimplePartitionItem> items;
    inputType = ReadPartitionTable(inputData, items);
    if (inputType == PartitionType::kNone) {
        fmt::print("Invalid or unknown partition table\n");
        return 1;
    }
    if (items.empty()) {
        fmt::print("Empty partition table\n");
        return 1;
    }
    bool hasError = false;
    {
        // check if output is valid
        // 1. range check
        for (const auto& item: items) {
            if (item.start >= item.end || item.start + item.size > item.end) {
                fmt::print("Invalid partition table, start={}, end={}, size={}\n", item.start, item.end, item.size);
                hasError = true;
            }
        }
        // 2. items should be sorted and not overlap
        uint32_t last = 0;
        for (const auto& item: items) {
            if (item.start < last) {
                fmt::print("Invalid partition table, items are not sorted or overlap\n");
                hasError = true;
                break;
            }
            last = item.end;
        }
    }
    if (hasError) {
        fmt::print("The partition table has errors, aborting\n");
        return 1;
    }
    // if not quiet, print the partition table
    if (!quiet) {
        std::string buf;
        buf += fmt::format("Input: {} partition(s), {}\n", items.size(), GetPartionTypeLongName(inputType));
        /**
         * Index     Start       End      Size
         *    31  00000000  00000000  00000000
         */
        buf += "Index     Start       End      Size\n";
        for (size_t i = 0; i < items.size(); i++) {
            const auto& item = items[i];
            buf += fmt::format("{:5}  {:08x}  {:08x}  {:08x}\n", i + 1, item.start, item.end, item.size);
        }
        fmt::print("{}", buf);
    }
    if (outputType != PartitionType::kNone && !outputFile.empty()) {
        std::vector<uint8_t> outputData = GeneratePartitionTable(items, outputType);
        if (int rc; !WriteFile(outputFile, outputData, rc)) {
            fmt::print("Unable to write '{}', rc={}, {}\n", outputFile, rc, GetErrorMsg(-rc));
            return 1;
        }
        if (!quiet) {
            fmt::print("Written to '{}', format: {}\n", outputFile, GetPartionTypeLongName(outputType));
        }
    }
    return 0;
}

// The real main entry
#if defined(_WIN32) || defined(_WIN64)

#include <windows.h>

int main(int argc, char** argv) {
    int argcW = 0;
    wchar_t** argvW = CommandLineToArgvW(GetCommandLineW(), &argcW);
    if (argvW == nullptr) {
        fmt::print("Unable to parse command line\n");
        return 1;
    }
    std::vector<std::string> args;
    args.reserve(argcW);
    for (int i = 0; i < argcW; i++) {
        args.push_back(utfcvt::WideToUTF8(argvW[i]));
    }
    std::vector<char*> argvC;
    argvC.reserve(args.size());
    for (auto& arg: args) {
        argvC.push_back(arg.data());
    }
    LocalFree(argvW);
    SetConsoleCP(CP_UTF8);
    SetConsoleOutputCP(CP_UTF8);
    return MainEntry((int) argvC.size(), argvC.data());
}

#else

int main(int argc, char** argv) {
    return MainEntry(argc, argv);
}

#endif