#include "tree_placeholder.h"

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

#include <asset_resolver.h>
#include <spt.h>

#include <godot_cpp/classes/image.hpp>
#include <godot_cpp/classes/image_texture.hpp>
#include <godot_cpp/classes/shader.hpp>
#include <godot_cpp/classes/shader_material.hpp>
#include <godot_cpp/classes/standard_material3d.hpp>
#include <godot_cpp/variant/packed_byte_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 <algorithm>
#include <cctype>
#include <cmath>
#include <cstdint>
#include <unordered_map>
#include <vector>

using namespace godot;

namespace mtgodot {

namespace {

constexpr float kPI = 3.14159265358979323846f;
constexpr float kTAU = 2.0f * kPI;

// SpeedTree 2 的叶片是中心点 + leaf-cluster table，在顶点 shader 中展开。
// proxy 已在 CPU 侧把叶簇展开成 card；这里只保留 alpha-test 和轻微、按实例错相的风摆。
const char *SRC_LEAF = R"(shader_type spatial;
render_mode cull_disabled, diffuse_lambert, specular_disabled, depth_prepass_alpha;
uniform sampler2D leaf_tex : source_color, filter_linear_mipmap_anisotropic;
uniform float wind_strength = 1.0;
void vertex() {
	vec3 wp = (MODEL_MATRIX * vec4(0.0, 0.0, 0.0, 1.0)).xyz;
	float ph = wp.x * 0.11 + wp.z * 0.13;
	float h = max(VERTEX.y, 0.0);
	VERTEX.x += sin(TIME * 1.3 + ph) * 0.025 * wind_strength * h;
	VERTEX.z += cos(TIME * 1.05 + ph) * 0.018 * wind_strength * h;
}
void fragment() {
	vec4 c = texture(leaf_tex, UV);
	if (c.a < 0.38) { discard; }
	ALBEDO = c.rgb;
	ROUGHNESS = 1.0;
}
)";

Ref<Shader> g_leaf_shader;
Ref<ImageTexture> g_fallback_broadleaf;
Ref<ImageTexture> g_fallback_conifer;
std::unordered_map<std::string, Ref<ImageTexture>> g_tree_texture_cache;
std::unordered_map<std::string, Ref<ArrayMesh>> g_tree_mesh_cache;

Ref<Shader> leaf_shader() {
	if (g_leaf_shader.is_null()) {
		g_leaf_shader.instantiate();
		g_leaf_shader->set_code(SRC_LEAF);
	}
	return g_leaf_shader;
}

Ref<ImageTexture> fallback_leaf_texture(bool conifer) {
	Ref<ImageTexture> &cached = conifer ? g_fallback_conifer : g_fallback_broadleaf;
	if (cached.is_valid())
		return cached;
	const int N = 96;
	PackedByteArray b;
	b.resize(N * N * 4);
	for (int y = 0; y < N; ++y) {
		for (int x = 0; x < N; ++x) {
			const float u = (x + 0.5f) / N * 2.0f - 1.0f;
			const float v = (y + 0.5f) / N * 2.0f - 1.0f;
			const float d = std::sqrt(u * u + v * v);
			float a = 1.0f - d;
			a = a <= 0 ? 0.0f : a * a * (3.0f - 2.0f * a);
			const float n = 0.5f + 0.5f * std::sin(x * 0.9f) * std::sin(y * 0.7f);
			const float g = conifer ? (0.28f + 0.14f * n) : (0.40f + 0.16f * n);
			const float r = conifer ? (0.11f + 0.06f * n) : (0.18f + 0.10f * n);
			const float bl = 0.10f + 0.06f * n;
			const int o = (y * N + x) * 4;
			b[o + 0] = uint8_t(std::min(255.0f, r * 255.0f));
			b[o + 1] = uint8_t(std::min(255.0f, g * 255.0f));
			b[o + 2] = uint8_t(std::min(255.0f, bl * 255.0f));
			b[o + 3] = uint8_t(std::min(255.0f, a * 255.0f));
		}
	}
	Ref<godot::Image> img =
			godot::Image::create_from_data(N, N, false, godot::Image::FORMAT_RGBA8, b);
	img->generate_mipmaps();
	cached = ImageTexture::create_from_image(img);
	return cached;
}

Ref<ImageTexture> load_dds_texture(const std::string &path) {
	if (path.empty())
		return Ref<ImageTexture>();
	auto found = g_tree_texture_cache.find(path);
	if (found != g_tree_texture_cache.end())
		return found->second;
	Ref<ImageTexture> result;
	mtgodot::Image d = mtgodot::dds_from_file(godot::String(path.c_str()));
	if (d.ok()) {
		// Bark / leaf-composite albedo (sRGB) -> mobile ASTC via
		// make_color_texture (no-op on desktop; keeps mipmaps). §F4.
		result = mtgodot::make_color_texture(d.w, d.h, d.rgba.data(), d.rgba.size(),
				/*mipmaps=*/true);
	}
	g_tree_texture_cache.emplace(path, result);
	return result;
}

std::string lower(std::string s) {
	std::transform(s.begin(), s.end(), s.begin(),
			[](unsigned char c) { return (char)std::tolower(c); });
	return s;
}

std::string basename(std::string path) {
	for (char &c : path)
		if (c == '\\') c = '/';
	const size_t slash = path.find_last_of('/');
	return slash == std::string::npos ? path : path.substr(slash + 1);
}

std::string as_dds(std::string path) {
	const size_t dot = path.find_last_of('.');
	if (dot != std::string::npos)
		path.resize(dot);
	return path + ".dds";
}

std::string resolve_sibling(const std::string &treefile, const std::string &texture,
		const fmt::AssetResolver &resolver) {
	if (texture.empty())
		return "";
	std::string parent = fmt::AssetResolver::normalize(treefile);
	const size_t slash = parent.find_last_of('/');
	if (slash != std::string::npos)
		parent.resize(slash);
	else
		parent.clear();
	const std::string name = as_dds(basename(texture));
	return resolver.resolve(parent.empty() ? name : parent + "/" + name, nullptr);
}

struct TreeTextures {
	Ref<ImageTexture> bark;
	Ref<ImageTexture> composite;
	std::string composite_name;
};

TreeTextures resolve_tree_textures(const std::string &treefile,
		const fmt::AssetResolver &resolver) {
	TreeTextures out;
	const std::string spt_path = resolver.resolve(treefile, nullptr);
	if (spt_path.empty())
		return out;
	fmt::SptInfo info;
	if (!fmt::sniff_spt_file(spt_path, info))
		return out;
	std::string branch;
	for (const std::string &ref : info.texture_refs) {
		if (lower(ref).find("bark") != std::string::npos) {
			branch = ref;
			break;
		}
	}
	if (branch.empty() && !info.texture_refs.empty())
		branch = info.texture_refs.front();
	out.bark = load_dds_texture(resolve_sibling(treefile, branch, resolver));
	out.composite_name = info.composite_texture;
	out.composite = load_dds_texture(
			resolve_sibling(treefile, info.composite_texture, resolver));
	return out;
}

struct UVRect {
	float u0 = 0, v0 = 0, u1 = 1, v1 = 1;
};

std::vector<UVRect> foliage_rects(const std::string &species,
		const std::string &composite, bool atlas) {
	if (!atlas)
		return {{0, 0, 1, 1}};
	// SPT leaf-cluster UV 尚未导出；这些区域只选择各 composite atlas 中的真实叶簇，
	// 不声称复原了具体树种的原始 UV。Windows exporter 接入后删除这组 proxy 布局。
	const std::string s = lower(species);
	const std::string c = lower(composite);
	const bool fall = s.find("fall") != std::string::npos;
	const bool winter = s.find("winter") != std::string::npos;
	if (c.find("b1") != std::string::npos) {
		if (fall) return {{0.00f, 0.05f, 0.25f, 0.25f}, {0.25f, 0.25f, 0.50f, 0.50f}};
		return {{0.25f, 0.02f, 0.50f, 0.23f}, {0.25f, 0.18f, 0.50f, 0.36f},
				{0.00f, 0.27f, 0.27f, 0.49f}};
	}
	if (c.find("b2") != std::string::npos) {
		if (fall) return {{0.00f, 0.00f, 0.25f, 0.25f}, {0.25f, 0.25f, 0.50f, 0.50f}};
		return {{0.50f, 0.38f, 0.75f, 0.63f}, {0.50f, 0.63f, 0.75f, 0.88f},
				{0.00f, 0.38f, 0.25f, 0.62f}};
	}
	if (c.find("b3") != std::string::npos) {
		if (fall) return {{0.25f, 0.25f, 0.50f, 0.50f}};
		return {{0.00f, 0.25f, 0.25f, 0.50f}, {0.00f, 0.50f, 0.25f, 0.75f},
				{0.25f, 0.50f, 0.50f, 0.75f}};
	}
	if (c.find("n1") != std::string::npos) {
		if (winter) return {{0.00f, 0.36f, 0.50f, 0.58f}, {0.25f, 0.55f, 0.52f, 0.75f}};
		return {{0.00f, 0.72f, 0.28f, 0.96f}, {0.25f, 0.74f, 0.53f, 0.97f}};
	}
	if (c.find("n2") != std::string::npos) {
		return {{0.00f, 0.48f, 0.27f, 0.75f}, {0.26f, 0.73f, 0.58f, 1.00f},
				{0.75f, 0.48f, 1.00f, 0.80f}};
	}
	return {{0, 0, 1, 1}};
}

struct Buf {
	PackedVector3Array v, n;
	PackedVector2Array uv;
	PackedInt32Array idx;

