# v5 bake — seamless limb sleeves + v3 torso plates:
#   limbs: the rectangular gen-v5 sleeve textures wrap the limb cylindrically
#          with a MIRRORED repeat (u ping-pongs around the circumference), so
#          the loop points are continuous by construction — no seam line, no
#          pattern mismatch. Every texel of the limb's UV island is painted
#          (full coverage), shorts texels skipped on legs.
#   torso: the v3 planar chest/back pieces (the look Can preferred), verbatim
#          from bake.gd.
# Output: baked_body.png. Run headless:
#   tinqs.console.exe --headless --path tattoo-test -s res://bake_v5.gd
extends SceneTree

const GLTF_PATH := "C:/Users/CAN/tinqs-ltd/ariki-game/assets/quaternius/source-models/Regular_Male_FullBody.gltf"
const SKIN_TEX_PATH := "C:/Users/CAN/tinqs-ltd/ariki-game/assets/quaternius/source-models/T_Regular_Male_Dark_BaseColor_png.png"
const MARKS_V3 := "C:/Users/CAN/tinqs-ltd/docs/conceptart/generated-images/tattoo-marks/"
const GEN_V5 := "C:/Users/CAN/tinqs-ltd/tattoo-test/gen-v5/"
const OUT_PATH := "C:/Users/CAN/tinqs-ltd/tattoo-test/baked_body.png"

const INK_COLOR := Color(0.10, 0.085, 0.08)
const INK_OPACITY := 0.88
const DILATE := 4

const LIMBS := [
	{"label": "arm_r", "tex": GEN_V5 + "v5_arm_r_chiefs-mark.png",
		"bones": ["upperarm_r", "lowerarm_r", "hand_r"], "center": "spine_03", "hem": false},
	{"label": "arm_l", "tex": GEN_V5 + "v5_arm_l_iron-spine.png",
		"bones": ["upperarm_l", "lowerarm_l", "hand_l"], "center": "spine_03", "hem": false},
	{"label": "leg_r", "tex": GEN_V5 + "v5_leg_r_voyagers-current.png",
		"bones": ["thigh_r", "calf_r", "foot_r"], "center": "pelvis", "hem": true},
	{"label": "leg_l", "tex": GEN_V5 + "v5_leg_l_storm-bearer.png",
		"bones": ["thigh_l", "calf_l", "foot_l"], "center": "pelvis", "hem": true},
]
# v3 torso pieces + framing constants, verbatim from bake.gd
const BACK_PATH := MARKS_V3 + "mark-06-hearth-warmth_back_v3.png"
const CHEST_PATH := MARKS_V3 + "mark-09-star-eye_chest_v3.png"
const BACK_HALF_WIDTH := 0.21
const BACK_Y_TOP_OFFSET := 0.04
const BACK_Y_CLIP := 1.04
const CHEST_HALF_WIDTH := 0.21
const CHEST_Y_TOP := 1.50
const CHEST_Y_CLIP := 1.15

var skel: Skeleton3D
var verts: PackedVector3Array
var normals: PackedVector3Array
var uvs: PackedVector2Array
var indices: PackedInt32Array
var armw_any: PackedFloat32Array
var skin_img: Image
var painted_mask: PackedByteArray
var tw: int
var th: int

