#include "Decompressor.h"
#include "TorchUtils.h"

#include <stdexcept>
#include <mutex>
#include "spdlog/spdlog.h"
#include <Companion.h>

extern "C" {
#include <libmio0/mio0.h>
#include <libyay0/yay0.h>
#include <libyay0/yay1.h>
#include <libyaz0/yaz0.h>
#include <libmio0/tkmk00.h>
}

#include <bk_zip/bk_unzip.h>

std::unordered_map<uint32_t, DataChunk*> gCachedChunks;
std::mutex gDecompCacheMutex;

DataChunk* Decompressor::Decode(const std::vector<uint8_t>& buffer, const uint32_t offset, const CompressionType type,
                                const uint32_t in_size, bool ignoreCache) {

    {
        std::lock_guard<std::mutex> lock(gDecompCacheMutex);
        if (!ignoreCache && Torch::contains(gCachedChunks, offset)) {
            return gCachedChunks[offset];
        }
    }

    const unsigned char* in_buf = buffer.data() + offset;

    switch (type) {
        case CompressionType::MIO0: {
            mio0_header_t head;
            if (!mio0_decode_header(in_buf, &head)) {
                throw std::runtime_error("Failed to decode MIO0 header");
            }

            const auto decompressed = new uint8_t[head.dest_size];
            mio0_decode(in_buf, decompressed, nullptr);
            {
                std::lock_guard<std::mutex> lock(gDecompCacheMutex);
                gCachedChunks[offset] = new DataChunk{ decompressed, head.dest_size };
                return gCachedChunks[offset];
            }
        }
        case CompressionType::YAY0: {
            uint32_t size = 0;
            uint8_t* decompressed = yay0_decode(in_buf, &size);

            if (!decompressed) {
                throw std::runtime_error("Failed to decode YAY0");
            }

            {
                std::lock_guard<std::mutex> lock(gDecompCacheMutex);
                gCachedChunks[offset] = new DataChunk{ decompressed, size };
                return gCachedChunks[offset];
            }
        }
        case CompressionType::YAY1: {
            uint32_t size = 0;
            uint8_t* decompressed = yay1_decode(in_buf, &size);

            if (!decompressed) {
                throw std::runtime_error("Failed to decode YAY1");
            }

            {
                std::lock_guard<std::mutex> lock(gDecompCacheMutex);
                gCachedChunks[offset] = new DataChunk{ decompressed, size };
                return gCachedChunks[offset];
            }
        }
        case CompressionType::BKZIP: {
            uint32_t size = in_size;
            try {
                uint8_t* decompressed = BK64::bk_unzip(in_buf, &size);

                if (!decompressed) {
                    throw std::runtime_error("bk_unzip returned null");
                }

                {
                    std::lock_guard<std::mutex> lock(gDecompCacheMutex);
                    gCachedChunks[offset] = new DataChunk{ decompressed, size };
                    return gCachedChunks[offset];
                }
            } catch (const std::exception& e) {
                throw std::runtime_error(std::string(e.what()) + " (ROM offset 0x" + Torch::to_hex(offset, false) +
                                         " compressed size 0x" + Torch::to_hex(in_size, false) + ")");
            }
        }
        case CompressionType::YAZ0: {
            uint32_t size = 0;
            uint8_t* decompressed = yaz0_decode(in_buf, &size);

            if (!decompressed) {
                throw std::runtime_error("Failed to decode YAZ0");
            }

            gCachedChunks[offset] = new DataChunk{ decompressed, size };
            return gCachedChunks[offset];
        }
        default:
            throw std::runtime_error("Unknown compression type");
    }
}

DataChunk* Decompressor::DecodeTKMK00(const std::vector<uint8_t>& buffer, const uint32_t offset, const uint32_t size,
                                      const uint32_t alpha) {
    {
        std::lock_guard<std::mutex> lock(gDecompCacheMutex);
        if (Torch::contains(gCachedChunks, offset)) {
            return gCachedChunks[offset];
        }
    }

    const uint8_t* in_buf = buffer.data() + offset;

    const auto decompressed = new uint8_t[size];
    const auto rgba = new uint8_t[size];
    tkmk00_decode(in_buf, decompressed, rgba, alpha);
    {
        std::lock_guard<std::mutex> lock(gDecompCacheMutex);
        gCachedChunks[offset] = new DataChunk{ rgba, size };
        return gCachedChunks[offset];
    }
}