	void quad(const Vector3 &a, const Vector3 &b, const Vector3 &c, const Vector3 &d,
			const UVRect &r, bool flip_u = false) {
		const int base = v.size();
		const Vector3 nn = (b - a).cross(d - a).normalized();
		v.push_back(a);
		v.push_back(b);
		v.push_back(c);
		v.push_back(d);
		for (int i = 0; i < 4; ++i)
			n.push_back(nn);
		const float l = flip_u ? r.u1 : r.u0;
		const float rr = flip_u ? r.u0 : r.u1;
		uv.push_back(Vector2(l, r.v1));
		uv.push_back(Vector2(rr, r.v1));
		uv.push_back(Vector2(rr, r.v0));
		uv.push_back(Vector2(l, r.v0));
		idx.push_back(base);
		idx.push_back(base + 1);
		idx.push_back(base + 2);
		idx.push_back(base);
		idx.push_back(base + 2);
		idx.push_back(base + 3);
	}

	void tube_quad(const Vector3 &b0, const Vector3 &b1, const Vector3 &t1,
			const Vector3 &t0, const Vector3 &n0, const Vector3 &n1,
			float u0, float u1, float v0, float v1) {
		const int base = v.size();
		v.push_back(b0);
		v.push_back(b1);
		v.push_back(t1);
		v.push_back(t0);
		n.push_back(n0);
		n.push_back(n1);
		n.push_back(n1);
		n.push_back(n0);
		uv.push_back(Vector2(u0, v0));
		uv.push_back(Vector2(u1, v0));
		uv.push_back(Vector2(u1, v1));
		uv.push_back(Vector2(u0, v1));
		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);
	}