func _init() -> void:
	var t_start := Time.get_ticks_msec()
	var doc := GLTFDocument.new()
	var state := GLTFState.new()
	if doc.append_from_file(GLTF_PATH, state) != OK:
		push_error("GLTF load failed")
		quit(1)
		return
	var scene := doc.generate_scene(state)

	var body: Node = null
	var stack: Array = [scene]
	while not stack.is_empty():
		var n: Node = stack.pop_back()
		if n is Skeleton3D and skel == null:
			skel = n
		if (n is MeshInstance3D or n.get_class() == "ImporterMeshInstance3D") \
				and String(n.name).to_lower().contains("regularmale"):
			body = n
		for c in n.get_children():
			stack.push_back(c)
	if skel == null or body == null:
		push_error("skeleton or body mesh not found")
		quit(1)
		return

	var mesh = body.mesh
	if mesh is ImporterMesh:
		mesh = mesh.get_mesh()
	var arrays: Array = mesh.surface_get_arrays(0)
	verts = arrays[Mesh.ARRAY_VERTEX]
	normals = arrays[Mesh.ARRAY_NORMAL]
	uvs = arrays[Mesh.ARRAY_TEX_UV]
	var vbones: PackedInt32Array = arrays[Mesh.ARRAY_BONES]
	var vweights = arrays[Mesh.ARRAY_WEIGHTS]
	indices = arrays[Mesh.ARRAY_INDEX]
	var nverts := verts.size()
	var influences := int(float(vbones.size()) / float(nverts))

	var skin: Skin = body.skin
	var bind_bone := PackedInt32Array()
	for i in skin.get_bind_count():
		var b := skin.get_bind_bone(i)
		if b < 0:
			b = skel.find_bone(skin.get_bind_name(i))
		bind_bone.append(b)

	skin_img = Image.load_from_file(SKIN_TEX_PATH)
	skin_img.convert(Image.FORMAT_RGBA8)
	tw = skin_img.get_width()
	th = skin_img.get_height()
	painted_mask = PackedByteArray()
	painted_mask.resize(tw * th)

	# per-vertex weights: per-limb (island pick) + any-arm (torso exclusion)
	var limb_w := {}
	for limb in LIMBS:
		limb_w[limb["label"]] = _weight_sum(bind_bone, vbones, vweights, influences, nverts,
			limb["bones"].slice(0, 2))
	armw_any = _weight_sum(bind_bone, vbones, vweights, influences, nverts,
		["upperarm_l", "lowerarm_l", "hand_l", "upperarm_r", "lowerarm_r", "hand_r"])

	var comp := _uv_components(nverts)
	for limb in LIMBS:
		var root := _dominant_component(comp, limb_w[limb["label"]], nverts)
		_bake_limb_island(limb, comp, root)

	var back_y_top := _bone_pos("neck_01").y + BACK_Y_TOP_OFFSET
	_bake_torso("back hearth-warmth", BACK_PATH, -1.0, BACK_HALF_WIDTH,
		back_y_top, BACK_Y_CLIP, "cover")
	_bake_torso("chest star-eye", CHEST_PATH, 1.0, CHEST_HALF_WIDTH,
		CHEST_Y_TOP, CHEST_Y_CLIP, "width")

	_dilate()
	skin_img.save_png(OUT_PATH)
	print("saved: ", OUT_PATH)
	print("bake took %d ms" % (Time.get_ticks_msec() - t_start))
	quit()

func _weight_sum(bind_bone: PackedInt32Array, vbones: PackedInt32Array, vweights,
		influences: int, nverts: int, names: Array) -> PackedFloat32Array:
	var bone_ids := {}
	for bn in names:
		var idx := skel.find_bone(bn)
		if idx < 0:
			push_error("bone not found: " + str(bn))
		bone_ids[idx] = true
	var binds := {}
	for i in bind_bone.size():
		if bone_ids.has(bind_bone[i]):
			binds[i] = true
	var w := PackedFloat32Array()
	w.resize(nverts)
	for v in nverts:
		var acc := 0.0
		for k in influences:
			if binds.has(vbones[v * influences + k]):
				acc += vweights[v * influences + k]
		w[v] = acc
	return w

func _uv_components(nverts: int) -> PackedInt32Array:
	var parent := PackedInt32Array()
	parent.resize(nverts)
	for i in nverts:
		parent[i] = i
	var ntris := indices.size() / 3
	for t in ntris:
		_union(parent, indices[t * 3], indices[t * 3 + 1])
		_union(parent, indices[t * 3 + 1], indices[t * 3 + 2])
	for i in nverts:
		parent[i] = _find(parent, i)
	return parent

func _find(parent: PackedInt32Array, i: int) -> int:
	while parent[i] != i:
		parent[i] = parent[parent[i]]
		i = parent[i]
	return i

