#include "adt_baker.h"
#include "blp_loader.h"

#include <godot_cpp/core/class_db.hpp>
#include <godot_cpp/variant/array.hpp>
#include <godot_cpp/variant/utility_functions.hpp>

#include <algorithm>
#include <cmath>

using namespace godot;

namespace {

constexpr float FALLBACK_R = 0.4f;
constexpr float FALLBACK_G = 0.55f;
constexpr float FALLBACK_B = 0.3f;

static float clamp01(float v) {
    return std::max(0.0f, std::min(1.0f, v));
}

static int32_t positive_mod(int32_t value, int32_t mod) {
    int32_t result = value % mod;
    return result < 0 ? result + mod : result;
}

static float wrap01(float v) {
    return v - std::floor(v);
}

} // namespace

Ref<Image> ADTBaker::bake_tile_albedo(
    const Dictionary &data,
    Dictionary image_cache,
    const String &extracted_dir,
    int32_t texture_size) {

    const int32_t size = get_baked_albedo_size(texture_size);
    const int32_t chunk_pixels = size / 16;
    if (size <= 0 || chunk_pixels <= 0) {
        return Ref<Image>();
    }

    PackedByteArray out;
    out.resize(size * size * 3);
    uint8_t *out_ptr = out.ptrw();
    for (int32_t i = 0; i < size * size; ++i) {
        out_ptr[i * 3 + 0] = static_cast<uint8_t>(FALLBACK_R * 255.0f);
        out_ptr[i * 3 + 1] = static_cast<uint8_t>(FALLBACK_G * 255.0f);
        out_ptr[i * 3 + 2] = static_cast<uint8_t>(FALLBACK_B * 255.0f);
    }

    const PackedStringArray tex_names = data.get("textures", PackedStringArray());
    const Array chunks = data.get("chunks", Array());

    for (int64_t ci = 0; ci < chunks.size(); ++ci) {
        const Dictionary chunk = chunks[ci];
        if (chunk.is_empty()) {
            continue;
        }

        const int32_t chunk_x = int32_t(chunk.get("index_x", -1));
        const int32_t chunk_y = int32_t(chunk.get("index_y", -1));
        if (chunk_x < 0 || chunk_x >= 16 || chunk_y < 0 || chunk_y >= 16) {
            continue;
        }

        const Array layers = chunk.get("layers", Array());
        if (layers.is_empty()) {
            continue;
        }

        TextureData textures[4];
        const int32_t layer_count = std::min<int32_t>(int32_t(layers.size()), 4);
        for (int32_t li = 0; li < layer_count; ++li) {
            const Dictionary layer = layers[li];
            const int32_t tex_id = int32_t(layer.get("texture_id", -1));
            if (tex_id >= 0 && tex_id < tex_names.size()) {
                textures[li] = load_texture_data(tex_names[tex_id], image_cache, extracted_dir);
            }
        }
        if (!textures[0].valid) {
            continue;
        }

        const Array alpha_maps = chunk.get("alpha_maps", Array());
        const int32_t base_x = chunk_x * chunk_pixels;
        const int32_t base_y = chunk_y * chunk_pixels;

        for (int32_t py = 0; py < chunk_pixels; ++py) {
            const float local_v = (float(py) + 0.5f) / float(chunk_pixels);
            const int32_t ay = std::max(0, std::min(63, int32_t(std::floor(local_v * 63.999f))));

            for (int32_t px = 0; px < chunk_pixels; ++px) {
                const float local_u = (float(px) + 0.5f) / float(chunk_pixels);
                const int32_t ax = std::max(0, std::min(63, int32_t(std::floor(local_u * 63.999f))));

                float weights[4] = {1.0f, 0.0f, 0.0f, 0.0f};
                for (int32_t li = 1; li < layer_count; ++li) {
                    PackedByteArray alpha;
                    if (li - 1 < alpha_maps.size()) {
                        alpha = alpha_maps[li - 1];
                    }
                    if (alpha.size() == 64 * 64) {
                        weights[li] = float(alpha[ay * 64 + ax]) / 255.0f;
                    }
                }

                weights[0] = std::max(0.0f, 1.0f - weights[1] - weights[2] - weights[3]);
                const float sum = weights[0] + weights[1] + weights[2] + weights[3];
                if (sum > 0.0f) {
                    for (float &weight : weights) {
                        weight /= sum;
                    }
                }

                const float tiled_u = local_u * 8.0f;
                const float tiled_v = local_v * 8.0f;
                float rgb[3] = {0.0f, 0.0f, 0.0f};
                for (int32_t li = 0; li < layer_count; ++li) {
                    if (!textures[li].valid || weights[li] <= 0.0f) {
                        continue;
                    }
                    float sample[3];
                    sample_image_repeat(textures[li], tiled_u, tiled_v, sample);
                    rgb[0] += sample[0] * weights[li];
                    rgb[1] += sample[1] * weights[li];
                    rgb[2] += sample[2] * weights[li];
                }

                const int32_t out_index = ((base_y + py) * size + (base_x + px)) * 3;
                out_ptr[out_index + 0] = static_cast<uint8_t>(clamp01(rgb[0]) * 255.0f + 0.5f);
                out_ptr[out_index + 1] = static_cast<uint8_t>(clamp01(rgb[1]) * 255.0f + 0.5f);
                out_ptr[out_index + 2] = static_cast<uint8_t>(clamp01(rgb[2]) * 255.0f + 0.5f);
            }
        }
    }

    return Image::create_from_data(size, size, false, Image::FORMAT_RGB8, out);
}

