# Bakes v3 tattoo marks onto the QuatSkin Lena body texture:
#   right arm  chiefs-mark v3        (360° wrap shoulder→wrist)
#   right leg  voyagers-current v3   (360° wrap below-underwear→ankle)
#   chest      star-eye v3           (planar +Z, width fit, clipped above navel)
#   back       hearth-warmth v3      (planar -Z, cover fit, tip at neck)
# Lena's arms AND legs are mirror-UV'd (~100% shared islands), so each limb
# design prints on BOTH sides — bake one side only, never two designs per pair.
# Painting happens in 3D bone space; her fragmented Tripo UV atlas receives the
# ink wherever her UVs put it. Fabric texels (painted bra/underwear: bright +
# low red-blue chroma) are skipped so ink stays on skin.
# Run headless:
#   tinqs.console.exe --headless --path tattoo-test -s res://bake_lena.gd
extends SceneTree

const GLB_PATH := "C:/Users/CAN/tinqs-ltd/tattoo-test/lena_clean_uv.glb"
const MARKS := "C:/Users/CAN/tinqs-ltd/design/conceptart/generated-images/tattoo-marks/"
const OUT_PATH := "C:/Users/CAN/tinqs-ltd/tattoo-test/baked_lena_body.png"

const INK_COLOR := Color(0.10, 0.085, 0.08)
const INK_OPACITY := 0.88
# Lena's baked Tripo texture has real shading — only skip near-black texels
const DARK_SKIP := 0.12

# limb sleeves; t0 = fraction of the bone chain where the design starts
# (legs start below the painted underwear hem)
const LIMBS := [
	{"label": "arm_r chiefs-mark", "tex": MARKS + "mark-01-chiefs-mark_arm_v3.png",
		"bones": ["upperarm_r", "lowerarm_r", "hand_r"], "center": "spine_03", "t0": 0.0},
	{"label": "arm_l iron-spine", "tex": MARKS + "mark-10-iron-spine_arm_v3.png",
		"bones": ["upperarm_l", "lowerarm_l", "hand_l"], "center": "spine_03", "t0": 0.0},
	{"label": "leg_r voyagers-current", "tex": MARKS + "mark-02-voyagers-current_leg_v3.png",
		"bones": ["thigh_r", "calf_r", "foot_r"], "center": "pelvis", "t0": 0.15},
	{"label": "leg_l storm-bearer", "tex": MARKS + "mark-07-storm-bearer_leg_v3.png",
		"bones": ["thigh_l", "calf_l", "foot_l"], "center": "pelvis", "t0": 0.15},
]

# torso boxes (heights from her skeleton: neck_01 1.413, pelvis 0.932)
const CHEST_PATH := MARKS + "mark-09-star-eye_chest_v3.png"
const CHEST_HALF_WIDTH := 0.17
const CHEST_Y_TOP := 1.38          # just under her clavicle line (1.397)
const CHEST_Y_CLIP := 1.05         # hard floor — navel stays empty
const BACK_PATH := MARKS + "mark-06-hearth-warmth_back_v3.png"
const BACK_HALF_WIDTH := 0.18
const BACK_Y_TOP_OFFSET := 0.03    # above neck_01 — flame tip reaches the neck
const BACK_Y_CLIP := 0.96          # waistline

var skel: Skeleton3D
var verts: PackedVector3Array
var normals: PackedVector3Array
var uvs: PackedVector2Array
var indices: PackedInt32Array
var limb_w := {}                 # limb label -> per-vertex weight array
var armw_any: PackedFloat32Array # both arms incl. fingers (torso-pass exclusion)
var skin_img: Image
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(GLB_PATH, state) != OK:
		push_error("GLB 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("lena"):
			body = n
		for c in n.get_children():
			stack.push_back(c)
	if skel == null or body == null:
		push_error("skeleton or Lena 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)

	var limb_binds := {}
	for limb in LIMBS:
		limb_binds[limb["label"]] = _bind_set(bind_bone, limb["bones"])
	# arm exclusion for torso passes — Lena has full finger bones, include them
	var arm_bones := []
	for side in ["l", "r"]:
		arm_bones += ["upperarm_%s" % side, "lowerarm_%s" % side, "hand_%s" % side]
		for f in ["index", "middle", "pinky", "ring", "thumb"]:
			for seg in ["01", "02", "03", "04_leaf"]:
				arm_bones.append("%s_%s_%s" % [f, seg, side])
	var binds_arm_any := _bind_set(bind_bone, arm_bones)

	for limb in LIMBS:
		var w := PackedFloat32Array()
		w.resize(nverts)
		limb_w[limb["label"]] = w
	armw_any = PackedFloat32Array()
	armw_any.resize(nverts)
	for v in nverts:
		for k in influences:
			var bi := vbones[v * influences + k]
			var vw: float = vweights[v * influences + k]
			if binds_arm_any.has(bi):
				armw_any[v] += vw
			for limb in LIMBS:
				if limb_binds[limb["label"]].has(bi):
					limb_w[limb["label"]][v] += vw

	# source texture straight from the GLB material — UVs are guaranteed to match
	var mat: Material = mesh.surface_get_material(0)
	skin_img = (mat as BaseMaterial3D).albedo_texture.get_image()
	skin_img.clear_mipmaps()
	skin_img.convert(Image.FORMAT_RGBA8)
	tw = skin_img.get_width()
	th = skin_img.get_height()
	print("texture %dx%d" % [tw, th])

	for limb in LIMBS:
		_bake_limb(limb)
	_bake_torso("chest star-eye", CHEST_PATH, 1.0, CHEST_HALF_WIDTH,
		CHEST_Y_TOP, CHEST_Y_CLIP, "width")
	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")

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

func _bind_set(bind_bone: PackedInt32Array, names: Array) -> Dictionary:
	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 out := {}
	for i in bind_bone.size():
		if bone_ids.has(bind_bone[i]):
			out[i] = true
	return out

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

# ink bounding box of a mark: normalized rect + pixel aspect (h/w) of the design
func _ink_bbox(img: Image) -> Dictionary:
	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)
	var rect := Rect2(float(minx) / w, float(miny) / h,
		float(maxx - minx + 1) / w, float(maxy - miny + 1) / h)
	return {"rect": rect, "aspect": float(maxy - miny + 1) / float(maxx - minx + 1)}

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)