func _union(parent: PackedInt32Array, a: int, b: int) -> void:
	var ra := _find(parent, a)
	var rb := _find(parent, b)
	if ra != rb:
		parent[ra] = rb

func _dominant_component(comp: PackedInt32Array, w: PackedFloat32Array, nverts: int) -> int:
	var sums := {}
	for v in nverts:
		sums[comp[v]] = sums.get(comp[v], 0.0) + w[v]
	var best := -1
	var best_w := 0.0
	for c in sums:
		if sums[c] > best_w:
			best_w = sums[c]
			best = c
	return best

func _bone_pos(bone: String) -> Vector3:
	return skel.get_bone_global_rest(skel.find_bone(bone)).origin

func _ink_bbox(img: Image) -> Rect2:
	var w := img.get_width()
	var h := img.get_height()
	var minx := w
	var maxx := -1
	var miny := h
	var maxy := -1
	for y in h:
		for x in w:
			var c := img.get_pixel(x, y)
			if (c.r + c.g + c.b) / 3.0 < 0.75:
				minx = mini(minx, x)
				maxx = maxi(maxx, x)
				miny = mini(miny, y)
				maxy = maxi(maxy, y)
	return Rect2(float(minx) / w, float(miny) / h,
		float(maxx - minx + 1) / w, float(maxy - miny + 1) / h)

func _sample_crop(img: Image, rect: Rect2, u: float, v: float) -> Color:
	return _sample_bilinear(img,
		rect.position.x + clampf(u, 0.0, 1.0) * rect.size.x,
		rect.position.y + clampf(v, 0.0, 1.0) * rect.size.y)

func _is_skin(px: int, py: int) -> bool:
	var c := skin_img.get_pixel(px, py)
	return (c.r + c.g + c.b) / 3.0 >= 0.30

# cylindrical coords of a 3D point against the limb's bone chain
func _limb_coords(pos: Vector3, segs: Array, total_len: float, outward: Vector3) -> Vector2:
	var best_dist := 1e9
	var t_along := 0.0
	var ang := 0.0
	for seg in segs:
		var a: Vector3 = seg[0]
		var axis: Vector3 = (seg[1] as Vector3) - a
		var seg_len := axis.length()
		axis /= seg_len
		var s := clampf((pos - a).dot(axis), 0.0, seg_len)
		var cp: Vector3 = a + axis * s
		var d := pos.distance_to(cp)
		if d < best_dist:
			best_dist = d
			t_along = ((seg[2] as float) + s) / total_len
			var uref: Vector3 = (outward - axis * outward.dot(axis)).normalized()
			var dv: Vector3 = pos - cp
			ang = atan2(dv.dot(axis.cross(uref)), dv.dot(uref))
	return Vector2(ang, t_along)

