#include "texture_atlas.h"
#include "third_party/stb_image.h"

#include <algorithm>
#include <array>
#include <cstring>
#include <filesystem>
#include <stdexcept>
#include <string>
#include <vector>

namespace {

using namespace TextureAtlasConfig;

std::vector<BlockType> supportedBlocks() {
    return {
        BlockType::Grass,
        BlockType::Dirt,
        BlockType::Stone,
        BlockType::Sand,
        BlockType::Gravel,
        BlockType::Snow,
        BlockType::Wood,
        BlockType::Leaves,
        BlockType::Sandstone,
        BlockType::Ice,
        BlockType::Water
    };
}

std::string blockFileName(BlockType type) {
    switch (type) {
        case BlockType::Grass:     return "grass";
        case BlockType::Dirt:      return "dirt";
        case BlockType::Stone:     return "stone";
        case BlockType::Sand:      return "sand";
        case BlockType::Gravel:    return "gravel";
        case BlockType::Snow:      return "snow";
        case BlockType::Wood:      return "wood";
        case BlockType::Leaves:    return "leaves";
        case BlockType::Sandstone: return "sandstone";
        case BlockType::Ice:       return "ice";
        case BlockType::Water:     return "water";
        case BlockType::Air:
        default:                   return "unknown";
    }
}

std::pair<int, int> faceCell(CubeNetFace face) {
    switch (face) {
        case CubeNetFace::Top:    return {1, 0};
        case CubeNetFace::Left:   return {0, 1};
        case CubeNetFace::Front:  return {1, 1};
        case CubeNetFace::Right:  return {2, 1};
        case CubeNetFace::Back:   return {3, 1};
        case CubeNetFace::Bottom: return {1, 2};
        default:                  return {0, 0};
    }
}

void fillErrorTile(std::vector<std::uint8_t>& atlasPixels, int tileIndex) {
    const int tileX = tileIndex % AtlasColumns;
    const int tileY = tileIndex / AtlasColumns;

    const int dstX0 = tileX * FaceSize;
    const int dstY0 = tileY * FaceSize;

    for (int y = 0; y < FaceSize; ++y) {
        for (int x = 0; x < FaceSize; ++x) {
            const bool a = ((x / 4) + (y / 4)) % 2 == 0;
            const std::uint8_t r = a ? 255 : 0;
            const std::uint8_t g = 0;
            const std::uint8_t b = a ? 255 : 0;
            const std::uint8_t alpha = 255;

            const int dst = ((dstY0 + y) * AtlasWidth + (dstX0 + x)) * 4;
            atlasPixels[dst + 0] = r;
            atlasPixels[dst + 1] = g;
            atlasPixels[dst + 2] = b;
            atlasPixels[dst + 3] = alpha;
        }
    }
}

void blitFaceToAtlas(
    std::vector<std::uint8_t>& atlasPixels,
    const std::uint8_t* srcPixels,
    CubeNetFace face,
    int tileIndex
) {
    const auto [cellX, cellY] = faceCell(face);

    const int srcX0 = cellX * FaceSize;
    const int srcY0 = cellY * FaceSize;

    const int tileX = tileIndex % AtlasColumns;
    const int tileY = tileIndex / AtlasColumns;

    const int dstX0 = tileX * FaceSize;
    const int dstY0 = tileY * FaceSize;

    for (int y = 0; y < FaceSize; ++y) {
        for (int x = 0; x < FaceSize; ++x) {
            const int src = ((srcY0 + y) * BlockNetWidth + (srcX0 + x)) * 4;
            const int dst = ((dstY0 + y) * AtlasWidth + (dstX0 + x)) * 4;

            atlasPixels[dst + 0] = srcPixels[src + 0];
            atlasPixels[dst + 1] = srcPixels[src + 1];
            atlasPixels[dst + 2] = srcPixels[src + 2];
            atlasPixels[dst + 3] = srcPixels[src + 3];
        }
    }
}

} // namespace

int blockTextureOrderIndex(BlockType type) {
    switch (type) {
        case BlockType::Grass:     return 0;
        case BlockType::Dirt:      return 1;
        case BlockType::Stone:     return 2;
        case BlockType::Sand:      return 3;
        case BlockType::Gravel:    return 4;
        case BlockType::Snow:      return 5;
        case BlockType::Wood:      return 6;
        case BlockType::Leaves:    return 7;
        case BlockType::Sandstone: return 8;
        case BlockType::Ice:       return 9;
        case BlockType::Water:     return 10;
        case BlockType::Air:
        default:                   return -1;
    }
}

int atlasTileIndex(BlockType type, CubeNetFace face) {
    const int blockIndex = blockTextureOrderIndex(type);
    if (blockIndex < 0) {
        return -1;
    }
    return blockIndex * TextureAtlasConfig::FacesPerBlock + static_cast<int>(face);
}

