import { CANONICAL_HIERARCHY, type CanonicalBone } from "@psx/shared";
import type { MeshData, Vec3Tuple } from "../mesh/MeshBuilder.js";
import { vertexCount } from "../mesh/MeshBuilder.js";
import type { Skeleton } from "./skeleton.js";

/**
 * Automatic skin weights — the auto-rig half of GDD §10.
 *
 * Weights come from distance to a bone's *segment* (parent joint to child
 * joint), not to the joint point. Distance-to-point is the obvious approach and
 * it is wrong: every vertex along a long bone would be pulled toward whichever
 * end it happened to sit nearer, so a forearm would tear in half at its middle.
 * Distance to the segment is uniform along the bone's length, which is what
 * "belongs to this bone" actually means.
 *
 * glTF allows four influences per vertex, which is also more than a PSX-era
 * character needs — most vertices end up with one or two.
 */

export interface SkinOptions {
  /**
   * Falloff exponent. Higher gives crisper, more rigid joints; lower gives
   * softer, more rubbery bending. 4 reads as "articulated puppet", which is
   * exactly the era's look.
   */
  falloff?: number;
  /** Influences below this fraction of the strongest are dropped. */
  cullRatio?: number;
}

const MAX_INFLUENCES = 4;

/** Bones that should never claim a vertex — they have no geometry of their own. */
const NON_DEFORMING: ReadonlySet<CanonicalBone> = new Set<CanonicalBone>([]);

/**
 * Extra distance charged per metre a vertex sits on the far side of the body
 * from a left/right bone.
 *
 * Without it the opposite leg still clears the cull ratio at a 17 cm stance —
 * bind pose looks fine, and then the walk cycle drags heel vertices toward the
 * other foot as the legs separate. Scaling with the offset rather than
 * rejecting outright keeps midline geometry (pelvis, groin) smooth.
 */
const MIRROR_PENALTY = 2;

interface SegmentShape {
  /** Extend the scoring segment backward past its own joint, in bone lengths. */
  headExtend?: number;
  /** Pull the scoring segment short of the child joint, in bone lengths. */
  tailTrim?: number;
}

/**
 * Where a bone's *scoring* segment differs from the bone itself.
 *
 * Distance-to-segment cannot separate two bones that meet at a joint: anything
 * near that joint is equidistant from both, and splits 50/50. Mid-limb that is
 * harmless. At a joint where the geometry continues past it in a direction only
 * the child owns, it tears — the heel sits behind and below the ankle, so half
 * of every heel vertex was driven by the shin and the boot sheared into a
 * wedge as the ankle pitched.
 *
 * The fix is to say so explicitly, per bone, rather than to invent a rule.
 */
const SEGMENT_SHAPE: Partial<Record<CanonicalBone, SegmentShape>> = {
  // The foot's mass extends back to the heel, ~45% of the ankle→toe length
  // behind the ankle. Claim it, and keep the shin off it.
  LeftFoot: { headExtend: 0.55 },
  RightFoot: { headExtend: 0.55 },
  LeftLeg: { tailTrim: 0.12 },
  RightLeg: { tailTrim: 0.12 },
  // Same story, milder: the palm continues past the wrist.
  LeftForeArm: { tailTrim: 0.1 },
  RightForeArm: { tailTrim: 0.1 },
};

/** +1 for a left-side bone, −1 for a right-side one, 0 for the spine. */
function boneSide(name: CanonicalBone): number {
  if (name.startsWith("Left")) return 1;
  if (name.startsWith("Right")) return -1;
  return 0;
}

interface BoneSegment {
  jointIndex: number;
  name: CanonicalBone;
  head: Vec3Tuple;
  tail: Vec3Tuple;
  /** +1 left, −1 right, 0 central. */
  side: number;
}

/**
 * A bone's influence region runs from its own joint to its child's. Leaf bones
 * (hands, toes, head) have no child, so they get a stub extending along the
 * direction they inherit from their parent.
 */
function buildSegments(skeleton: Skeleton): BoneSegment[] {
  const segments: BoneSegment[] = [];

  for (const [jointIndex, joint] of skeleton.joints.entries()) {
    if (NON_DEFORMING.has(joint.name)) continue;

    const children = skeleton.joints.filter((other) => CANONICAL_HIERARCHY[other.name] === joint.name);

    if (children.length > 0) {
      // Average the children so a fork (Hips → both legs, Spine2 → both
      // shoulders) points down the body's axis rather than at one branch.
      const tail: Vec3Tuple = [0, 0, 0];
      for (const child of children) {
        tail[0] += child.world[0] / children.length;
        tail[1] += child.world[1] / children.length;
        tail[2] += child.world[2] / children.length;
      }
      segments.push(shape(jointIndex, joint.name, joint.world, tail));
      continue;
    }

    const parent = joint.parent ? skeleton.joints[skeleton.index.get(joint.parent)!] : null;
    const direction: Vec3Tuple = parent
      ? [
          joint.world[0] - parent.world[0],
          joint.world[1] - parent.world[1],
          joint.world[2] - parent.world[2],
        ]
      : [0, 1, 0];

    const length = Math.hypot(...direction) || 1;
    const stub = skeleton.height * 0.05;
    segments.push(
      shape(jointIndex, joint.name, joint.world, [
        joint.world[0] + (direction[0] / length) * stub,
        joint.world[1] + (direction[1] / length) * stub,
        joint.world[2] + (direction[2] / length) * stub,
      ]),
    );
  }

  return segments;
}

