import { PNG } from "pngjs";
import { vertexCount, type MeshData } from "../mesh/MeshBuilder.js";

export interface UvGuideOptions {
  size?: number;
  /** Draw wireframe edges over the filled islands. */
  wireframe?: boolean;
}

/**
 * Rasterise the mesh's UVs into a guide sheet for Imagine:
 *  - fill each triangle with its vertex colour (linear→approx sRGB)
 *  - optional white wireframe so islands stay legible
 *  - skip degenerate / zero-area UV tris (chart borders)
 *  - dark background outside islands
 */
export function bakeUvGuide(mesh: MeshData, options: UvGuideOptions = {}): Buffer {
  if (mesh.uvs.length < vertexCount(mesh) * 2) {
    throw new Error("Mesh has no UVs — run unwrapCharacterAtlas first");
  }

  const size = options.size ?? 512;
  const wireframe = options.wireframe ?? true;
  const png = new PNG({ width: size, height: size });
  const data = png.data;

  // Dark charcoal background.
  for (let i = 0; i < size * size; i++) {
    data[i * 4] = 22;
    data[i * 4 + 1] = 24;
    data[i * 4 + 2] = 26;
    data[i * 4 + 3] = 255;
  }

  const toPx = (u: number, v: number): [number, number] => {
    // Match glTF: (0,0) = upper-left of the image, V increases downward.
    // Our unwrap puts the head at low V and feet at high V, so the guide
    // sheet reads head-at-top the same way the mesh samples the albedo.
    const x = Math.min(size - 1, Math.max(0, Math.floor(u * (size - 1) + 0.5)));
    const y = Math.min(size - 1, Math.max(0, Math.floor(v * (size - 1) + 0.5)));
    return [x, y];
  };

  const linToByte = (c: number) => {
    // Rough linear→sRGB so skin/cloth read closer to the vertex palette.
    const s = c <= 0.0031308 ? 12.92 * c : 1.055 * Math.pow(c, 1 / 2.4) - 0.055;
    return Math.max(0, Math.min(255, Math.round(s * 255)));
  };

  // Painter's algorithm: draw back-facing / deeper tris first so the front
  // silhouette of each island wins in the guide (cleaner for Imagine).
  type Tri = { ia: number; ib: number; ic: number; depth: number };
  const tris: Tri[] = [];
  for (let t = 0; t < mesh.indices.length; t += 3) {
    const ia = mesh.indices[t]!;
    const ib = mesh.indices[t + 1]!;
    const ic = mesh.indices[t + 2]!;
    const ua = mesh.uvs[ia * 2]!;
    const va = mesh.uvs[ia * 2 + 1]!;
    const ub = mesh.uvs[ib * 2]!;
    const vb = mesh.uvs[ib * 2 + 1]!;
    const uc = mesh.uvs[ic * 2]!;
    const vc = mesh.uvs[ic * 2 + 1]!;
    if (Math.max(ua, ub, uc) - Math.min(ua, ub, uc) > 0.4) continue;
    if (Math.max(va, vb, vc) - Math.min(va, vb, vc) > 0.4) continue;
    // Depth = average model Z (front = larger Z draws later).
    const depth =
      (mesh.positions[ia * 3 + 2]! + mesh.positions[ib * 3 + 2]! + mesh.positions[ic * 3 + 2]!) / 3;
    tris.push({ ia, ib, ic, depth });
  }
  tris.sort((a, b) => a.depth - b.depth);

  for (const tri of tris) {
    const { ia, ib, ic } = tri;
    const [ax, ay] = toPx(mesh.uvs[ia * 2]!, mesh.uvs[ia * 2 + 1]!);
    const [bx, by] = toPx(mesh.uvs[ib * 2]!, mesh.uvs[ib * 2 + 1]!);
    const [cx, cy] = toPx(mesh.uvs[ic * 2]!, mesh.uvs[ic * 2 + 1]!);

    const area = (bx - ax) * (cy - ay) - (cx - ax) * (by - ay);
    if (Math.abs(area) < 0.5) continue;

    const r = linToByte((mesh.colors[ia * 4]! + mesh.colors[ib * 4]! + mesh.colors[ic * 4]!) / 3);
    const g = linToByte(
      (mesh.colors[ia * 4 + 1]! + mesh.colors[ib * 4 + 1]! + mesh.colors[ic * 4 + 1]!) / 3,
    );
    const b = linToByte(
      (mesh.colors[ia * 4 + 2]! + mesh.colors[ib * 4 + 2]! + mesh.colors[ic * 4 + 2]!) / 3,
    );

    fillTriangle(data, size, ax, ay, bx, by, cx, cy, r, g, b);
    if (wireframe) {
      drawLine(data, size, ax, ay, bx, by, 235, 235, 230);
      drawLine(data, size, bx, by, cx, cy, 235, 235, 230);
      drawLine(data, size, cx, cy, ax, ay, 235, 235, 230);
    }
  }

  return PNG.sync.write(png);
}

function setPx(
  data: Buffer,
  size: number,
  x: number,
  y: number,
  r: number,
  g: number,
  b: number,
): void {
  if (x < 0 || y < 0 || x >= size || y >= size) return;
  const i = (y * size + x) * 4;
  data[i] = r;
  data[i + 1] = g;
  data[i + 2] = b;
  data[i + 3] = 255;
}

/** Barycentric fill. */
function fillTriangle(
  data: Buffer,
  size: number,
  ax: number,
  ay: number,
  bx: number,
  by: number,
  cx: number,
  cy: number,
  r: number,
  g: number,
  b: number,
): void {
  const minX = Math.max(0, Math.min(ax, bx, cx));
  const maxX = Math.min(size - 1, Math.max(ax, bx, cx));
  const minY = Math.max(0, Math.min(ay, by, cy));
  const maxY = Math.min(size - 1, Math.max(ay, by, cy));
  const area = (bx - ax) * (cy - ay) - (cx - ax) * (by - ay);
  if (Math.abs(area) < 1e-6) return;

  for (let y = minY; y <= maxY; y++) {
    for (let x = minX; x <= maxX; x++) {
      const w0 = (bx - ax) * (y - ay) - (by - ay) * (x - ax);
      const w1 = (cx - bx) * (y - by) - (cy - by) * (x - bx);
      const w2 = (ax - cx) * (y - cy) - (ay - cy) * (x - cx);
      if (
        (w0 >= 0 && w1 >= 0 && w2 >= 0 && area > 0) ||
        (w0 <= 0 && w1 <= 0 && w2 <= 0 && area < 0)
      ) {
        setPx(data, size, x, y, r, g, b);
      }
    }
  }
}

function drawLine(
  data: Buffer,
  size: number,
  x0: number,
  y0: number,
  x1: number,
  y1: number,
  r: number,
  g: number,
  b: number,
): void {
  let dx = Math.abs(x1 - x0);
  let dy = Math.abs(y1 - y0);
  const sx = x0 < x1 ? 1 : -1;
  const sy = y0 < y1 ? 1 : -1;
  let err = dx - dy;
  let x = x0;
  let y = y0;
  for (;;) {
    setPx(data, size, x, y, r, g, b);
    if (x === x1 && y === y1) break;
    const e2 = 2 * err;
    if (e2 > -dy) {
      err -= dy;
      x += sx;
    }
    if (e2 < dx) {
      err += dx;
      y += sy;
    }
  }
}
