#include <DevCommand/WTestCommand/WTestCommand.hpp>
#include <Log/Log.hpp>
#include <AtomicDataStore/AtomicDataStore.hpp>
#include <iomanip>

void WTestCommand::testHashTable(FileIO &file, uint64_t pos, uint64_t size, std::string hashPath, std::vector<Region> &regions)
{
    std::vector<uint8_t> compressedData(size);
    if (!file.read(pos, compressedData.data(), size))
    {
        regions.push_back({pos, size, formatString("Hash table at $+: cannot be read", hashPath), true});
        return;
    }

    std::vector<uint8_t> uncompressedData;
    if (!uncompressBytes(compressedData, uncompressedData))
    {
        regions.push_back({pos, size, formatString("Hash table at $+: cannot uncompress", hashPath), true});
        return;
    }

    AtomicDataStore::HashTable hashtab;
    if (uncompressedData.size() != sizeof(hashtab))
    {
        regions.push_back({pos, size, formatString("Hash table at $+: invalid size", hashPath), true});
        return;
    }

    memcpy(&hashtab, uncompressedData.data(), sizeof(hashtab));

    regions.push_back({pos, size, formatString("Hash table at $+: OK", hashPath), false});

    for (int i = 0; i < AtomicDataStore::numHashTableEntries; i++)
    {
        auto &entry = hashtab.entries[i];

        if (entry.type == AtomicDataStore::HashTableEntryType::Unmapped)
        {
            // NOP
        }
        else if (entry.type == AtomicDataStore::HashTableEntryType::NestedHashTable)
        {
            testHashTable(file, entry.pos, entry.dataSize, formatString("$+ -> $+", hashPath, i), regions);
        }
        else if (entry.type == AtomicDataStore::HashTableEntryType::FileDescriptor)
        {
            std::vector<char> nameBytes;
            nameBytes.resize(entry.fileNameLen);

            if (!file.read(entry.pos, nameBytes.data(), nameBytes.size()))
            {
                regions.push_back({entry.pos, entry.dataSize + entry.fileNameLen, formatString("File $+ -> $+: cannot read name", hashPath, i), true});
                return;
            }

            std::string name((char*) nameBytes.data(), entry.fileNameLen);

            std::vector<uint8_t> payload;
            payload.resize(entry.dataSize);
            if (!file.read(entry.pos + entry.fileNameLen, payload.data(), payload.size()))
            {
                regions.push_back({entry.pos, entry.dataSize + entry.fileNameLen, formatString("File $+ -> $+ '$+': cannot read data", hashPath, i, name), true});
                return;
            }

            std::vector<uint8_t> data;
            if (!uncompressBytes(payload, data))
            {
                regions.push_back({entry.pos, entry.dataSize + entry.fileNameLen, formatString("File $+ -> $+ '$+': cannot uncompress", hashPath, i, name), true});
                return;
            }

            regions.push_back({entry.pos, entry.dataSize + entry.fileNameLen, formatString("File $+ -> $+ '$+': OK", hashPath, i, name), false});
        }
        else
        {
            Log::warn(formatString("Unrecognised hash table entry type: ", static_cast<int>(entry.type)));
        }
    }
}