DecompressedData Decompressor::AutoDecode(YAML::Node& node, std::vector<uint8_t>& buffer,
                                          std::optional<size_t> manualSize) {
    auto offset = GetSafeNode<uint32_t>(node, "offset");

    CompressionType type = Companion::Instance->GetCurrCompressionType();

    auto fileOffset = TranslateAddr(offset, true);
    auto offsetFromFile = TranslateAddr(offset, false) - fileOffset;

    // Check if an asset in a yaml file is mio0 compressed and extract.
    if (node["mio0"]) {
        auto assetPtr = ASSET_PTR(offset);
        auto gameSize = Companion::Instance->GetRomData().size();

        auto decoded = Decode(buffer, fileOffset + offsetFromFile, CompressionType::MIO0);
        size_t decodedSize = decoded->size - offset;
        size_t size;

        if (node["size"]) {
            size = node["size"].as<size_t>();
        } else if (manualSize.has_value()) {
            size = manualSize.value();
        } else {
            size = decodedSize;
        }

        if (size > decodedSize) {
            SPDLOG_WARN("Requested size 0x{:X} exceeds decoded MIO0 asset size 0x{:X} at offset 0x{:X}. Reducing to "
                        "available size.",
                        size, decodedSize, assetPtr);
            size = decodedSize;
        }

        return { .root = decoded, .segment = { decoded->data, size } };
    }

    // Check if an asset in a yaml file is tkmk00 compressed and extract (mk64).
    if (node["tkmk00"]) {
        const auto alpha = GetSafeNode<uint32_t>(node, "alpha");
        const auto width = GetSafeNode<uint32_t>(node, "width");
        const auto height = GetSafeNode<uint32_t>(node, "height");
        const auto textureSize = width * height * 2;

        auto fileOffset = TranslateAddr(offset, true);
        offset = ASSET_PTR(offset);

        auto assetPtr = fileOffset + offset;
        auto decoded = DecodeTKMK00(buffer, assetPtr, textureSize, alpha);
        size_t decodedSize = decoded->size - offset;
        size_t size;

        if (node["size"]) {
            size = node["size"].as<size_t>();
        } else if (manualSize.has_value()) {
            size = manualSize.value();
        } else {
            size = decodedSize;
        }

        if (size > decodedSize) {
            SPDLOG_WARN("Requested size 0x{:X} exceeds decoded TKMK00 asset size 0x{:X} at offset 0x{:X}. Reducing to "
                        "available size.",
                        size, decodedSize, assetPtr);
            size = decodedSize;
        }
        return { .root = decoded, .segment = { decoded->data, size } };
    }
    if (node["bkzip"]) {
        const auto compressedSize = GetSafeNode<uint32_t>(node, "compressed_size");
        auto decoded = Decode(buffer, fileOffset + offsetFromFile, CompressionType::BKZIP, compressedSize);
        auto size = node["size"] ? node["size"].as<size_t>() : manualSize.value_or(decoded->size);
        return { .root = decoded, .segment = { decoded->data, size } };
    }

    // Extract a compressed file which contains many assets.
    switch (type) {
        case CompressionType::YAY0:
        case CompressionType::YAY1:
        case CompressionType::MIO0:
        case CompressionType::YAZ0: {
            offset = ASSET_PTR(offset);

            auto decoded = Decode(buffer, fileOffset, type);
            auto availableSize = decoded->size - offset;
            size_t size;

            if (node["size"]) {
                size = node["size"].as<size_t>();
            } else if (manualSize.has_value()) {
                size = manualSize.value();
            } else {
                size = availableSize;
            }

            if (size > availableSize) {
                SPDLOG_WARN("Requested size 0x{:X} exceeds decoded asset size 0x{:X} at offset 0x{:X}. Reducing to "
                            "available size.",
                            size, availableSize, fileOffset);
                size = availableSize;
            }

            return { .root = decoded, .segment = { decoded->data + offset, size } };
        }
        case CompressionType::BKZIP: {
            const auto sizeEntry = Companion::Instance->GetCurrentCompressedSize();

            if (!sizeEntry.has_value()) {
                throw std::runtime_error("Auto decode missing file compressed size.");
            }

            auto decoded = Decode(buffer, fileOffset, type, sizeEntry.value());
            // Rom data can reference offsets past the decoded file; clamp instead of
            // letting decoded->size - offsetFromFile underflow into a giant bogus segment.
            if (offsetFromFile > decoded->size) {
                SPDLOG_WARN("AutoDecode BKZIP: offset 0x{:X} lies past decoded size 0x{:X}; clamping to empty segment",
                            offsetFromFile, decoded->size);
                offsetFromFile = decoded->size;
            }
            auto size = node["size"] ? node["size"].as<size_t>() : manualSize.value_or(decoded->size - offsetFromFile);
            if (size > decoded->size - offsetFromFile) {
                SPDLOG_WARN("AutoDecode BKZIP: requested size 0x{:X} exceeds available 0x{:X}; reducing", size,
                            decoded->size - offsetFromFile);
                size = decoded->size - offsetFromFile;
            }
            SPDLOG_INFO("AutoDecode BKZIP: offset=0x{:X} fileOffset=0x{:X} offsetFromFile=0x{:X} compSize=0x{:X} "
                        "decodedSize=0x{:X} segSize=0x{:X}",
                        offset, fileOffset, offsetFromFile, sizeEntry.value(), decoded->size, size);
            return { .root = decoded, .segment = { decoded->data + offsetFromFile, size } };
        }
        case CompressionType::None: // The data does not have compression
        {
            fileOffset = TranslateAddr(offset, false);

            // An offset past the file must not underflow availableSize.
            if (fileOffset > buffer.size()) {
                SPDLOG_WARN("AutoDecode: offset 0x{:X} lies past file size 0x{:X}; clamping to empty segment",
                            fileOffset, buffer.size());
                fileOffset = buffer.size();
            }

            auto availableSize = buffer.size() - fileOffset;
            size_t size;

            if (node["size"]) {
                size = node["size"].as<size_t>();
            } else if (manualSize.has_value()) {
                size = manualSize.value();
            } else {
                size = availableSize;
            }

            if (size > availableSize) {
                SPDLOG_WARN("Requested size 0x{:X} exceeds available asset size 0x{:X} at offset 0x{:X}. Reducing to "
                            "available size.",
                            size, availableSize, fileOffset);
                size = availableSize;
            }

            return { .root = nullptr, .segment = { buffer.data() + fileOffset, size } };
        }
        default:
            throw std::runtime_error("Auto decode could not find a supported compression type.");
    }
}

