# Cuts per-body-part UV island masks for the Quaternius Regular Male:
#   mask_arm_r / mask_arm_l   — full sleeve islands (shoulder→wrist)
#   mask_leg_r / mask_leg_l   — leg islands, visible skin only (shorts excluded)
#   mask_chest / mask_back    — torso island split by rest-pose normal Z sign
# White = tattooable skin, black = everything else. 2048², same UV space as
# T_Regular_Male_Dark_BaseColor. Also prints each mask's ink bbox (px) for the
# gen/composite scripts.
# Run headless:
#   tinqs.console.exe --headless --path tattoo-test -s res://masks.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 OUT_DIR := "C:/Users/CAN/tinqs-ltd/tattoo-test/masks/"
const SIZE := 2048

const LIMB_MASKS := [
	{"name": "arm_r", "bones": ["upperarm_r", "lowerarm_r"], "skip_shorts": false},
	{"name": "arm_l", "bones": ["upperarm_l", "lowerarm_l"], "skip_shorts": false},
	{"name": "leg_r", "bones": ["thigh_r", "calf_r"], "skip_shorts": true},
	{"name": "leg_l", "bones": ["thigh_l", "calf_l"], "skip_shorts": true},
]
const TORSO_BONES := ["spine_01", "spine_02", "spine_03", "clavicle_l", "clavicle_r", "pelvis"]
const ARM_BONES := ["upperarm_l", "lowerarm_l", "hand_l", "upperarm_r", "lowerarm_r", "hand_r"]
const NECK_BONES := ["neck_01", "Head"]

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

func _init() -> void:
	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()

	DirAccess.make_dir_recursive_absolute(OUT_DIR)

	# per-vertex weight sums per group
	var groups := {}
	for lm in LIMB_MASKS:
		groups[lm["name"]] = _weight_sum(bind_bone, vbones, vweights, influences, nverts, lm["bones"])
	groups["torso"] = _weight_sum(bind_bone, vbones, vweights, influences, nverts, TORSO_BONES)
	groups["arm_any"] = _weight_sum(bind_bone, vbones, vweights, influences, nverts, ARM_BONES)
	groups["neck"] = _weight_sum(bind_bone, vbones, vweights, influences, nverts, NECK_BONES)

	# limb masks are ISLAND-based, not weight-based: the whole connected UV island
	# whose vertices carry the most limb weight. The arm island IS shoulder→wrist
	# skin (hand is a separate island), so this leaves no bare rim at the
	# weight-blend zones — full-sleeve coverage by construction.
	var comp := _uv_components(nverts)
	for lm in LIMB_MASKS:
		var w: PackedFloat32Array = groups[lm["name"]]
		var best := _dominant_component(comp, w, nverts)
		_bake_island_mask(lm["name"], lm["skip_shorts"], comp, best)

	# torso split: chest = rest normal z > 0.15, back = z < -0.15
	var tor: PackedFloat32Array = groups["torso"]
	var armw: PackedFloat32Array = groups["arm_any"]
	var neckw: PackedFloat32Array = groups["neck"]
	_bake_torso_mask("chest", 1.0, tor, armw, neckw)
	_bake_torso_mask("back", -1.0, tor, armw, neckw)
	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

# union-find over shared vertex indices: connected components = UV islands
# (vertices are split at UV seams in the GLTF, so indices only connect within an island)
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

# component root whose vertices carry the highest total group weight
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

# rasterize every triangle of one UV island (optionally minus shorts texels)
func _bake_island_mask(mname: String, skip_shorts: bool, comp: PackedInt32Array, root: int) -> void:
	var img := Image.create(SIZE, SIZE, false, Image.FORMAT_L8)
	img.fill(Color(0, 0, 0))
	var ntris := indices.size() / 3
	for t in ntris:
		var i0 := indices[t * 3]
		if comp[i0] != root:
			continue
		_raster_mask(img, i0, indices[t * 3 + 1], indices[t * 3 + 2], skip_shorts,
			func(_b0: float, _b1: float, _b2: float, _pos: Vector3, _nrm: Vector3) -> bool:
				return true)
	_save_mask(img, mname)

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