int32_t ADTBaker::get_baked_albedo_size(int32_t texture_size) {
    const int32_t safe_size = std::max(texture_size, 16);
    const int32_t chunk_pixels = std::max(8, int32_t(std::ceil(float(safe_size) / 16.0f)));
    return chunk_pixels * 16;
}

String ADTBaker::normalize_rel_path(const String &path) {
    return path.replace("\\", "/");
}

ADTBaker::TextureData ADTBaker::load_texture_data(
    const String &blp_path,
    Dictionary &image_cache,
    const String &extracted_dir) {

    TextureData texture;
    if (blp_path.is_empty()) {
        return texture;
    }

    Ref<Image> image;
    if (image_cache.has(blp_path)) {
        image = image_cache[blp_path];
    } else {
        const String rel_path = normalize_rel_path(blp_path);
        String abs_path = extracted_dir;
        if (!abs_path.ends_with("/") && !abs_path.ends_with("\\")) {
            abs_path += "/";
        }
        abs_path += rel_path;

        Ref<BLPLoader> loader;
        loader.instantiate();
        image = loader->load_image(abs_path);
        image_cache[blp_path] = image;
    }

    if (image.is_null() || image->is_empty()) {
        return texture;
    }

    if (image->get_format() != Image::FORMAT_RGBA8) {
        image = image->duplicate();
        image->convert(Image::FORMAT_RGBA8);
    }

    texture.image = image;
    texture.width = image->get_width();
    texture.height = image->get_height();
    texture.bytes = image->get_data();
    texture.valid = texture.width > 0 && texture.height > 0 && texture.bytes.size() >= texture.width * texture.height * 4;
    return texture;
}

void ADTBaker::sample_image_repeat(
    const TextureData &texture,
    float u,
    float v,
    float out_rgb[3]) {

    if (!texture.valid) {
        out_rgb[0] = 1.0f;
        out_rgb[1] = 1.0f;
        out_rgb[2] = 1.0f;
        return;
    }

    const float fu = wrap01(u);
    const float fv = wrap01(v);
    const float px = fu * float(texture.width) - 0.5f;
    const float py = fv * float(texture.height) - 0.5f;
    const int32_t x0 = positive_mod(int32_t(std::floor(px)), texture.width);
    const int32_t y0 = positive_mod(int32_t(std::floor(py)), texture.height);
    const int32_t x1 = positive_mod(x0 + 1, texture.width);
    const int32_t y1 = positive_mod(y0 + 1, texture.height);
    const float tx = px - std::floor(px);
    const float ty = py - std::floor(py);

    const uint8_t *bytes = texture.bytes.ptr();
    const int32_t i00 = (y0 * texture.width + x0) * 4;
    const int32_t i10 = (y0 * texture.width + x1) * 4;
    const int32_t i01 = (y1 * texture.width + x0) * 4;
    const int32_t i11 = (y1 * texture.width + x1) * 4;

    for (int32_t c = 0; c < 3; ++c) {
        const float c00 = float(bytes[i00 + c]) / 255.0f;
        const float c10 = float(bytes[i10 + c]) / 255.0f;
        const float c01 = float(bytes[i01 + c]) / 255.0f;
        const float c11 = float(bytes[i11 + c]) / 255.0f;
        const float cx0 = c00 + (c10 - c00) * tx;
        const float cx1 = c01 + (c11 - c01) * tx;
        out_rgb[c] = cx0 + (cx1 - cx0) * ty;
    }
}

void ADTBaker::_bind_methods() {
    ClassDB::bind_method(
        D_METHOD("bake_tile_albedo", "data", "image_cache", "extracted_dir", "texture_size"),
        &ADTBaker::bake_tile_albedo);
}
