//
//  SkeletonRenderer.swift
//  HelloVisionOS - Niall George Horn - June 2026
//
//  Builds and updates a RealityKit entity hierarchy representing the
//  CoreIK skeleton output. Joint positions are rendered as small spheres;
//  bones as thin box segments connecting parent→child.
//

import RealityKit
import simd
import UIKit

// SkeletonRenderer

class SkeletonRenderer {

    let rootEntity = Entity()

    private var jointEntities: [ModelEntity] = []
    private var boneEntities:  [ModelEntity] = []

    // Input joints: only wrist joints highlighted (head is internal, not a user-driven EE)
    private let inputIndices: Set<Int> = [38, 66]

    // Build

    func buildSkeleton(jointCount: Int) {
        jointEntities.removeAll()
        boneEntities.removeAll()
        for child in rootEntity.children { child.removeFromParent() }

        // Joint spheres
        for i in 0..<jointCount {
            let isInput = inputIndices.contains(i)
            let radius: Float = isInput ? 0.04 : 0.02  // 4cm EE, 2cm regular

            let mesh = MeshResource.generateSphere(radius: radius)
            let tint: UIColor = isInput
                ? UIColor(red: 1,   green: 0.3, blue: 0,   alpha: 1)  // orange — anchor joints
                : UIColor(red: 0.1, green: 1,   blue: 0.1, alpha: 1)  // bright green

            var mat = PhysicallyBasedMaterial()
            mat.baseColor = .init(tint: tint)
            mat.roughness = .init(floatLiteral: 0.2)
            mat.emissiveColor = .init(color: tint)
            mat.emissiveIntensity = isInput ? 3.0 : 1.5  // self-lit — visible in bright room
            let entity = ModelEntity(mesh: mesh, materials: [mat])
            entity.name = "joint_\(i)"
            rootEntity.addChild(entity)
            jointEntities.append(entity)
        }

        // Bone segments (thin box per bone, excluding root)
        for i in 1..<jointCount {
            let mesh = MeshResource.generateBox(size: SIMD3(0.018, 0.05, 0.018))
            var mat = PhysicallyBasedMaterial()
            mat.baseColor = .init(tint: UIColor(white: 1, alpha: 1))
            mat.emissiveColor = .init(color: UIColor(white: 0.9, alpha: 1))
            mat.emissiveIntensity = 1.0
            mat.roughness = .init(floatLiteral: 0.3)
            let entity = ModelEntity(mesh: mesh, materials: [mat])
            entity.name = "bone_\(i)"
            rootEntity.addChild(entity)
            boneEntities.append(entity)
        }
    }

    // Joint visibility

    private var permanentlyHidden: Set<Int> = []

    /// Permanently hide a set of joint indices and their parent bones.
    /// Survives repeated updateSkeleton calls.
    func hideJoints(_ indices: Set<Int>) {
        permanentlyHidden = indices
        for i in indices {
            if i < jointEntities.count { jointEntities[i].isEnabled = false }
            if i >= 1 && (i - 1) < boneEntities.count { boneEntities[i - 1].isEnabled = false }
        }
    }

    // Update

    /// World position of the hip joint (j1) used as scene anchor.
    private(set) var hipOffset: SIMD3<Float> = .zero

    func updateSkeleton(joints: [JointTransform]) {
        guard joints.count == jointEntities.count else {
            print("[SkeletonRenderer] count mismatch: joints=\(joints.count) entities=\(jointEntities.count)")
            return
        }

        // Normalise by joint 1 (hips) — the stable body centre.
        // Joint 0 is a virtual VIO anchor that moves with the head; normalising
        // by j0 causes apparent body translation as VIO changes because j0 shifts
        // independently of the rest of the skeleton chain.
        // Normalising by hips keeps the body centred regardless of head VIO position.
        guard joints.count > 1 else { return }
        let norm = joints[1].worldPosition   // hips world pos in solver frame
        hipOffset = norm

        // Hide the virtual root sphere (j0); it has no anatomical meaning
        // and overlaps the head marker in scene space.
        jointEntities[0].isEnabled = false

        for i in 1..<joints.count {
            if permanentlyHidden.contains(i) {
                jointEntities[i].isEnabled = false
                continue
            }
            jointEntities[i].position = joints[i].worldPosition - norm
        }

        // Position + orient bone segments
        for i in 1..<joints.count {
            // Skip permanently hidden joints and their incoming bone
            if permanentlyHidden.contains(i) {
                boneEntities[i - 1].isEnabled = false
                continue
            }

            let p = kPersonaParent[i]
            guard p >= 0 && p < joints.count else {
                boneEntities[i - 1].isEnabled = false
                continue
            }

            // Skip the virtual-root → hips bone (i==1, p==0); it has no meaning
            // and creates a confusing stem from the hip origin up to (0,0,0).
            if i == 1 {
                boneEntities[i - 1].isEnabled = false
                continue
            }

            let pPos = joints[p].worldPosition - norm
            let cPos = joints[i].worldPosition - norm

            let diff = cPos - pPos
            let length = simd_length(diff)

            guard length > 0.001 else {
                boneEntities[i - 1].isEnabled = false
                continue
            }
            boneEntities[i - 1].isEnabled = true

            // Position at midpoint
            boneEntities[i - 1].position = (pPos + cPos) * 0.5

            // Scale box along Y to match bone length
            boneEntities[i - 1].scale = SIMD3(1, length / 0.05, 1)

            // Orient box's Y axis along bone direction
            let up = SIMD3<Float>(0, 1, 0)
            let dir = simd_normalize(diff)
            let dot = simd_dot(up, dir)
            if abs(dot) > 0.9999 {
                if dot > 0 {
                    boneEntities[i - 1].orientation = simd_quatf(ix: 0, iy: 0, iz: 0, r: 1)
                } else {
                    boneEntities[i - 1].orientation = simd_quatf(angle: .pi, axis: SIMD3(1, 0, 0))
                }
            } else {
                let axis = simd_normalize(simd_cross(up, dir))
                let angle = acos(max(-1, min(1, dot)))
                boneEntities[i - 1].orientation = simd_quatf(angle: angle, axis: axis)
            }
        }
    }
}
