#include "water_builder.h"
#include "water_motion.h"

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

#include <godot_cpp/classes/image.hpp>
#include <godot_cpp/classes/shader.hpp>
#include <godot_cpp/classes/shader_material.hpp>
#include <godot_cpp/classes/texture2d_array.hpp>
#include <godot_cpp/classes/time.hpp>
#include <godot_cpp/variant/utility_functions.hpp>
#include <godot_cpp/variant/packed_byte_array.hpp>
#include <godot_cpp/variant/packed_color_array.hpp>
#include <godot_cpp/variant/packed_int32_array.hpp>
#include <godot_cpp/variant/packed_vector2_array.hpp>
#include <godot_cpp/variant/packed_vector3_array.hpp>
#include <godot_cpp/variant/typed_array.hpp>

#include <asset_resolver.h>
#include <m2_coord.h>
#include <terrain_mesh.h> // terrain_height_at

#include <algorithm>

using namespace godot;

namespace mtgodot {

namespace {

// 30 帧序列 + 逐顶点水深 alpha(顶点 COLOR.a) + 轻微高度浮动。
// UV = 世界米；平铺频率在 shader 里按 1/(CELLSCALE*4 cm) = 1/8m。
const char *SRC_WATER = R"(shader_type spatial;
render_mode blend_mix, cull_disabled, depth_draw_never, unshaded;

uniform sampler2DArray frames : source_color, filter_linear_mipmap, repeat_enable;
uniform float uv_per_meter = 0.125;   // 1/8m，= 原客户端 1/(CELLSCALE*4)
uniform float height_offset = 0.0; // Shared reference water translation, metres.

void vertex() {
	VERTEX.y += height_offset;
}

void fragment() {
	int f = int(mod(TIME * 1000.0 / 70.0, 30.0)); // 70ms/帧
	vec2 uv = UV * uv_per_meter;
	vec3 tex = texture(frames, vec3(uv, float(f))).rgb;
	// MapOutdoorWater selects texture RGB and vertex diffuse alpha directly.
	// Do not introduce PBR lighting, invented tint or view-dependent opacity.
	ALBEDO = tex;
	ALPHA = COLOR.a;
}
)";

Ref<Shader> g_water_shader;
Ref<ShaderMaterial> g_water_mat; // 30 帧数组只建一次
WaterMotion g_water_motion;

Ref<ShaderMaterial> water_material(const fmt::AssetResolver &res) {
	if (g_water_mat.is_valid())
		return g_water_mat;
	if (g_water_shader.is_null()) {
		g_water_shader.instantiate();
		g_water_shader->set_code(SRC_WATER);
	}
	Ref<ShaderMaterial> m;
	m.instantiate();
	m->set_shader(g_water_shader);

	// special/water/01..30.dds
	TypedArray<godot::Image> imgs;
	int W = 0, H = 0;
	for (int i = 1; i <= 30; ++i) {
		char nm[64];
		std::snprintf(nm, sizeof(nm), "d:/ymir work/special/water/%02d.dds", i);
		std::string rp = res.resolve(nm, nullptr);
		mtgodot::Image d = rp.empty() ? mtgodot::Image{} : mtgodot::dds_from_file(godot::String(rp.c_str()));
		if (!d.ok())
			break;
		if (W == 0) {
			W = d.w;
			H = d.h;
		}
		PackedByteArray b;
		b.resize((int64_t)d.rgba.size());
		std::copy(d.rgba.begin(), d.rgba.end(), b.ptrw());
		Ref<godot::Image> img =
				godot::Image::create_from_data(d.w, d.h, false, godot::Image::FORMAT_RGBA8, b);
		if (img.is_valid() && (d.w != W || d.h != H))
			img->resize(W, H, godot::Image::INTERPOLATE_BILINEAR);
		if (img.is_valid())
			img->generate_mipmaps();
		imgs.push_back(img);
	}
	if (imgs.size() == 30) {
		Ref<Texture2DArray> arr;
		arr.instantiate();
		arr->create_from_images(imgs);
		m->set_shader_parameter("frames", arr);
	}
	g_water_mat = m;
	return m;
}

} // namespace

void cleanup_water_shader() {
	g_water_shader.unref();
	g_water_mat.unref();
	g_water_motion = WaterMotion{};
}

void update_water_animation() {
	if (g_water_mat.is_null()) return;
	const uint64_t now = Time::get_singleton()->get_ticks_msec();
	uint64_t duration = 1000;
	double depth = 0;
	if (g_water_motion.initialized && now - g_water_motion.start_ms > g_water_motion.duration_ms) {
		duration = UtilityFunctions::randi_range(1000, 3000);
		if (g_water_motion.end == 0) depth = UtilityFunctions::randi_range(0, 15);
	}
	g_water_mat->set_shader_parameter("height_offset", g_water_motion.sample(now, duration, depth));
}