# ── limb sleeve: cropped design wraps the full circumference ────────────────
func _bake_limb(limb: Dictionary) -> void:
	var tat := Image.load_from_file(limb["tex"])
	tat.convert(Image.FORMAT_RGBA8)
	var bbox := _ink_bbox(tat)
	var crop: Rect2 = bbox["rect"]
	var w: PackedFloat32Array = limb_w[limb["label"]]
	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 len2 := p1.distance_to(p2)
	var total_len := len1 + len2
	var outward := (p0 - center).normalized()
	# one global angular frame for the whole limb — per-segment frames jump at the
	# joint when the rest pose bends there (Lena's elbow ~17°), printing a split line
	var limb_axis := (p2 - p0).normalized()
	var uref_g := (outward - limb_axis * outward.dot(limb_axis)).normalized()
	var vref_g := limb_axis.cross(uref_g)
	var ax1 := (p1 - p0) / len1
	var ax2 := (p2 - p1) / len2
	# blend width for the seg1/seg2 handover: a hard closest-segment pick is
	# discontinuous on the bend's concave (inner-joint) side
	const JUNCTION_BLEND := 0.03
	var t0: float = limb["t0"]
	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 max(w[i0], max(w[i1], w[i2])) < 0.3:
			continue
		painted += _raster_tri(i0, i1, i2, func(pos: Vector3, _nrm: Vector3, bary: Vector3) -> float:
			var aw: float = bary.x * w[i0] + bary.y * w[i1] + bary.z * w[i2]
			if aw < 0.30:
				return -1.0
			var s1 := clampf((pos - p0).dot(ax1), 0.0, len1)
			var cp1 := p0 + ax1 * s1
			var d1 := pos.distance_to(cp1)
			var s2 := clampf((pos - p1).dot(ax2), 0.0, len2)
			var cp2 := p1 + ax2 * s2
			var d2 := pos.distance_to(cp2)
			var wb := smoothstep(-JUNCTION_BLEND, JUNCTION_BLEND, d1 - d2)
			var t_along := lerpf(s1 / total_len, (len1 + s2) / total_len, wb)
			var dv: Vector3 = pos - cp1.lerp(cp2, wb)
			var ang := atan2(dv.dot(vref_g), dv.dot(uref_g))
			var v_norm := (t_along - t0) / (1.0 - t0)
			if v_norm <= 0.001 or v_norm >= 0.999:
				return -1.0
			var c := _sample_crop(tat, crop, fposmod(0.5 + ang / TAU, 1.0), v_norm)
			var ink := 1.0 - smoothstep(0.45, 0.80, (c.r + c.g + c.b) / 3.0)
			var alpha := ink * INK_OPACITY * smoothstep(0.30, 0.50, aw)
			alpha *= smoothstep(0.0, 0.03, v_norm) * (1.0 - smoothstep(0.97, 1.0, v_norm))
			return alpha)
	print(limb["label"], " painted px: ", painted)

# ── torso piece, planar projection along Z (facing: -1 back, +1 chest) ──────
# fit "width": design spans the full box width, height follows its aspect.
# fit "cover": design fills the whole box (top-anchored), overflow is clipped.
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 bbox := _ink_bbox(tat)
	var crop: Rect2 = bbox["rect"]
	var aspect: float = bbox["aspect"]
	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, x ±%.2f" %
		[label, design_w, design_h, y_bot, y_top, hw])
	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]
		# prefilter: skip triangles clearly outside the box / on arms
		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, 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
			# u flips with facing so the design reads unmirrored from the viewer's side
			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)

# rasterizes one triangle in UV space; shade(pos, normal, bary) → alpha (<=0 skips).
# skips near-black texels and fabric texels (painted bra/underwear: bright with
# low red-blue chroma; skin is warm — (r-b)/r ≈ 0.5) so ink stays on skin.
func _raster_tri(i0: int, i1: int, i2: int, 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
			var base_col := skin_img.get_pixel(px, py)
			var lum := (base_col.r + base_col.g + base_col.b) / 3.0
			if lum < DARK_SKIP:
				continue
			if lum > 0.60 and (base_col.r - base_col.b) / maxf(base_col.r, 0.001) < 0.27:
				continue  # fabric
			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))
			if alpha <= 0.003:
				continue
			skin_img.set_pixel(px, py, base_col.lerp(INK_COLOR, alpha))
			painted += 1
	return painted

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