TextureAtlasData TextureAtlasBuilder::buildFromFolder(const std::string& folder) {
    using namespace TextureAtlasConfig;

    TextureAtlasData atlas;
    atlas.width = AtlasWidth;
    atlas.height = AtlasHeight;
    atlas.pixels.resize(static_cast<std::size_t>(AtlasWidth * AtlasHeight * 4), 0);

    const auto blocks = supportedBlocks();

    for (BlockType type : blocks) {
        const std::string path = folder + "/" + blockFileName(type) + ".png";

        int w = 0;
        int h = 0;
        int channels = 0;
        stbi_uc* pixels = stbi_load(path.c_str(), &w, &h, &channels, STBI_rgb_alpha);

        if (!pixels) {
            for (int f = 0; f < FacesPerBlock; ++f) {
                fillErrorTile(atlas.pixels, atlasTileIndex(type, static_cast<CubeNetFace>(f)));
            }
            continue;
        }

        if (w != BlockNetWidth || h != BlockNetHeight) {
            stbi_image_free(pixels);
            for (int f = 0; f < FacesPerBlock; ++f) {
                fillErrorTile(atlas.pixels, atlasTileIndex(type, static_cast<CubeNetFace>(f)));
            }
            continue;
        }

        blitFaceToAtlas(atlas.pixels, pixels, CubeNetFace::Top,    atlasTileIndex(type, CubeNetFace::Top));
        blitFaceToAtlas(atlas.pixels, pixels, CubeNetFace::Left,   atlasTileIndex(type, CubeNetFace::Left));
        blitFaceToAtlas(atlas.pixels, pixels, CubeNetFace::Front,  atlasTileIndex(type, CubeNetFace::Front));
        blitFaceToAtlas(atlas.pixels, pixels, CubeNetFace::Right,  atlasTileIndex(type, CubeNetFace::Right));
        blitFaceToAtlas(atlas.pixels, pixels, CubeNetFace::Back,   atlasTileIndex(type, CubeNetFace::Back));
        blitFaceToAtlas(atlas.pixels, pixels, CubeNetFace::Bottom, atlasTileIndex(type, CubeNetFace::Bottom));

        stbi_image_free(pixels);
    }

    // Generate procedural item icons (1 tile each, starting at tile 66)
    {
        auto setPixel = [&](int tile, int px, int py, uint8_t r, uint8_t g, uint8_t b, uint8_t a=255) {
            if (px < 0 || px >= FaceSize || py < 0 || py >= FaceSize) return;
            int col = tile % AtlasColumns;
            int row = tile / AtlasColumns;
            int ax  = col * FaceSize + px;
            int ay  = row * FaceSize + py;
            int idx = (ay * AtlasWidth + ax) * 4;
            if (idx + 3 >= (int)atlas.pixels.size()) return;
            atlas.pixels[idx+0] = r; atlas.pixels[idx+1] = g;
            atlas.pixels[idx+2] = b; atlas.pixels[idx+3] = a;
        };

        // Stick (tile 66) — brown diagonal line, 2px thick
        for (int i = 0; i < 14; ++i) {
            setPixel(66, 12-i, i+1, 120, 60, 10);
            setPixel(66, 13-i, i+1, 120, 60, 10);
            setPixel(66, 12-i, i+2, 100, 50,  8);
        }

        // WoodPickaxe (tile 67) — 3 brown blocks on top row, handle diag
        uint8_t wr=160, wg=100, wb=40; // wood color
        // Pick head (top)
        for (int x=2;x<12;++x) for(int y=0;y<3;++y) setPixel(67,x,y,wr,wg,wb);
        // Handle going diagonal down-right
        for (int i=0;i<9;++i) setPixel(67,3+i,3+i,wr,wg,wb);
        for (int i=0;i<9;++i) setPixel(67,4+i,3+i,wr,wg,wb);

        // StonePickaxe (tile 68) — same shape, gray
        uint8_t sr=130,sg=130,sb=130;
        for (int x=2;x<12;++x) for(int y=0;y<3;++y) setPixel(68,x,y,sr,sg,sb);
        for (int i=0;i<9;++i) setPixel(68,3+i,3+i,sr,sg,sb);
        for (int i=0;i<9;++i) setPixel(68,4+i,3+i,sr,sg,sb);

        // WoodAxe (tile 69) — axe head + handle
        for (int x=1;x<8;++x) for(int y=0;y<6;++y) setPixel(69,x,y,wr,wg,wb);
        for (int i=0;i<8;++i) setPixel(69,5+i,5+i,wr,wg,wb);
        for (int i=0;i<8;++i) setPixel(69,6+i,5+i,wr,wg,wb);

        // StoneAxe (tile 70) — same shape, gray
        for (int x=1;x<8;++x) for(int y=0;y<6;++y) setPixel(70,x,y,sr,sg,sb);
        for (int i=0;i<8;++i) setPixel(70,5+i,5+i,sr,sg,sb);
        for (int i=0;i<8;++i) setPixel(70,6+i,5+i,sr,sg,sb);

        // WoodSword (tile 71) — two wood blocks stacked, stick handle below
        for (int y=0;y<4;++y) for(int x=5;x<11;++x) setPixel(71,x,y,wr,wg,wb);   // blade body
        for (int y=4;y<7;++y) for(int x=6;x<10;++x) setPixel(71,x,y,wr,wg,wb);   // blade tip
        for (int i=0;i<6;++i) { setPixel(71,7,7+i,120,60,10); setPixel(71,8,7+i,120,60,10); } // handle

        // StoneSword (tile 72) — same shape, gray
        for (int y=0;y<4;++y) for(int x=5;x<11;++x) setPixel(72,x,y,sr,sg,sb);
        for (int y=4;y<7;++y) for(int x=6;x<10;++x) setPixel(72,x,y,sr,sg,sb);
        for (int i=0;i<6;++i) { setPixel(72,7,7+i,120,60,10); setPixel(72,8,7+i,120,60,10); }
    }

    return atlas;
}