import * as THREE from "three";

/**
 * Unit-frame portraits: each character's real head, drawn ONCE by a close camera into a 128x128 target and kept as a
 * data URL. The HUD reads it through `portraits` (a tiny store). Never runs per frame: `flush()` only does work for a
 * portrait that was requested and has waited for its body to load and settle.
 */
const SIZE = 128;
const SETTLE = 0.45; // seconds after the body is ready, so the rig has left its rest pose

type Listener = () => void;
interface Want {
  key: string;
  root: THREE.Object3D;
  age: number;
  tries: number;
}

class PortraitStore {
  private urls = new Map<string, string>();
  private ver = 0;
  private subs = new Set<Listener>();
  private want: Want[] = [];
  private rt: THREE.WebGLRenderTarget | null = null;
  private cam = new THREE.PerspectiveCamera(22, 1, 0.05, 20);
  private buf = new Uint8Array(SIZE * SIZE * 4);
  private head = new THREE.Vector3();
  private fwd = new THREE.Vector3();
  private canvas: HTMLCanvasElement | null = null;
  private last = 0;

  get = (key: string): string | null => this.urls.get(key) ?? null;
  version = () => this.ver;
  subscribe = (fn: Listener) => {
    this.subs.add(fn);
    return () => this.subs.delete(fn);
  };

  /** Ask for a portrait of `root` (a loaded-or-loading walker root). Ignored when the key already has one. */
  request(key: string, root: THREE.Object3D, ready: Promise<void>) {
    if (this.urls.has(key) || this.want.some((w) => w.key === key)) return;
    const w: Want = { key, root, age: -1, tries: 0 };
    this.want.push(w);
    void ready.then(() => (w.age = 0), () => (w.age = -2));
  }

  /** Call once per frame after the main render; cheap when nothing is waiting. */
  flush(r: THREE.WebGLRenderer, scene: THREE.Scene, now: number) {
    if (!this.want.length) return;
    const dt = Math.min(0.2, Math.max(0, (now - this.last) / 1000));
    this.last = now;
    for (let i = this.want.length - 1; i >= 0; i--) {
      const w = this.want[i];
      if (w.age === -2) {
        console.warn(`[portraits] body of ${w.key} failed to load`);
        this.want.splice(i, 1);
        continue;
      }
      if (w.age < 0) continue;
      w.age += dt;
      if (w.age < SETTLE) continue;
      if (this.snap(r, scene, w)) this.want.splice(i, 1);
      else if (++w.tries > 20) {
        console.warn(`[portraits] no head bone for ${w.key}`);
        this.want.splice(i, 1);
      }
      return; // one capture per frame
    }
  }

  private findHead(root: THREE.Object3D): THREE.Object3D | null {
    let found: THREE.Object3D | null = null;
    root.traverse((o) => {
      if (!found && (o as THREE.Bone).isBone && o.name.toLowerCase() === "head") found = o;
    });
    return found;
  }

  private snap(r: THREE.WebGLRenderer, scene: THREE.Scene, w: Want): boolean {
    const bone = this.findHead(w.root);
    if (!bone) return false;
    w.root.updateMatrixWorld(true);
    bone.getWorldPosition(this.head);
    this.head.y += 0.07;
    w.root.getWorldDirection(this.fwd);
    this.fwd.y = 0;
    this.fwd.normalize();
    const cam = this.cam;
    cam.position.copy(this.head).addScaledVector(this.fwd, 1.05);
    cam.position.y += 0.02;
    cam.lookAt(this.head);
    cam.updateMatrixWorld(true);
    if (!this.rt) this.rt = new THREE.WebGLRenderTarget(SIZE, SIZE, { depthBuffer: true });
    const prevRT = r.getRenderTarget();
    const prevBg = scene.background;
    const prevTone = r.toneMapping;
    scene.background = null;
    r.toneMapping = THREE.NoToneMapping;
    r.setRenderTarget(this.rt);
    r.setClearColor(0x0b1118, 1);
    r.clear();
    r.render(scene, cam);
    r.readRenderTargetPixels(this.rt, 0, 0, SIZE, SIZE, this.buf);
    r.setRenderTarget(prevRT);
    r.toneMapping = prevTone;
    scene.background = prevBg;
    this.finish(w.key);
    return true;
  }

  /** Linear bytes to sRGB, flipped upright, with a soft vignette, then a data URL. */
  private finish(key: string) {
    const c = (this.canvas ??= document.createElement("canvas"));
    c.width = c.height = SIZE;
    const ctx = c.getContext("2d");
    if (!ctx) return;
    const img = ctx.createImageData(SIZE, SIZE);
    const b = this.buf;
    for (let y = 0; y < SIZE; y++) {
      const sy = SIZE - 1 - y;
      for (let x = 0; x < SIZE; x++) {
        const si = (sy * SIZE + x) * 4, di = (y * SIZE + x) * 4;
        const dx = (x - SIZE / 2) / (SIZE / 2), dy = (y - SIZE / 2) / (SIZE / 2);
        const vig = 1 - 0.35 * Math.min(1, dx * dx + dy * dy);
        for (let k = 0; k < 3; k++) {
          const lin = Math.min(1, (b[si + k] / 255) * 1.5);
          img.data[di + k] = Math.round(255 * Math.pow(lin, 1 / 2.2) * vig);
        }
        img.data[di + 3] = 255;
      }
    }
    ctx.putImageData(img, 0, 0);
    this.urls.set(key, c.toDataURL("image/jpeg", 0.88));
    this.ver++;
    for (const f of this.subs) f();
  }

  dispose() {
    this.rt?.dispose();
    this.rt = null;
    this.want.length = 0;
  }
}

export const portraits = new PortraitStore();