int WTestCommand::execute(const std::vector<std::string> &args)
{
    if (args.size() != 3)
    {
        Log::error("SYNTAX: wtest <world-file>");
        return 1;
    }

    std::string worldFilePath = args[2];

    AtomicDataStore store;
    store.open(worldFilePath);
    const std::string &error = store.getError();

    if (!error.empty())
    {
        Log::error(error);
        return 1;
    }

    AtomicDataStore::IndexBlock indexBlocks[2];
    if (!store.file.read(sizeof(AtomicDataStore::MainHeader), indexBlocks, sizeof(indexBlocks)))
    {
        Log::error("Failed to read the index blocks.");
        return 1;
    }

    size_t validIndex = 0;
    for (size_t i = 0; i < 2; i++)
    {
        if (indexBlocks[i].isValid())
        {
            validIndex = i;
        }
    }

    if (!indexBlocks[0].isValid() && !indexBlocks[1].isValid())
    {
        std::cerr << "There are no valid index blocks in this file." << std::endl;
        return 1;
    }

    if (indexBlocks[0].isValid() == indexBlocks[1].isValid())
    {
        std::cerr << "Currently we only handle the case where only 1 index block is valid." << std::endl;
        return 1;
    }

    std::vector<Region> regions;
    regions.push_back({
        indexBlocks[validIndex].freeRegionTablePos,
        sizeof(AtomicDataStore::Region) * indexBlocks[validIndex].numFreeRegions,
        "Free region table",
        false,
    });

    regions.push_back({
        0,
        sizeof(AtomicDataStore::MainHeader) + 2 * sizeof(AtomicDataStore::IndexBlock),
        "Header and index blocks",
        false,
    });

    for (size_t regionIndex = 0; regionIndex < indexBlocks[validIndex].numFreeRegions; regionIndex++)
    {
        AtomicDataStore::Region region;
        if (!store.file.read(indexBlocks[validIndex].freeRegionTablePos + sizeof(region) * regionIndex, &region, sizeof(region)))
        {
            std::cout << "# Failed to read free region table entry " << regionIndex << std::endl;
        }

        if (region.size == 0)
        {
            if (region.basePos != 0)
            {
                regions.push_back({region.basePos, region.size, "Free region: redundant", true});
            }

            continue;
        }

        regions.push_back({region.basePos, region.size, "Free region", false});
    }

    testHashTable(store.file, store.rootHashTablePos, store.rootHashTableSize, "ROOT", regions);

    std::sort(regions.begin(), regions.end(), [](const Region& a, const Region& b) {
        return a.start < b.start;
    });

    std::vector<Region> updatedRegions;
    for (size_t i = 0; i < regions.size() - 1; i++)
    {
        Region& current = regions[i];
        Region& next = regions[i + 1];
        
        // Check for overlap
        if ((current.start + current.size) == 0 || current.start + current.size > next.start)
        {
            current.usage += " [OVERLAPS BELOW]";
            current.bad = true;

            next.usage += " [OVERLAPS ABOVE]";
            next.bad = true;
        }

        // Add the current region to the updated list
        updatedRegions.push_back(current);

        // Check for unaccounted space
        if ((current.start + current.size) != 0 && current.start + current.size < next.start)
        {
            Region unaccountedRegion;
            unaccountedRegion.start = current.start + current.size;
            unaccountedRegion.size = next.start - unaccountedRegion.start;
            unaccountedRegion.usage = "[UNACCOUNTED]";
            unaccountedRegion.bad = true;
            updatedRegions.push_back(unaccountedRegion);
        }
    }

    // Add the last region to the updated list
    updatedRegions.push_back(regions.back());

    regions = updatedRegions;

    std::cout << std::endl << "Map of '" << worldFilePath << "':" << std::endl;
    std::cout << "┌──────────────────┬──────────────────┬────────────────────────────────────────────────────────────────────────────────┐" << std::endl;
    std::cout << "│Base              │Size              │                                                                                │" << std::endl;
    std::cout << "├──────────────────┼──────────────────┼────────────────────────────────────────────────────────────────────────────────┤" << std::endl;

    int problems = 0;

    for (auto &region : regions)
    {
        const char *colorstr = region.bad ? "\033[1;33m" : "";

        if (region.bad)
        {
            problems++;
        }

        std::cout << "│" << std::hex << std::setfill('0') << std::setw(16) << region.start << "  │ "
            << std::setfill(' ') << std::setw(17) << formatDataSize(region.size) << "│ " << colorstr
            << std::setw(79) << region.usage << "\033[0m│" << std::endl;
    }

    std::cout << "└──────────────────┴──────────────────┴────────────────────────────────────────────────────────────────────────────────┘" << std::endl;

    std::cout << std::dec;
    std::cout << "Found " << problems << " problems in this file." << std::endl;

    return problems;
}