# ── limb sleeve: full island painted, mirrored wrap (seamless loop) ─────────
func _bake_limb_island(limb: Dictionary, comp: PackedInt32Array, root: int) -> void:
	var tat := Image.load_from_file(limb["tex"])
	tat.convert(Image.FORMAT_RGBA8)
	var crop := _ink_bbox(tat)
	var bones: Array = limb["bones"]
	var p0 := _bone_pos(bones[0])
	var p1 := _bone_pos(bones[1])
	var p2 := _bone_pos(bones[2])
	var center := _bone_pos(limb["center"])
	var len1 := p0.distance_to(p1)
	var total_len := len1 + p1.distance_to(p2)
	var outward := (p0 - center).normalized()
	var segs := [[p0, p1, 0.0], [p1, p2, len1]]
	var skip_shorts: bool = limb["hem"]

	# v range over the island's SKIN vertices (legs: hem→ankle, arms: cap→wrist)
	var t_min := 1e9
	var t_max := -1e9
	for v in verts.size():
		if comp[v] != root:
			continue
		if skip_shorts:
			var sx := clampi(int(uvs[v].x * tw), 0, tw - 1)
			var sy := clampi(int(uvs[v].y * th), 0, th - 1)
			if not _is_skin(sx, sy):
				continue
		var t := _limb_coords(verts[v], segs, total_len, outward).y
		t_min = minf(t_min, t)
		t_max = maxf(t_max, t)
	print("%s v-range %.3f..%.3f" % [limb["label"], t_min, t_max])

	var painted := 0
	var ntris := indices.size() / 3
	for t in ntris:
		var i0 := indices[t * 3]
		if comp[i0] != root:
			continue
		var i1 := indices[t * 3 + 1]
		var i2 := indices[t * 3 + 2]
		painted += _raster_tri(i0, i1, i2, skip_shorts,
			func(pos: Vector3, _nrm: Vector3, _bary: Vector3) -> float:
				var cc := _limb_coords(pos, segs, total_len, outward)
				var v_norm := clampf((cc.y - t_min) / (t_max - t_min), 0.0, 1.0)
				var frac := fposmod(0.5 + cc.x / TAU, 1.0)
				var u := 1.0 - absf(2.0 * frac - 1.0)  # mirrored repeat: seamless loop
				var c := _sample_crop(tat, crop, u, v_norm)
				var ink := 1.0 - smoothstep(0.45, 0.80, (c.r + c.g + c.b) / 3.0)
				return ink * INK_OPACITY)
	print(limb["label"], " painted px: ", painted)

# ── torso piece, v3 planar projection (verbatim from bake.gd) ────────────────
func _bake_torso(label: String, tex_path: String, facing: float, hw: float,
		y_top: float, y_clip: float, fit: String) -> void:
	var tat := Image.load_from_file(tex_path)
	tat.convert(Image.FORMAT_RGBA8)
	var crop := _ink_bbox(tat)
	var aspect := crop.size.y / crop.size.x  # normalized rect → recompute px aspect
	var timg_aspect := float(tat.get_height()) / float(tat.get_width())
	aspect = aspect * timg_aspect / 1.0
	var box_w := 2.0 * hw
	var design_w := box_w
	if fit == "cover":
		design_w = maxf(box_w, (y_top - y_clip) / aspect)
	var design_h := design_w * aspect
	var y_bot := maxf(y_clip, y_top - design_h)
	print("%s: design %.2fw x %.2fh, y %.2f..%.2f" % [label, design_w, design_h, y_bot, y_top])
	var painted := 0
	var ntris := indices.size() / 3
	for t in ntris:
		var i0 := indices[t * 3]
		var i1 := indices[t * 3 + 1]
		var i2 := indices[t * 3 + 2]
		if min(armw_any[i0], min(armw_any[i1], armw_any[i2])) > 0.5:
			continue
		var ymax := maxf(verts[i0].y, maxf(verts[i1].y, verts[i2].y))
		var ymin := minf(verts[i0].y, minf(verts[i1].y, verts[i2].y))
		if ymin > y_top + 0.05 or ymax < y_bot - 0.05:
			continue
		painted += _raster_tri(i0, i1, i2, true,
			func(pos: Vector3, nrm: Vector3, bary: Vector3) -> float:
				if pos.y < y_bot or pos.y > y_top or absf(pos.x) > hw:
					return -1.0
				var faceness := nrm.z * facing
				if faceness < 0.25:
					return -1.0
				var aw: float = bary.x * armw_any[i0] + bary.y * armw_any[i1] + bary.z * armw_any[i2]
				if aw > 0.4:
					return -1.0
				var x_norm := (pos.x * facing + hw) / box_w
				var u := 0.5 + (x_norm - 0.5) * (box_w / design_w)
				var v := (y_top - pos.y) / design_h
				if u < 0.0 or u > 1.0 or v > 1.0:
					return -1.0
				var c := _sample_crop(tat, crop, u, v)
				var ink := 1.0 - smoothstep(0.45, 0.80, (c.r + c.g + c.b) / 3.0)
				var alpha := ink * INK_OPACITY * smoothstep(0.25, 0.45, faceness)
				alpha *= 1.0 - smoothstep(0.2, 0.4, aw)
				return alpha)
	print(label, " painted px: ", painted)

