#include "terrain_splat.h"

#include "asset_io.h"
#include "dxt.h"

#include <godot_cpp/classes/image.hpp>
#include <godot_cpp/classes/image_texture.hpp>
#include <godot_cpp/classes/shader.hpp>
#include <godot_cpp/classes/texture2d_array.hpp>
#include <godot_cpp/variant/color.hpp>
#include <godot_cpp/variant/packed_byte_array.hpp>
#include <godot_cpp/variant/typed_array.hpp>

#include <asset_resolver.h>

#include <algorithm>
#include <unordered_map>

using namespace godot;

namespace mtgodot {

// 最多 16 层：A1 单区块实测最多 12 个活动图层（`000004`/`003004`），旧的 8 层上限会静默丢层。
static const int MAX_LAYERS = 16;

namespace {

const char *SRC_TERRAIN = R"(shader_type spatial;
render_mode diffuse_lambert, specular_disabled, cull_back;

uniform sampler2DArray layers : source_color, filter_linear_mipmap_anisotropic, repeat_enable;
uniform sampler2D weights0 : filter_linear;   // RGBA = layer 0..3 alpha
uniform sampler2D weights1 : filter_linear;   // RGBA = layer 4..7 alpha
uniform sampler2D weights2 : filter_linear;   // RGBA = layer 8..11 alpha
uniform sampler2D weights3 : filter_linear;   // RGBA = layer 12..15 alpha
uniform vec4 layer_uv[16];                     // xy = 每区块平铺频率 (8*Scale)，zw = offset
uniform int layer_count = 0;
uniform sampler2D shadowmap : source_color, filter_linear;
uniform bool use_shadowmap = false;

void fragment() {
	vec3 col = vec3(0.32, 0.30, 0.24);
	vec4 w[4];
	w[0] = texture(weights0, UV);
	w[1] = texture(weights1, UV);
	w[2] = texture(weights2, UV);
	w[3] = texture(weights3, UV);
	for (int i = 0; i < 16; i++) {
		if (i >= layer_count) { break; }
		float wi = w[i >> 2][i & 3];
		if (wi <= 0.003) { continue; }
		vec2 tuv = UV * layer_uv[i].xy + layer_uv[i].zw;
		vec3 lc = texture(layers, vec3(tuv, float(i))).rgb;
		col = mix(col, lc, wi);
	}
	if (use_shadowmap) {
		col *= texture(shadowmap, UV).rgb;
	}
	ALBEDO = col;
	ROUGHNESS = 1.0;
}
)";

Ref<Shader> g_terrain_shader;
Ref<Shader> terrain_shader() {
	if (g_terrain_shader.is_null()) {
		g_terrain_shader.instantiate();
		g_terrain_shader->set_code(SRC_TERRAIN);
	}
	return g_terrain_shader;
}

// DDS -> RGBA8 Image，resize 到 size×size。缓存（非函数静态 —— 见 cleanup）。
std::unordered_map<std::string, Ref<godot::Image>> g_layer_cache;
Ref<godot::Image> layer_image(const std::string &real_path, int size) {
	auto &cache = g_layer_cache;
	std::string key = real_path + "@" + std::to_string(size);
	auto it = cache.find(key);
	if (it != cache.end())
		return it->second;
	Ref<godot::Image> out;
	mtgodot::Image d = mtgodot::dds_from_file(godot::String(real_path.c_str()));
	if (d.ok()) {
		PackedByteArray b;
		b.resize((int64_t)d.rgba.size());
		std::copy(d.rgba.begin(), d.rgba.end(), b.ptrw());
		out = godot::Image::create_from_data(d.w, d.h, false, godot::Image::FORMAT_RGBA8, b);
		if (out.is_valid() && (d.w != size || d.h != size))
			out->resize(size, size, godot::Image::INTERPOLATE_BILINEAR);
		if (out.is_valid())
			out->generate_mipmaps();
	}
	cache.emplace(key, out);
	return out;
}

// SplatLayer.alpha (258²) -> 256² 的某个通道
void pack_channel(uint8_t *dst /*256*256*4*/, int ch, const std::vector<uint8_t> &alpha258) {
	const int S = fmt::SPLAT_RAW_XY; // 258
	for (int y = 0; y < 256; ++y)
		for (int x = 0; x < 256; ++x)
			dst[(y * 256 + x) * 4 + ch] = alpha258[size_t(y + 1) * S + (x + 1)];
}

Ref<ImageTexture> weight_tex(const std::vector<const fmt::SplatLayer *> &four) {
	std::vector<uint8_t> buf(size_t(256) * 256 * 4, 0);
	for (int c = 0; c < 4 && c < (int)four.size(); ++c)
		if (four[c])
			pack_channel(buf.data(), c, four[c]->alpha);
	PackedByteArray b;
	b.resize((int64_t)buf.size());
	std::copy(buf.begin(), buf.end(), b.ptrw());
	Ref<godot::Image> img =
			godot::Image::create_from_data(256, 256, false, godot::Image::FORMAT_RGBA8, b);
	return ImageTexture::create_from_image(img);
}

} // namespace