DecompressedData Decompressor::AutoDecode(uint32_t offset, std::optional<size_t> size, std::vector<uint8_t>& buffer) {
    YAML::Node node;
    node["offset"] = offset;

    return AutoDecode(node, buffer, size);
}

uint32_t Decompressor::TranslateAddr(uint32_t addr, bool baseAddress) {
    if (IS_SEGMENTED(addr) || IS_VIRTUAL_SEGMENT(addr)) {
        const auto segment = Companion::Instance->GetFileOffsetFromSegmentedAddr(SEGMENT_NUMBER(addr));
        if (!segment.has_value()) {
            const auto compressedSegmentPair =
                Companion::Instance->GetFileOffsetFromCompressedSegmentedAddr(SEGMENT_NUMBER(addr));
            if (!compressedSegmentPair.has_value()) {
                SPDLOG_ERROR("Segment data missing from game config\nPlease add an entry for segment {}",
                             SEGMENT_NUMBER(addr));
                return 0;
            }
            return compressedSegmentPair.value().first +
                   (!baseAddress ? (compressedSegmentPair.value().second + SEGMENT_OFFSET(addr)) : 0);
        }

        return segment.value() + (!baseAddress ? SEGMENT_OFFSET(addr) : 0);
    }

    const auto vramEntry = Companion::Instance->GetCurrentVRAM();

    if (vramEntry.has_value()) {
        const auto vram = vramEntry.value();

        if (addr >= vram.addr) {
            return vram.offset + (!baseAddress ? (addr - vram.addr) : 0);
        }
    }

    return addr;
}

CompressionType Decompressor::GetCompressionType(std::vector<uint8_t>& buffer, const uint32_t offset) {
    if (offset) {
        LUS::BinaryReader reader((char*)buffer.data() + offset, sizeof(uint32_t));
        reader.SetEndianness(Torch::Endianness::Big);

        const std::string header = reader.ReadCString();

        // Check if a compressed header exists
        if (header == "MIO0") {
            return CompressionType::MIO0;
        }

        if (header == "Yay0" || header == "PERS") {
            return CompressionType::YAY0;
        }

        if (header == "Yay1") {
            return CompressionType::YAY1;
        }

        if (header == "Yaz0") {
            return CompressionType::YAZ0;
        }
    }
    return CompressionType::None;
}

bool Decompressor::IsSegmented(uint32_t addr) {
    if (IS_SEGMENTED(addr)) {
        const auto segment = Companion::Instance->GetFileOffsetFromSegmentedAddr(SEGMENT_NUMBER(addr));

        if (!segment.has_value()) {
            const auto compressedSegmentPair =
                Companion::Instance->GetFileOffsetFromCompressedSegmentedAddr(SEGMENT_NUMBER(addr));
            if (!compressedSegmentPair.has_value()) {
                SPDLOG_ERROR("Segment data missing from game config\nPlease add an entry for segment {}",
                             SEGMENT_NUMBER(addr));
                return false;
            }
        }
        return true;
    }

    return false;
}

void Decompressor::ClearCache() {
    std::lock_guard<std::mutex> lock(gDecompCacheMutex);
    for (auto& [key, value] : gCachedChunks) {
        delete[] value->data;
    }
    gCachedChunks.clear();
}