	Array arrays() const {
		Array a;
		a.resize(Mesh::ARRAY_MAX);
		a[Mesh::ARRAY_VERTEX] = v;
		a[Mesh::ARRAY_NORMAL] = n;
		a[Mesh::ARRAY_TEX_UV] = uv;
		a[Mesh::ARRAY_INDEX] = idx;
		return a;
	}
};

void tube(Buf &m, const Vector3 &from, const Vector3 &to,
		float r0, float r1, int seg, float bark_repeat = 1.0f) {
	const Vector3 axis = (to - from).normalized();
	if (axis.length_squared() < 0.5f)
		return;
	const Vector3 helper = std::fabs(axis.y) > 0.9f ? Vector3(1, 0, 0) : Vector3(0, 1, 0);
	const Vector3 u = axis.cross(helper).normalized();
	const Vector3 w = axis.cross(u).normalized();
	for (int i = 0; i < seg; ++i) {
		const float a0 = float(i) / seg * kTAU;
		const float a1 = float(i + 1) / seg * kTAU;
		const Vector3 n0 = u * std::cos(a0) + w * std::sin(a0);
		const Vector3 n1 = u * std::cos(a1) + w * std::sin(a1);
		m.tube_quad(from + n0 * r0, from + n1 * r0, to + n1 * r1, to + n0 * r1,
				n0, n1, float(i) / seg, float(i + 1) / seg, bark_repeat, 0.0f);
	}
}

uint32_t hash32(uint32_t s) {
	s ^= s >> 16;
	s *= 0x7feb352dU;
	s ^= s >> 15;
	s *= 0x846ca68bU;
	s ^= s >> 16;
	return s;
}