void cleanup_terrain_shader() {
	g_terrain_shader.unref();
	g_layer_cache.clear(); // 释放缓存的 Ref<Image>，别拖到 __cxa_finalize（那时引擎已析构）
}

Ref<ShaderMaterial> build_chunk_terrain_material(const fmt::SplatSet &splat,
		const fmt::TextureSet &tset, const fmt::AssetResolver &res,
		const String &shadowmap_path) {
	// Texture2DArray 要求各 slice 同尺寸 —— 取用到的图层里的最大源边长（上限 1024），
	// 只放大不缩小最大源，避免把 512² 地表贴图硬降采样（PARITY-GAP §3.4）。
	int src_max = 256;
	for (const auto &L : splat.layers) {
		if (L.layer >= 1 && L.layer <= (int)tset.layers.size()) {
			std::string rp = res.resolve(tset.layers[L.layer - 1].texture, nullptr);
			if (rp.empty())
				continue;
			mtgodot::Image d = mtgodot::dds_from_file(godot::String(rp.c_str()));
			if (d.ok())
				src_max = std::max<int>(src_max, std::max<int>(d.w, d.h));
		}
	}
	const int LSIZE = std::min(1024, src_max);

	int n = std::min<int>(MAX_LAYERS, (int)splat.layers.size());
	if (n == 0)
		return Ref<ShaderMaterial>();

	// 颜色数组
	TypedArray<godot::Image> imgs;
	std::vector<Color> uvparm(MAX_LAYERS, Color(40, 40, 0, 0));
	Ref<godot::Image> fallback;
	{
		PackedByteArray b;
		b.resize(LSIZE * LSIZE * 4);
		for (int i = 0; i < LSIZE * LSIZE * 4; i += 4) {
			b[i] = 90;
			b[i + 1] = 110;
			b[i + 2] = 70;
			b[i + 3] = 255;
		}
		fallback = godot::Image::create_from_data(LSIZE, LSIZE, false, godot::Image::FORMAT_RGBA8, b);
		fallback->generate_mipmaps();
	}
	for (int i = 0; i < n; ++i) {
		const fmt::SplatLayer &L = splat.layers[i];
		Ref<godot::Image> img;
		if (L.layer >= 1 && L.layer <= (int)tset.layers.size()) {
			const fmt::TextureLayer &tl = tset.layers[L.layer - 1];
			std::string rp = res.resolve(tl.texture, nullptr);
			if (!rp.empty())
				img = layer_image(rp, LSIZE);
			// 原客户端 TextureSet.cpp:185：u' = (TexCoordBase*UScale)*vtx_cm + UOffset，
			// TexCoordBase = 1/(PATCH_XSIZE*CELLSCALE) = 1/3200；区块归一化 UV -> 平铺频率 = 8*Scale。
			float us = tl.u_scale > 0.01f ? tl.u_scale : 1.0f;
			float vs = tl.v_scale > 0.01f ? tl.v_scale : 1.0f;
			uvparm[i] = Color(8.0f * us, -8.0f * vs, tl.u_offset, -tl.v_offset);
		}
		imgs.push_back(img.is_valid() ? img : fallback);
	}

	Ref<Texture2DArray> arr;
	arr.instantiate();
	arr->create_from_images(imgs);

	// 权重贴图：ceil(n/4) 张 RGBA8（每通道一层 alpha）
	Ref<ImageTexture> wtex[4];
	for (int g = 0; g < 4; ++g) {
		std::vector<const fmt::SplatLayer *> grp(4, nullptr);
		for (int k = 0; k < 4; ++k) {
			int li = g * 4 + k;
			if (li < n)
				grp[k] = &splat.layers[li];
		}
		wtex[g] = weight_tex(grp);
	}

	Ref<ShaderMaterial> mat;
	mat.instantiate();
	mat->set_shader(terrain_shader());
	mat->set_shader_parameter("layers", arr);
	mat->set_shader_parameter("weights0", wtex[0]);
	mat->set_shader_parameter("weights1", wtex[1]);
	mat->set_shader_parameter("weights2", wtex[2]);
	mat->set_shader_parameter("weights3", wtex[3]);
	mat->set_shader_parameter("layer_count", n);
	{
		Array uva;
		for (int i = 0; i < MAX_LAYERS; ++i)
			uva.push_back(Plane(uvparm[i].r, uvparm[i].g, uvparm[i].b, uvparm[i].a));
		mat->set_shader_parameter("layer_uv", uva);
	}

	if (!shadowmap_path.is_empty()) {
		mtgodot::Image sm = mtgodot::dds_from_file(shadowmap_path);
		if (sm.ok()) {
			PackedByteArray b;
			b.resize((int64_t)sm.rgba.size());
			std::copy(sm.rgba.begin(), sm.rgba.end(), b.ptrw());
			Ref<godot::Image> smi = godot::Image::create_from_data(
					sm.w, sm.h, false, godot::Image::FORMAT_RGBA8, b);
			mat->set_shader_parameter("shadowmap", ImageTexture::create_from_image(smi));
			mat->set_shader_parameter("use_shadowmap", true);
		}
	}
	return mat;
}

} // namespace mtgodot