std::vector<WaterPiece> build_chunk_water(const fmt::WaterMap &wm, const fmt::HeightMap &hm,
		int tile_x, int tile_y, double height_scale, const fmt::AssetResolver &res) {
	std::vector<WaterPiece> out;
	if (wm.layer_count == 0 || wm.ids.size() != size_t(fmt::WATERMAP_XY) * fmt::WATERMAP_XY)
		return out;

	const int W = fmt::WATERMAP_XY; // 128
	const double CELL_M = double(fmt::m2coord::CELLSCALE) * fmt::m2coord::CM_TO_M; // 2m/texel
	const double X0 = double(tile_x) * fmt::m2coord::CHUNK_CM * fmt::m2coord::CM_TO_M;
	const double Z0 = double(tile_y) * fmt::m2coord::CHUNK_CM * fmt::m2coord::CM_TO_M;

	auto depth_alpha = [&](int tex_x, int tex_y, double water_h_cm) -> float {
		// 水 texel (tex_x,tex_y) 对应的区块本地 cm（texel = 1 格 = CELLSCALE）
		double lx = tex_x * double(fmt::m2coord::CELLSCALE);
		double ly = tex_y * double(fmt::m2coord::CELLSCALE);
		double th = fmt::terrain_height_at(hm, lx, ly, height_scale); // cm
		double depth_cm = water_h_cm - th;
		// AreaTerrain uses RAW height differences, before HeightScale, with
		// OpaqueWaterDepth=400 and clamps at 80% (not a 12% shoreline floor).
		return water_depth_alpha(depth_cm, height_scale);
	};

	for (int layer = 0; layer < wm.layer_count; ++layer) {
		double h_cm = double(layer < (int)wm.heights.size() ? wm.heights[layer] : 0) * height_scale;
		float gy = float(h_cm * fmt::m2coord::CM_TO_M);

		PackedVector3Array v;
		PackedVector3Array n;
		PackedVector2Array uv;
		PackedColorArray col;
		PackedInt32Array idx;
		for (int y = 0; y < W; ++y) {
			int x = 0;
			while (x < W) {
				if (wm.ids[y * W + x] != layer) {
					++x;
					continue;
				}
				int xs = x;
				// Keep each cell's four shoreline samples. Merging a whole row
				// discarded interior terrain peaks and bridged dry banks with water.
				++x;
				float aLL = depth_alpha(xs, y, h_cm), aLR = depth_alpha(x, y, h_cm);
				float aUR = depth_alpha(x, y + 1, h_cm), aUL = depth_alpha(xs, y + 1, h_cm);
				// AreaTerrain skips the quad when every diffuse alpha is zero.
				if (aLL == 0 && aLR == 0 && aUR == 0 && aUL == 0)
					continue;
				double x0 = X0 + xs * CELL_M, x1 = X0 + x * CELL_M;
				double z0 = Z0 + y * CELL_M, z1 = Z0 + (y + 1) * CELL_M;
				int base = v.size();
				v.push_back(Vector3(x0, gy, z0));
				v.push_back(Vector3(x1, gy, z0));
				v.push_back(Vector3(x1, gy, z1));
				v.push_back(Vector3(x0, gy, z1));
				for (int k = 0; k < 4; ++k)
					n.push_back(Vector3(0, 1, 0));
				uv.push_back(Vector2(float(x0), float(z0)));
				uv.push_back(Vector2(float(x1), float(z0)));
				uv.push_back(Vector2(float(x1), float(z1)));
				uv.push_back(Vector2(float(x0), float(z1)));
				col.push_back(Color(1, 1, 1, aLL));
				col.push_back(Color(1, 1, 1, aLR));
				col.push_back(Color(1, 1, 1, aUR));
				col.push_back(Color(1, 1, 1, aUL));
				idx.push_back(base);
				idx.push_back(base + 2);
				idx.push_back(base + 1);
				idx.push_back(base);
				idx.push_back(base + 3);
				idx.push_back(base + 2);
			}
		}
		if (v.is_empty())
			continue;

		Array arr;
		arr.resize(Mesh::ARRAY_MAX);
		arr[Mesh::ARRAY_VERTEX] = v;
		arr[Mesh::ARRAY_NORMAL] = n;
		arr[Mesh::ARRAY_TEX_UV] = uv;
		arr[Mesh::ARRAY_COLOR] = col;
		arr[Mesh::ARRAY_INDEX] = idx;

		Ref<godot::ArrayMesh> mesh;
		mesh.instantiate();
		mesh->add_surface_from_arrays(Mesh::PRIMITIVE_TRIANGLES, arr);
		mesh->surface_set_material(0, water_material(res));
		out.push_back({mesh, gy});
	}
	return out;
}

} // namespace mtgodot