# rasterize one triangle in UV space; shade(pos, nrm, bary) → alpha (<=0 skips
# ink but still marks coverage). skip_shorts guards dark shorts texels.
func _raster_tri(i0: int, i1: int, i2: int, skip_shorts: bool, shade: Callable) -> int:
	var p0 := Vector2(uvs[i0].x * tw, uvs[i0].y * th)
	var p1 := Vector2(uvs[i1].x * tw, uvs[i1].y * th)
	var p2 := Vector2(uvs[i2].x * tw, uvs[i2].y * th)
	var denom := (p1.y - p2.y) * (p0.x - p2.x) + (p2.x - p1.x) * (p0.y - p2.y)
	if absf(denom) < 1e-9:
		return 0
	var minx := clampi(int(floor(min(p0.x, min(p1.x, p2.x)))), 0, tw - 1)
	var maxx := clampi(int(ceil(max(p0.x, max(p1.x, p2.x)))), 0, tw - 1)
	var miny := clampi(int(floor(min(p0.y, min(p1.y, p2.y)))), 0, th - 1)
	var maxy := clampi(int(ceil(max(p0.y, max(p1.y, p2.y)))), 0, th - 1)
	var painted := 0
	for py in range(miny, maxy + 1):
		for px in range(minx, maxx + 1):
			var fx := px + 0.5
			var fy := py + 0.5
			var b0 := ((p1.y - p2.y) * (fx - p2.x) + (p2.x - p1.x) * (fy - p2.y)) / denom
			var b1 := ((p2.y - p0.y) * (fx - p2.x) + (p0.x - p2.x) * (fy - p2.y)) / denom
			var b2 := 1.0 - b0 - b1
			if b0 < -0.001 or b1 < -0.001 or b2 < -0.001:
				continue
			if skip_shorts and not _is_skin(px, py):
				continue
			var pos: Vector3 = verts[i0] * b0 + verts[i1] * b1 + verts[i2] * b2
			var nrm: Vector3 = (normals[i0] * b0 + normals[i1] * b1 + normals[i2] * b2).normalized()
			var alpha: float = shade.call(pos, nrm, Vector3(b0, b1, b2))
			painted_mask[py * tw + px] = 1
			if alpha <= 0.003:
				continue
			skin_img.set_pixel(px, py, skin_img.get_pixel(px, py).lerp(INK_COLOR, alpha))
			painted += 1
	return painted

# pull composited colors DILATE px into unpainted texels (seam gutters)
func _dilate() -> void:
	for pass_i in DILATE:
		var grown := PackedInt32Array()
		for y in th:
			for x in tw:
				if painted_mask[y * tw + x] == 1:
					continue
				var acc := Color(0, 0, 0)
				var n := 0
				for d in [[1, 0], [-1, 0], [0, 1], [0, -1]]:
					var nx: int = x + d[0]
					var ny: int = y + d[1]
					if nx < 0 or nx >= tw or ny < 0 or ny >= th:
						continue
					if painted_mask[ny * tw + nx] == 1:
						acc += skin_img.get_pixel(nx, ny)
						n += 1
				if n > 0:
					skin_img.set_pixel(x, y, acc / n)
					grown.append(y * tw + x)
		for p in grown:
			painted_mask[p] = 1

func _sample_bilinear(img: Image, x: float, y: float) -> Color:
	var w := img.get_width()
	var h := img.get_height()
	var fx := clampf(x * w - 0.5, 0.0, w - 1.001)
	var fy := clampf(y * h - 0.5, 0.0, h - 1.001)
	var x0 := int(fx)
	var y0 := int(fy)
	var x1 := mini(x0 + 1, w - 1)
	var y1 := mini(y0 + 1, h - 1)
	var tx := fx - x0
	var ty := fy - y0
	return img.get_pixel(x0, y0).lerp(img.get_pixel(x1, y0), tx) \
		.lerp(img.get_pixel(x0, y1).lerp(img.get_pixel(x1, y1), tx), ty)