func _bake_mask(mname: String, skip_shorts: bool, accept: Callable, w: PackedFloat32Array) -> void:
	var img := Image.create(SIZE, SIZE, false, Image.FORMAT_L8)
	img.fill(Color(0, 0, 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 max(w[i0], max(w[i1], w[i2])) < 0.3:
			continue
		_raster_mask(img, i0, i1, i2, skip_shorts,
			func(b0: float, b1: float, b2: float, pos: Vector3, nrm: Vector3) -> bool:
				return accept.call(b0 * w[i0] + b1 * w[i1] + b2 * w[i2], pos, nrm))
	_save_mask(img, mname)

func _bake_torso_mask(mname: String, facing: float, tor: PackedFloat32Array,
		armw: PackedFloat32Array, neckw: PackedFloat32Array) -> void:
	var img := Image.create(SIZE, SIZE, false, Image.FORMAT_L8)
	img.fill(Color(0, 0, 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 max(tor[i0], max(tor[i1], tor[i2])) < 0.3:
			continue
		_raster_mask(img, i0, i1, i2, true,
			func(b0: float, b1: float, b2: float, _pos: Vector3, nrm: Vector3) -> bool:
				var tw_ := b0 * tor[i0] + b1 * tor[i1] + b2 * tor[i2]
				var aw := b0 * armw[i0] + b1 * armw[i1] + b2 * armw[i2]
				var nw := b0 * neckw[i0] + b1 * neckw[i1] + b2 * neckw[i2]
				return tw_ > 0.5 and aw < 0.35 and nw < 0.35 and nrm.z * facing > 0.10)
	_save_mask(img, mname)

func _raster_mask(img: Image, i0: int, i1: int, i2: int, skip_shorts: bool, accept: Callable) -> void:
	var p0 := Vector2(uvs[i0].x * SIZE, uvs[i0].y * SIZE)
	var p1 := Vector2(uvs[i1].x * SIZE, uvs[i1].y * SIZE)
	var p2 := Vector2(uvs[i2].x * SIZE, uvs[i2].y * SIZE)
	var denom := (p1.y - p2.y) * (p0.x - p2.x) + (p2.x - p1.x) * (p0.y - p2.y)
	if absf(denom) < 1e-9:
		return
	var minx := clampi(int(floor(min(p0.x, min(p1.x, p2.x)))), 0, SIZE - 1)
	var maxx := clampi(int(ceil(max(p0.x, max(p1.x, p2.x)))), 0, SIZE - 1)
	var miny := clampi(int(floor(min(p0.y, min(p1.y, p2.y)))), 0, SIZE - 1)
	var maxy := clampi(int(ceil(max(p0.y, max(p1.y, p2.y)))), 0, SIZE - 1)
	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:
				var sx := clampi(int(float(px) / SIZE * tw), 0, tw - 1)
				var sy := clampi(int(float(py) / SIZE * th), 0, th - 1)
				if _is_shorts(sx, sy):
					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()
			if accept.call(b0, b1, b2, pos, nrm):
				img.set_pixel(px, py, Color(1, 1, 1))

# keeps only the largest 4-connected white component (kills stray slivers from
# shoulder-cap / wrist texels that bleed into neighbouring islands)
func _largest_component(img: Image) -> void:
	var lbl := PackedInt32Array()
	lbl.resize(SIZE * SIZE)
	var sizes := [0]  # component 0 = background
	var next_lbl := 1
	var queue := PackedInt32Array()
	for start in SIZE * SIZE:
		if lbl[start] != 0 or img.get_pixel(start % SIZE, start / SIZE).r < 0.5:
			continue
		var comp := next_lbl
		next_lbl += 1
		var count := 0
		queue.clear()
		queue.append(start)
		lbl[start] = comp
		var head := 0
		while head < queue.size():
			var p := queue[head]
			head += 1
			count += 1
			var x := p % SIZE
			var y := p / SIZE
			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 >= SIZE or ny < 0 or ny >= SIZE:
					continue
				var np := ny * SIZE + nx
				if lbl[np] == 0 and img.get_pixel(nx, ny).r > 0.5:
					lbl[np] = comp
					queue.append(np)
		sizes.append(count)
	var best := 1
	for c in range(2, sizes.size()):
		if sizes[c] > sizes[best]:
			best = c
	for p in SIZE * SIZE:
		if lbl[p] != 0 and lbl[p] != best:
			img.set_pixel(p % SIZE, p / SIZE, Color(0, 0, 0))

func _save_mask(img: Image, mname: String) -> void:
	_largest_component(img)
	# bbox of white pixels
	var minx := SIZE
	var maxx := -1
	var miny := SIZE
	var maxy := -1
	var count := 0
	for y in SIZE:
		for x in SIZE:
			if img.get_pixel(x, y).r > 0.5:
				minx = mini(minx, x)
				maxx = maxi(maxx, x)
				miny = mini(miny, y)
				maxy = maxi(maxy, y)
				count += 1
	img.convert(Image.FORMAT_RGBA8)
	var path := OUT_DIR + "mask_" + mname + ".png"
	img.save_png(path)
	print("MASK %s px=%d bbox=%d,%d,%d,%d saved=%s" % [mname, count, minx, miny, maxx - minx + 1, maxy - miny + 1, path])