float hash01(uint32_t s) {
	return float(hash32(s) & 0x00FFFFFFU) / float(0x01000000U);
}

uint32_t species_seed(const std::string &s) {
	uint32_t h = 2166136261U;
	for (unsigned char c : s) {
		h ^= c;
		h *= 16777619U;
	}
	return h;
}

bool species_is_palm(const std::string &hint) {
	const std::string h = lower(hint);
	static const char *kw[] = {"palm", "banana", "aloe", "fern", "joshua"};
	for (const char *k : kw)
		if (h.find(k) != std::string::npos)
			return true;
	return false;
}

void add_branches(Buf &wood, const std::string &species, float H, bool conifer, bool palm) {
	const float trunk_top = H * (palm ? 0.82f : (conifer ? 0.90f : 0.76f));
	const float trunk_r = H * (palm ? 0.028f : 0.035f);
	tube(wood, Vector3(0, 0, 0), Vector3(0, trunk_top, 0),
			trunk_r * 1.35f, trunk_r * 0.42f, 9, H * 0.22f);
	if (palm)
		return;
	const int count = conifer ? 9 : 8;
	const uint32_t seed = species_seed(species);
	for (int i = 0; i < count; ++i) {
		const float f = (i + 1.0f) / (count + 1.0f);
		const float y = H * (conifer ? (0.28f + f * 0.52f) : (0.32f + f * 0.34f));
		const float angle = kTAU * (f * 1.6180339f + hash01(seed + i * 17U));
		const float len = H * (conifer ? (0.24f * (1.0f - f * 0.55f)) :
				(0.18f + 0.08f * hash01(seed + i * 29U)));
		const Vector3 from(0, y, 0);
		const Vector3 to(std::cos(angle) * len,
				y + H * (conifer ? 0.06f : (0.10f + 0.06f * hash01(seed + i * 31U))),
				std::sin(angle) * len);
		tube(wood, from, to, trunk_r * (0.55f - 0.20f * f), trunk_r * 0.12f, 6,
				H * 0.08f);
		if (!conifer && (i % 2 == 0)) {
			const float side = angle + (hash01(seed + i * 37U) > 0.5f ? 0.65f : -0.65f);
			const Vector3 tip = to + Vector3(std::cos(side), 0.65f, std::sin(side)) * (len * 0.42f);
			tube(wood, to, tip, trunk_r * 0.16f, trunk_r * 0.05f, 5, H * 0.04f);
		}
	}
}

void add_leaf_cards(Buf &leaves, const std::string &species, float H,
		bool conifer, bool palm, const std::vector<UVRect> &rects) {
	const uint32_t seed = species_seed(species);
	const int count = palm ? 16 : (conifer ? 24 : 24);
	for (int i = 0; i < count; ++i) {
		const float a = kTAU * (float(i) * 0.6180339f + hash01(seed + i * 101U) * 0.15f);
		Vector3 center;
		float width = 1.0f, height = 1.0f;
		if (palm) {
			const float radial = H * (0.10f + 0.18f * hash01(seed + i * 103U));
			center = Vector3(std::cos(a) * radial, H * (0.78f + 0.12f * hash01(seed + i * 107U)),
					std::sin(a) * radial);
			width = H * 0.32f;
			height = H * 0.18f;
		} else if (conifer) {
			const float yf = 0.30f + 0.62f * (float(i) + 0.5f) / count;
			const float radial = H * 0.23f * (1.0f - yf * 0.70f) *
					(0.35f + 0.65f * hash01(seed + i * 109U));
			center = Vector3(std::cos(a) * radial, H * yf, std::sin(a) * radial);
			width = H * (0.18f + 0.10f * (1.0f - yf));
			height = H * 0.18f;
		} else {
			const float yf = hash01(seed + i * 109U);
			const float yn = yf * 2.0f - 1.0f;
			const float radial = H * 0.34f * std::sqrt(std::max(0.05f, 1.0f - yn * yn)) *
					(0.25f + 0.75f * std::sqrt(hash01(seed + i * 113U)));
			center = Vector3(std::cos(a) * radial, H * (0.58f + yf * 0.34f),
					std::sin(a) * radial);
			width = H * (0.23f + 0.10f * hash01(seed + i * 127U));
			height = H * (0.15f + 0.08f * hash01(seed + i * 131U));
		}
		const Vector3 right(std::cos(a + kPI * 0.5f), 0, std::sin(a + kPI * 0.5f));
		const Vector3 up(0, 1, 0);
		const UVRect &uv = rects[size_t(i) % rects.size()];
		auto card = [&](const Vector3 &r, bool flip) {
			leaves.quad(center - r * (width * 0.5f) - up * (height * 0.5f),
					center + r * (width * 0.5f) - up * (height * 0.5f),
					center + r * (width * 0.5f) + up * (height * 0.5f),
					center - r * (width * 0.5f) + up * (height * 0.5f), uv, flip);
		};
		card(right, (i & 1) != 0);
		// 原 SpeedTree leaf cluster 始终面向相机；静态 proxy 用交叉 card 保证任意视角
		// 都不会只看到一条边。离线 exporter 接入后由真实 leaf table 替代。
		const Vector3 crossed(std::cos(a), 0, std::sin(a));
		card(crossed, (i & 1) == 0);
	}
}

