#include "core/TileSet.h"

#include <cstdio>
#include <fstream>
#include <stdexcept>
#include <unordered_map>

#include "core/compressors/ApLib.h"
#include "core/compressors/PsGaiden.h"
#include "core/compressors/RawBpp.h"
#include "core/compressors/Stc0.h"
#include "core/compressors/Stc4.h"
#include "core/compressors/Stm.h"
#include "core/compressors/Zx7.h"

namespace bmp2tile
{

namespace
{

TilePixels ExtractBlock(const IndexedImage& image, int tileX, int tileY)
{
    TilePixels block{};
    for (int y = 0; y < TileSize; ++y)
    {
        for (int x = 0; x < TileSize; ++x)
        {
            const int srcX = tileX * TileSize + x;
            const int srcY = tileY * TileSize + y;
            uint8_t value = 0;
            if (srcX < image.Width && srcY < image.Height)
            {
                value = image.At(srcX, srcY);
            }
            block[static_cast<size_t>(y) * TileSize + x] = value;
        }
    }
    return block;
}

TilePixels FlipHorizontal(const TilePixels& t)
{
    TilePixels r{};
    for (int y = 0; y < TileSize; ++y)
    {
        for (int x = 0; x < TileSize; ++x)
        {
            r[static_cast<size_t>(y) * TileSize + x] = t[static_cast<size_t>(y) * TileSize + (TileSize - 1 - x)];
        }
    }
    return r;
}

TilePixels FlipVertical(const TilePixels& t)
{
    TilePixels r{};
    for (int y = 0; y < TileSize; ++y)
    {
        for (int x = 0; x < TileSize; ++x)
        {
            r[static_cast<size_t>(y) * TileSize + x] = t[static_cast<size_t>(TileSize - 1 - y) * TileSize + x];
        }
    }
    return r;
}

// FNV-1a over the 64 pixel bytes -- fine for a hash map key, no need for
// anything fancier here.
struct TilePixelsHash
{
    size_t operator()(const TilePixels& t) const
    {
        uint64_t h = 1469598103934665603ULL;
        for (uint8_t b : t)
        {
            h ^= b;
            h *= 1099511628211ULL;
        }
        return static_cast<size_t>(h);
    }
};

// Top+bottom 8x8 halves of one 8x16 sprite cell.
struct SpritePixels
{
    TilePixels Top;
    TilePixels Bottom;
    bool operator==(const SpritePixels& o) const { return Top == o.Top && Bottom == o.Bottom; }
};

struct SpritePixelsHash
{
    size_t operator()(const SpritePixels& s) const
    {
        uint64_t h = 1469598103934665603ULL;
        for (uint8_t b : s.Top)    { h ^= b; h *= 1099511628211ULL; }
        for (uint8_t b : s.Bottom) { h ^= b; h *= 1099511628211ULL; }
        return static_cast<size_t>(h);
    }
};

SpritePixels SpriteFlipH(const SpritePixels& s)
{
    return {FlipHorizontal(s.Top), FlipHorizontal(s.Bottom)};
}

// V-flip of an 8x16 sprite reverses all 16 rows vertically, so the two
// 8x8 halves swap AND each is individually flipped vertically.
SpritePixels SpriteFlipV(const SpritePixels& s)
{
    return {FlipVertical(s.Bottom), FlipVertical(s.Top)};
}

// Packs one tilemap cell into devkitSMS/SMSlib's native 2-byte PNT entry
// format: tile index in bits 0-8, TILE_FLIPPED_X in bit 9, TILE_FLIPPED_Y
// in bit 10 (SMSlib.h). Priority and sprite-palette bits (11-12) aren't
// modelled by TilemapEntry, so are always 0.
std::array<uint8_t, 2> PackTilemapEntry(const TilemapEntry& entry)
{
    uint16_t word = static_cast<uint16_t>(entry.TileIndex & 0x1FF);
    if (entry.bHFlip)            word |= 0x0200;
    if (entry.bVFlip)            word |= 0x0400;
    if (entry.bUseSpritePalette) word |= 0x0800;
    if (entry.bHighPriority)     word |= 0x1000;
    return {static_cast<uint8_t>(word & 0xFF), static_cast<uint8_t>((word >> 8) & 0xFF)};
}

// Flattens a list of packable items (tiles or tilemap entries) into one
// concatenated byte buffer, via `pack(item)` returning a fixed-size
// std::array. Used to build the input to CompressZx7, which operates on a
// flat byte stream rather than a tile/entry list.
template <typename T, typename PackFn>
std::vector<uint8_t> Flatten(const std::vector<T>& items, PackFn pack)
{
    std::vector<uint8_t> flat;
    if (!items.empty())
    {
        flat.reserve(items.size() * pack(items[0]).size());
    }
    for (const T& item : items)
    {
        const auto packed = pack(item);
        flat.insert(flat.end(), packed.begin(), packed.end());
    }
    return flat;
}

std::ofstream OpenForBinaryWrite(const std::string& path)
{
    std::ofstream f(path, std::ios::binary | std::ios::trunc);
    if (!f)
    {
        throw std::runtime_error("Couldn't open file for writing: " + path);
    }
    return f;
}

void WriteAll(std::ofstream& f, const std::vector<uint8_t>& data)
{
    if (!data.empty())
    {
        f.write(reinterpret_cast<const char*>(data.data()), static_cast<std::streamsize>(data.size()));
    }
}

// Defensive fallback for SaveTilesBin/SaveTilemapBin's switch statements.
// Every CompressionFormat is implemented and has an explicit case, so
// reaching this can only mean SupportsTiles/SupportsTilemap have drifted
// out of sync with those switches (e.g. a new format added to one but
// not the other) -- a real bug, not a "not implemented yet" gap.
[[noreturn]] void ThrowUnhandledFormat(CompressionFormat format)
{
    throw std::runtime_error(std::string(CompressionFormatName(format)) +
                              " has no case in this switch -- SupportsTiles/SupportsTilemap must be out of sync "
                              "with it.");
}

TileSet ExtractTiles8x16(const IndexedImage& image, bool removeDuplicates, bool useMirroring, bool useSpritePalette,
                          bool highPriority)
{
    TileSet result;
    result.TilesWide = (image.Width + TileSize - 1) / TileSize;
    result.TilesHigh = (image.Height + TileSize * 2 - 1) / (TileSize * 2);
    result.Tilemap.resize(static_cast<size_t>(result.TilesWide) * result.TilesHigh);

    std::unordered_map<SpritePixels, int, SpritePixelsHash> seen;

    for (int ty = 0; ty < result.TilesHigh; ++ty)
    {
        for (int tx = 0; tx < result.TilesWide; ++tx)
        {
            const SpritePixels sprite{ExtractBlock(image, tx, ty * 2), ExtractBlock(image, tx, ty * 2 + 1)};
            TilemapEntry entry;

            if (!removeDuplicates)
            {
                entry = TilemapEntry{static_cast<int>(result.Tiles.size()), false, false};
                result.Tiles.push_back(Tile{sprite.Top});
                result.Tiles.push_back(Tile{sprite.Bottom});
            }
            else if (auto it = seen.find(sprite); it != seen.end())
            {
                entry = TilemapEntry{it->second, false, false};
            }
            else
            {
                bool matched = false;
                if (useMirroring)
                {
                    if (auto it = seen.find(SpriteFlipH(sprite)); it != seen.end())
                    {
                        entry = TilemapEntry{it->second, true, false};
                        matched = true;
                    }
                    else if (auto it2 = seen.find(SpriteFlipV(sprite)); it2 != seen.end())
                    {
                        entry = TilemapEntry{it2->second, false, true};
                        matched = true;
                    }
                    else if (auto it3 = seen.find(SpriteFlipH(SpriteFlipV(sprite))); it3 != seen.end())
                    {
                        entry = TilemapEntry{it3->second, true, true};
                        matched = true;
                    }
                }
                if (!matched)
                {
                    entry = TilemapEntry{static_cast<int>(result.Tiles.size()), false, false};
                    result.Tiles.push_back(Tile{sprite.Top});
                    result.Tiles.push_back(Tile{sprite.Bottom});
                    seen.emplace(sprite, entry.TileIndex);
                }
            }

            result.Tilemap[static_cast<size_t>(ty) * result.TilesWide + tx] = entry;
        }
    }

    for (TilemapEntry& e : result.Tilemap)
    {
        e.bUseSpritePalette = useSpritePalette;
        e.bHighPriority = highPriority;
    }

    return result;
}

}  // namespace

std::array<uint8_t, 32> PackTilePlanar(const Tile& tile)
{
    std::array<uint8_t, 32> packed{};
    for (int y = 0; y < TileSize; ++y)
    {
        for (int plane = 0; plane < 4; ++plane)
        {
            uint8_t byte = 0;
            for (int x = 0; x < TileSize; ++x)
            {
                const uint8_t index = tile.Pixels[static_cast<size_t>(y) * TileSize + x];
                const uint8_t bit = (index >> plane) & 0x1;
                byte = static_cast<uint8_t>(byte | (bit << (TileSize - 1 - x)));
            }
            packed[static_cast<size_t>(y) * 4 + plane] = byte;
        }
    }
    return packed;
}

std::string FormatTilesAsSmsHex(const std::vector<Tile>& tiles)
{
    std::string out;
    char buf[8];
    for (size_t i = 0; i < tiles.size(); ++i)
    {
        std::snprintf(buf, sizeof(buf), "%03X", static_cast<int>(i));
        out += "; Tile index $";
        out += buf;
        out += "\n.db";
        for (const uint8_t b : PackTilePlanar(tiles[i]))
        {
            std::snprintf(buf, sizeof(buf), " $%02X", b);
            out += buf;
        }
        out += '\n';
    }
    return out;
}

std::string FormatTilemapAsSmsHex(const std::vector<TilemapEntry>& tilemap, int tilesWide)
{
    std::string out;
    char buf[8];
    const int total = static_cast<int>(tilemap.size());
    const int w = (tilesWide > 0) ? tilesWide : total;
    for (int i = 0; i < total; ++i)
    {
        if (i % w == 0)
        {
            if (i > 0) out += '\n';
            out += ".dw";
        }
        const auto packed = PackTilemapEntry(tilemap[static_cast<size_t>(i)]);
        const uint16_t word = static_cast<uint16_t>(packed[0]) | (static_cast<uint16_t>(packed[1]) << 8);
        std::snprintf(buf, sizeof(buf), " $%04X", word);
        out += buf;
    }
    if (!out.empty()) out += '\n';
    return out;
}

TileSet ExtractTiles(const IndexedImage& image, bool removeDuplicates, bool useMirroring, bool sprite8x16,
                     bool useSpritePalette, bool highPriority)
{
    if (sprite8x16)
        return ExtractTiles8x16(image, removeDuplicates, useMirroring, useSpritePalette, highPriority);

    TileSet result;
    result.TilesWide = (image.Width + TileSize - 1) / TileSize;
    result.TilesHigh = (image.Height + TileSize - 1) / TileSize;
    result.Tilemap.resize(static_cast<size_t>(result.TilesWide) * result.TilesHigh);

    // Maps a tile's exact pixel content -> its index in result.Tiles, so we
    // can recognise "we've seen this exact tile before" in O(1). When
    // mirroring is enabled we additionally probe this map with the
    // candidate's horizontal/vertical/both-axis flips.
    std::unordered_map<TilePixels, int, TilePixelsHash> seen;

    for (int ty = 0; ty < result.TilesHigh; ++ty)
    {
        for (int tx = 0; tx < result.TilesWide; ++tx)
        {
            const TilePixels block = ExtractBlock(image, tx, ty);
            TilemapEntry entry;

            if (!removeDuplicates)
            {
                entry = TilemapEntry{static_cast<int>(result.Tiles.size()), false, false};
                result.Tiles.push_back(Tile{block});
            }
            else if (auto it = seen.find(block); it != seen.end())
            {
                entry = TilemapEntry{it->second, false, false};
            }
            else
            {
                bool matched = false;
                if (useMirroring)
                {
                    if (auto it = seen.find(FlipHorizontal(block)); it != seen.end())
                    {
                        entry = TilemapEntry{it->second, true, false};
                        matched = true;
                    }
                    else if (auto it2 = seen.find(FlipVertical(block)); it2 != seen.end())
                    {
                        entry = TilemapEntry{it2->second, false, true};
                        matched = true;
                    }
                    else if (auto it3 = seen.find(FlipVertical(FlipHorizontal(block))); it3 != seen.end())
                    {
                        entry = TilemapEntry{it3->second, true, true};
                        matched = true;
                    }
                }
                if (!matched)
                {
                    entry = TilemapEntry{static_cast<int>(result.Tiles.size()), false, false};
                    result.Tiles.push_back(Tile{block});
                    seen.emplace(block, entry.TileIndex);
                }
            }

            result.Tilemap[static_cast<size_t>(ty) * result.TilesWide + tx] = entry;
        }
    }

    for (TilemapEntry& e : result.Tilemap)
    {
        e.bUseSpritePalette = useSpritePalette;
        e.bHighPriority = highPriority;
    }

    return result;
}

const char* CompressionFormatName(CompressionFormat format)
{
    switch (format)
    {
        case CompressionFormat::Uncompressed:
            return "Uncompressed";
        case CompressionFormat::STC0:
            return "STC0";
        case CompressionFormat::STC4:
            return "STC4";
        case CompressionFormat::PSGaiden:
            return "PSGaiden";
        case CompressionFormat::ZX7:
            return "ZX7";
        case CompressionFormat::APLib:
            return "aPLib";
        case CompressionFormat::Raw1Bpp:
            return "1bpp raw";
        case CompressionFormat::Raw2Bpp:
            return "2bpp raw";
        case CompressionFormat::STM:
            return "STM";
    }
    return "Unknown";
}

const char* CompressionFormatFileExtension(CompressionFormat format)
{
    switch (format)
    {
        case CompressionFormat::Uncompressed:
            return "bin";
        case CompressionFormat::STC0:
            return "stc0compr";
        case CompressionFormat::STC4:
            return "stc4compr";
        case CompressionFormat::PSGaiden:
            return "psgcompr";
        case CompressionFormat::ZX7:
            return "zx7";
        case CompressionFormat::APLib:
            return "aPLib";
        case CompressionFormat::Raw1Bpp:
            return "1bpp";
        case CompressionFormat::Raw2Bpp:
            return "2bpp";
        case CompressionFormat::STM:
            return "stmcompr";
    }
    return "bin";
}

bool SupportsTiles(CompressionFormat format)
{
    return format != CompressionFormat::STM;
}

bool SupportsTilemap(CompressionFormat format)
{
    // Confirmed against bmp2tilecompressors' actual plugin source: ZX7 and
    // aPLib both have a compressTilemap() that flat-compresses the
    // width*height*2 tilemap bytes, exactly like compressTiles() does for
    // tile data. SMSlib doesn't have a *direct* compressed-tilemap-to-VRAM
    // loader for either of them the way it does for STM -- only generic
    // decompress-to-RAM (SMS_decompressZX7/aPLib) followed by a normal
    // SMS_loadTileMap call -- but the file format itself is exactly as
    // valid and produceable as it is for tiles, so both are included here.
    return format == CompressionFormat::Uncompressed || format == CompressionFormat::STM ||
           format == CompressionFormat::ZX7 || format == CompressionFormat::APLib;
}

void SaveTilesBin(const std::string& path, const std::vector<Tile>& tiles, CompressionFormat format)
{
    if (!SupportsTiles(format))
    {
        throw std::runtime_error(std::string(CompressionFormatName(format)) + " doesn't support tile data.");
    }

    std::vector<uint8_t> output;
    switch (format)
    {
        case CompressionFormat::Uncompressed:
            output = Flatten(tiles, PackTilePlanar);
            break;
        case CompressionFormat::ZX7:
            output = CompressZx7(Flatten(tiles, PackTilePlanar));
            break;
        case CompressionFormat::PSGaiden:
            output = CompressPsGaidenTiles(tiles);
            break;
        case CompressionFormat::STC0:
            output = CompressStc0(Flatten(tiles, PackTilePlanar));
            break;
        case CompressionFormat::STC4:
            output = CompressStc4(Flatten(tiles, PackTilePlanar));
            break;
        case CompressionFormat::APLib:
            output = CompressApLib(Flatten(tiles, PackTilePlanar));
            break;
        case CompressionFormat::Raw1Bpp:
            output = CompressRaw1Bpp(Flatten(tiles, PackTilePlanar));
            break;
        case CompressionFormat::Raw2Bpp:
            output = CompressRaw2Bpp(Flatten(tiles, PackTilePlanar));
            break;
        default:
            ThrowUnhandledFormat(format);
    }

    std::ofstream f = OpenForBinaryWrite(path);
    WriteAll(f, output);
    if (!f)
    {
        throw std::runtime_error("Couldn't write file: " + path);
    }
}

std::vector<TilemapEntry> OffsetTilemapTileIndices(const std::vector<TilemapEntry>& tilemap, int offset)
{
    std::vector<TilemapEntry> result = tilemap;
    for (TilemapEntry& entry : result)
    {
        const int newIndex = entry.TileIndex + offset;
        if (newIndex < 0 || newIndex > 0x1FF)
        {
            throw std::runtime_error("Tile index " + std::to_string(entry.TileIndex) + " + offset " + std::to_string(offset) +
                                      " = " + std::to_string(newIndex) +
                                      ", which doesn't fit in the tilemap's 9-bit tile index field (0-511).");
        }
        entry.TileIndex = newIndex;
    }
    return result;
}

void SaveTilemapBin(const std::string& path, const std::vector<TilemapEntry>& tilemap, int tilesWide, CompressionFormat format,
                     int tileIndexOffset)
{
    if (!SupportsTilemap(format))
    {
        throw std::runtime_error(std::string(CompressionFormatName(format)) + " doesn't support tilemap data.");
    }

    const std::vector<TilemapEntry> offsetTilemap =
        (tileIndexOffset != 0) ? OffsetTilemapTileIndices(tilemap, tileIndexOffset) : tilemap;

    std::vector<uint8_t> output;
    switch (format)
    {
        case CompressionFormat::Uncompressed:
            output = Flatten(offsetTilemap, PackTilemapEntry);
            break;
        case CompressionFormat::ZX7:
            output = CompressZx7(Flatten(offsetTilemap, PackTilemapEntry));
            break;
        case CompressionFormat::STM:
            output = CompressStm(Flatten(offsetTilemap, PackTilemapEntry), tilesWide);
            break;
        case CompressionFormat::APLib:
            output = CompressApLib(Flatten(offsetTilemap, PackTilemapEntry));
            break;
        default:
            ThrowUnhandledFormat(format);
    }

    std::ofstream f = OpenForBinaryWrite(path);
    WriteAll(f, output);
    if (!f)
    {
        throw std::runtime_error("Couldn't write file: " + path);
    }
}

}  // namespace bmp2tile