/** Apply this bone's {@link SEGMENT_SHAPE} entry to its raw head/tail. */
function shape(
  jointIndex: number,
  name: CanonicalBone,
  head: Vec3Tuple,
  tail: Vec3Tuple,
): BoneSegment {
  const adjust = SEGMENT_SHAPE[name];
  const side = boneSide(name);
  if (!adjust) return { jointIndex, name, head, tail, side };

  const dx = tail[0] - head[0];
  const dy = tail[1] - head[1];
  const dz = tail[2] - head[2];
  const headExtend = adjust.headExtend ?? 0;
  const tailTrim = adjust.tailTrim ?? 0;

  return {
    jointIndex,
    name,
    side,
    head: [head[0] - dx * headExtend, head[1] - dy * headExtend, head[2] - dz * headExtend],
    tail: [tail[0] - dx * tailTrim, tail[1] - dy * tailTrim, tail[2] - dz * tailTrim],
  };
}

/** Shortest distance from a point to a finite segment. */
function distanceToSegment(point: Vec3Tuple, head: Vec3Tuple, tail: Vec3Tuple): number {
  const dx = tail[0] - head[0];
  const dy = tail[1] - head[1];
  const dz = tail[2] - head[2];
  const lengthSq = dx * dx + dy * dy + dz * dz;

  if (lengthSq < 1e-12) return Math.hypot(point[0] - head[0], point[1] - head[1], point[2] - head[2]);

  let t =
    ((point[0] - head[0]) * dx + (point[1] - head[1]) * dy + (point[2] - head[2]) * dz) / lengthSq;
  t = t < 0 ? 0 : t > 1 ? 1 : t;

  return Math.hypot(
    point[0] - (head[0] + dx * t),
    point[1] - (head[1] + dy * t),
    point[2] - (head[2] + dz * t),
  );
}

export function computeSkinWeights(
  mesh: MeshData,
  skeleton: Skeleton,
  options: SkinOptions = {},
): void {
  const falloff = options.falloff ?? 4;
  const cullRatio = options.cullRatio ?? 0.08;
  const segments = buildSegments(skeleton);

  // Keeps the inverse-distance weight finite for a vertex sitting exactly on a
  // bone, and sets the scale at which "close" stops meaning anything.
  const epsilon = skeleton.height * 0.004;

  const count = vertexCount(mesh);
  mesh.joints = new Array<number>(count * 4).fill(0);
  mesh.weights = new Array<number>(count * 4).fill(0);

  const scored: Array<{ jointIndex: number; weight: number }> = [];

  for (let v = 0; v < count; v++) {
    const point: Vec3Tuple = [
      mesh.positions[v * 3]!,
      mesh.positions[v * 3 + 1]!,
      mesh.positions[v * 3 + 2]!,
    ];

    scored.length = 0;
    let strongest = 0;

    for (const segment of segments) {
      // Charge for being on the wrong side of the body, so the far leg cannot
      // claim a share of this one and drag it across during a stride.
      const wrongSide = segment.side === 0 ? 0 : Math.max(0, -segment.side * point[0]);
      const distance =
        distanceToSegment(point, segment.head, segment.tail) + wrongSide * MIRROR_PENALTY + epsilon;
      const weight = 1 / Math.pow(distance, falloff);
      if (weight > strongest) strongest = weight;
      scored.push({ jointIndex: segment.jointIndex, weight });
    }

    // Keep the four strongest, discarding anything negligible next to the best.
    const kept = scored
      .filter((entry) => entry.weight >= strongest * cullRatio)
      .sort((a, b) => b.weight - a.weight)
      .slice(0, MAX_INFLUENCES);

    const total = kept.reduce((sum, entry) => sum + entry.weight, 0) || 1;

    for (let i = 0; i < MAX_INFLUENCES; i++) {
      const entry = kept[i];
      mesh.joints[v * 4 + i] = entry?.jointIndex ?? 0;
      mesh.weights[v * 4 + i] = entry ? entry.weight / total : 0;
    }
  }
}

/**
 * Inverse bind matrices: the transform taking a vertex from model space into
 * each joint's local space at rest. Rest poses here are translation-only, so
 * the inverse is a translation by the negated world position — no general
 * matrix inversion required.
 */
export function computeInverseBindMatrices(skeleton: Skeleton): Float32Array {
  const matrices = new Float32Array(skeleton.joints.length * 16);

  for (const [i, joint] of skeleton.joints.entries()) {
    const offset = i * 16;
    // Column-major identity.
    matrices[offset] = 1;
    matrices[offset + 5] = 1;
    matrices[offset + 10] = 1;
    matrices[offset + 15] = 1;
    // Translation occupies the last column in column-major order.
    matrices[offset + 12] = -joint.world[0];
    matrices[offset + 13] = -joint.world[1];
    matrices[offset + 14] = -joint.world[2];
  }

  return matrices;
}

/** Diagnostics for the pipeline report and for tests. */
export function skinStatistics(mesh: MeshData): {
  maxInfluences: number;
  averageInfluences: number;
  unweighted: number;
} {
  const count = vertexCount(mesh);
  let maxInfluences = 0;
  let totalInfluences = 0;
  let unweighted = 0;

  for (let v = 0; v < count; v++) {
    let influences = 0;
    let sum = 0;
    for (let i = 0; i < 4; i++) {
      const weight = mesh.weights[v * 4 + i]!;
      if (weight > 0) influences++;
      sum += weight;
    }
    if (sum < 0.999) unweighted++;
    totalInfluences += influences;
    maxInfluences = Math.max(maxInfluences, influences);
  }

  return {
    maxInfluences,
    averageInfluences: count > 0 ? totalInfluences / count : 0,
    unweighted,
  };
}