Ref<ArrayMesh> build_proxy_impl(const std::string &species, float height_m,
		const TreeTextures &textures) {
	const String hint(species.c_str());
	const bool conifer = species_is_conifer(hint);
	const bool palm = species_is_palm(species);
	const float H = std::max(2.0f, height_m);
	const bool atlas = textures.composite.is_valid();
	const std::vector<UVRect> rects = foliage_rects(species, textures.composite_name, atlas);

	Buf wood, leaves;
	add_branches(wood, species, H, conifer, palm);
	add_leaf_cards(leaves, species, H, conifer, palm, rects);

	Ref<ArrayMesh> mesh;
	mesh.instantiate();
	mesh->add_surface_from_arrays(Mesh::PRIMITIVE_TRIANGLES, wood.arrays());
	mesh->add_surface_from_arrays(Mesh::PRIMITIVE_TRIANGLES, leaves.arrays());

	Ref<StandardMaterial3D> bark;
	bark.instantiate();
	bark->set_albedo(textures.bark.is_valid() ? Color(1, 1, 1) : Color(0.30f, 0.21f, 0.13f));
	bark->set_roughness(1.0f);
	bark->set_texture_filter(StandardMaterial3D::TEXTURE_FILTER_LINEAR_WITH_MIPMAPS_ANISOTROPIC);
	if (textures.bark.is_valid())
		bark->set_texture(StandardMaterial3D::TEXTURE_ALBEDO, textures.bark);
	mesh->surface_set_material(0, bark);

	Ref<ShaderMaterial> leaf;
	leaf.instantiate();
	leaf->set_shader(leaf_shader());
	leaf->set_shader_parameter("leaf_tex",
			textures.composite.is_valid() ? textures.composite : fallback_leaf_texture(conifer));
	leaf->set_shader_parameter("wind_strength", 1.0f);
	mesh->surface_set_material(1, leaf);
	return mesh;
}

} // namespace

void cleanup_tree_shader() {
	g_tree_mesh_cache.clear();
	g_tree_texture_cache.clear();
	g_fallback_broadleaf.unref();
	g_fallback_conifer.unref();
	g_leaf_shader.unref();
}

bool species_is_conifer(const String &hint) {
	const String h = hint.to_lower();
	static const char *kw[] = {"cedar", "cypress", "pine", "fir", "spruce", "conifer", "juniper",
			"christmastree"};
	for (const char *k : kw)
		if (h.find(k) != -1)
			return true;
	return false;
}

Ref<ArrayMesh> build_placeholder_tree(const String &species_hint, float height_m) {
	TreeTextures empty;
	return build_proxy_impl(std::string(species_hint.utf8().get_data()), height_m, empty);
}

Ref<ArrayMesh> get_tree_proxy_mesh(const std::string &treefile,
		const fmt::AssetResolver &resolver, float height_m) {
	const std::string key = resolver.assets_root + "|" + fmt::AssetResolver::normalize(treefile) +
			"#" + std::to_string(height_m);
	auto found = g_tree_mesh_cache.find(key);
	if (found != g_tree_mesh_cache.end())
		return found->second;
	const TreeTextures textures = resolve_tree_textures(treefile, resolver);
	Ref<ArrayMesh> mesh = build_proxy_impl(treefile, height_m, textures);
	g_tree_mesh_cache.emplace(key, mesh);
	return mesh;
}

} // namespace mtgodot