src/avatar.js

// @ts-check
/**
 * The avatar stage: a TalkingHead 3D avatar with real-time audio-driven
 * lip-sync (HeadAudio), plus static model fallback for custom GLBs.
 *
 * GLB ARCHITECTURE (avaturn.me exports):
 *   Type A (With morph targets) — ~13MB, has Head_Mesh with 72 ARKit
 *   viseme blend shapes, separate Eye/Mouth/Tongue meshes, Mixamo skeleton.
 *   → Load via TalkingHead.showAvatar() for lip-sync + expressions.
 *
 *   Type B (Without morph targets) — ~4MB, single avaturn_look_0 block,
 *   no blend shapes at all. showAvatar() throws 'Blend shapes not found'.
 *   → Load via THREE.GLTFLoader directly into TalkingHead's scene.
 *
 * CLASSIFICATION: Pre-fetch GLB as ArrayBuffer, parse with GLTFLoader,
 * inspect for morphTargetInfluences. One network call, zero double-load.
 */

import { TalkingHead } from "@met4citizen/talkinghead";
import { HeadAudio } from "./vendor/headaudio.min.mjs";
import { Box3, Vector3, HemisphereLight, DirectionalLight } from "three";
import { GLTFLoader } from "three/addons/loaders/GLTFLoader.js";

const HEADAUDIO_WORKLET_URL = "/vendor/headworklet.min.mjs";
const HEADAUDIO_MODEL_URL = "/vendor/model-en-mixed.bin";

export const AVATAR_MOODS = [
  "neutral", "happy", "angry", "sad", "fear", "disgust", "love", "sleep",
];

export const AVATAR_GESTURES = [
  "handup", "index", "ok", "thumbup", "thumbdown", "side", "shrug",
];

const HEAD_CONFIG = {
  ttsEndpoint: "N/A",
  lipsyncModules: [],
  cameraView: "upper",
  cameraDistance: -1.4,
  cameraY: -0.15,
  cameraX: -0.18,
  cameraRotateEnable: false,
  lightAmbientIntensity: 0,
  lightDirectIntensity: 0,
  lightSpotIntensity: 0,
  avatarIdleEyeContact: 0.3,
  avatarSpeakingEyeContact: 0.7,
};

/** Fully dispose all Three.js + AudioContext resources. */
function _fullyDispose(head) {
  try {
    if (head.scene) {
      head.scene.traverse((obj) => {
        if (obj.isMesh) {
          obj.geometry?.dispose();
          if (obj.material) {
            (Array.isArray(obj.material) ? obj.material : [obj.material]).forEach((m) => m.dispose());
          }
        }
      });
    }
    while (head.scene?.children?.length) head.scene.remove(head.scene.children[0]);
    if (head.renderer) {
      head.renderer.dispose();
      const canvas = head.renderer.domElement;
      if (canvas?.parentNode) canvas.parentNode.removeChild(canvas);
    }
    if (head.audioCtx && head.audioCtx.state !== "closed") {
      head.audioCtx.close().catch(() => {});
    }
    head.dispose?.();
  } catch (e) {
    console.warn("[AvatarStage] dispose:", e);
  }
}

export class AvatarStage {
  /** @param {HTMLElement} container */
  constructor(container) {
    this._container = container;
    /** @type {TalkingHead | null} */
    this.head = null;
    /** @type {any | null} */
    this._headaudio = null;
    this._lastSpeechEnded = 0;
    this._loadingEl = document.getElementById("loading");
    this._lastError = null;
    /** @type {boolean} True when model is static (no morph targets) */
    this._staticModel = false;
    /** @type {THREE.Object3D | null} */
    this._modelRoot = null;
    /** @type {THREE.Light[]} */
    this._staticLights = [];
  }

  /**
   * Initialize the avatar. Returns true on success, false on failure.
   *
   * Loading strategy:
   *   Default model (vuong.glb) — has morph targets → showAvatar() directly
   *   Custom models — pre-fetch + classify:
   *     Type A (has morphs) → Blob URL → showAvatar()
   *     Type B (no morphs)  → add to scene from buffer directly, skip showAvatar()
   */
  async init(opt = {}) {
    // ── FULL DISPOSE ────────────────────────────────────────────────
    if (this.head) {
      this._removeStaticLights();
      _fullyDispose(this.head);
      this.head = null;
      this._headaudio = null;
      this._modelRoot = null;
      this._staticModel = false;
    }
    while (this._container.firstChild) this._container.removeChild(this._container.firstChild);

    this._lastError = null;
    const isDefault = !opt.avatarUrl;
    const avatarUrl = opt.avatarUrl ?? "/avatars/vuong.glb";

    /** @type {boolean} */
    let hasMorphs;
    /** @type {ArrayBuffer | null} */
    let buffer = null;

    if (isDefault) {
      // vuong.glb is known to have morph targets — skip fetch/parse
      hasMorphs = true;
    } else {
      // Custom model: classify via fetch + parseAsync
      try {
        const result = await this._classifyGLB(avatarUrl);
        hasMorphs = result.hasMorphs;
        buffer = result.buffer;
      } catch (err) {
        console.warn("[AvatarStage] Classification failed, assuming morph:", err);
        hasMorphs = true;
      }
    }

    this.head = new TalkingHead(this._container, {
      ...HEAD_CONFIG,
      cameraDistance: isDefault ? HEAD_CONFIG.cameraDistance : -1.25,
      cameraY: isDefault ? HEAD_CONFIG.cameraY : 0.0,
      cameraX: isDefault ? HEAD_CONFIG.cameraX : 0.0,
    });

    console.log("[AvatarStage] Loading:", avatarUrl, `(morphs=${hasMorphs})`);

    if (hasMorphs) {
      // ── Type A: Has morph targets → via showAvatar() for lip-sync ──
      let morphOk = false;
      try {
        if (isDefault) {
          await this._loadShowAvatar(avatarUrl, opt);
        } else {
          const blob = new Blob([/** @type {ArrayBuffer} */(buffer)], { type: "model/gltf-binary" });
          const blobUrl = URL.createObjectURL(blob);
          await this._loadShowAvatar(blobUrl, opt);
          URL.revokeObjectURL(blobUrl);
        }
        morphOk = true;
        console.log("[AvatarStage] Model loaded WITH morph targets");
      } catch (err) {
        console.error("[AvatarStage] showAvatar failed:", err?.message);
        if (isDefault) {
          this._lastError = err?.message || String(err);
          _fullyDispose(this.head);
          this.head = null;
          return false;
        }
      }

      if (morphOk) {
        try { await this._initLipsync(); } catch (err) { console.warn("[AvatarStage] Lipsync:", err); }
      } else if (buffer) {
        try {
          await this._loadStaticFromBuffer(buffer, opt);
          this._staticModel = true;
          this._addStaticLights();
        } catch (err2) {
          this._lastError = err2?.message || String(err2);
          _fullyDispose(this.head);
          this.head = null;
          return false;
        }
      } else {
        this._lastError = "showAvatar failed and no buffer available";
        _fullyDispose(this.head);
        this.head = null;
        return false;
      }
    } else {
      // ── Type B: No morph targets → load as static ─────────────────
      try {
        await this._loadStaticFromBuffer(/** @type {ArrayBuffer} */(buffer), opt);
        this._staticModel = true;
        this._addStaticLights();
        console.log("[AvatarStage] Model loaded as STATIC (no morph targets)");
      } catch (err) {
        this._lastError = err?.message || String(err);
        console.error("[AvatarStage] Static load failed:", this._lastError);
        _fullyDispose(this.head);
        this.head = null;
        return false;
      }
    }

    this._adjustCamera(avatarUrl);
    return true;
  }

  /**
   * Fetch a GLB and classify it: does it have morph targets or not?
   * Returns { hasMorphs, buffer } — one fetch, one parse.
   */
  async _classifyGLB(url) {
    const response = await fetch(url);
    const buffer = await response.arrayBuffer();

    const loader = new GLTFLoader();
    const gltf = await loader.parseAsync(buffer, "");

    let hasMorphs = false;
    gltf.scene.traverse((child) => {
      if (child.isMesh && child.morphTargetInfluences) {
        hasMorphs = child.morphTargetInfluences.length > 0;
      }
      if (!hasMorphs && child.isMesh && child.geometry && child.geometry.morphAttributes) {
        const attrs = child.geometry.morphAttributes;
        hasMorphs = !!(attrs.position && attrs.position.length > 0);
      }
    });

    gltf.scene.traverse((child) => {
      if (child.isMesh && child.geometry) child.geometry.dispose();
    });

    return { hasMorphs, buffer };
  }

  /** Load morph-target model via TalkingHead showAvatar */
  async _loadShowAvatar(url, opt) {
    const head = /** @type {TalkingHead} */ (this.head);
    await head.showAvatar(
      { url, body: opt.body ?? "F", avatarMood: "neutral" },
      (ev) => {
        if (ev.lengthComputable) {
          const pct = Math.min(100, Math.round((ev.loaded / ev.total) * 100));
          if (this._loadingEl) this._loadingEl.textContent = `Loading avatar ${pct}%`;
        }
        opt.onprogress?.(ev);
      },
    );
  }

  /** Load GLB as a static model via THREE.GLTFLoader from pre-fetched ArrayBuffer.
   *  Only called for Type B models (no morph targets).
   *  Uses the buffer already fetched during classification — zero extra network calls. */
  _loadStaticFromBuffer(buffer, opt) {
    return new Promise((resolve, reject) => {
      const loader = new GLTFLoader();
      const head = /** @type {TalkingHead} */ (this.head);

      loader.parse(
        buffer,
        "",
        (gltf) => {
          const model = gltf.scene;
          this._modelRoot = model;

          for (let i = head.scene.children.length - 1; i >= 0; i--) {
            const c = head.scene.children[i];
            if (c !== head.camera) head.scene.remove(c);
          }

          model.visible = true;
          head.scene.add(model);

          if (this._loadingEl) this._loadingEl.textContent = "Positioning model...";
          opt?.onprogress?.({ lengthComputable: true, loaded: 1, total: 1 });
          resolve();
        },
        (err) => reject(new Error(`GLTF parse: ${err?.message || err}`)),
      );
    });
  }

  /** Add lighting for static models (TalkingHead lights are disabled at intensity 0) */
  _addStaticLights() {
    this._removeStaticLights();
    const head = /** @type {TalkingHead} */ (this.head);
    const lights = [
      new HemisphereLight(0xffffff, 0x444444, 0.7),
      new DirectionalLight(0xffffff, 0.9),
    ];
    lights[1].position.set(5, 8, 5);
    const fill = new DirectionalLight(0x8888ff, 0.3);
    fill.position.set(-3, 2, -4);
    lights.push(fill);
    for (const l of lights) { head.scene.add(l); this._staticLights.push(l); }
    console.log("[AvatarStage] Added", lights.length, "lights for static model");
  }

  _removeStaticLights() {
    for (const l of this._staticLights) {
      l.parent?.remove(l);
      l.dispose?.();
    }
    this._staticLights = [];
  }

  /** Auto-adjust camera + model position based on bounding box */
  _adjustCamera(avatarUrl) {
    const head = /** @type {TalkingHead} */ (this.head);
    if (!head.scene || !head.camera) return;

    if (!avatarUrl && !this._staticModel) {
      requestAnimationFrame(() => {
        try { head.renderer?.render(head.scene, head.camera); head.camera.updateProjectionMatrix?.(); } catch {}
      });
      return;
    }

    try {
      const box = new Box3().setFromObject(head.scene);
      const size = box.getSize(new Vector3());
      const center = box.getCenter(new Vector3());
      console.log("[AvatarStage] BBox:", { size, center });

      if (size.x === 0 && size.y === 0 && size.z === 0) {
        console.warn("[AvatarStage] Empty bbox, skipping");
        return;
      }

      if (size.y > 2.0) {
        head.scene.position.y = -(center.y - size.y * 0.35);
        head.camera.position.z = -Math.max(size.z, size.y, 1.0) * 1.2;
        head.camera.position.y = size.y * 0.3;
        console.log("[AvatarStage] Full-body adjusted");
      } else if (size.y > 0.3) {
        head.camera.position.z = -Math.max(size.z, 0.5) * 2.5;
        head.camera.position.y = center.y * 0.5;
        console.log("[AvatarStage] Head model adjusted");
      }

      requestAnimationFrame(() => {
        try { head.renderer?.render(head.scene, head.camera); head.camera.updateProjectionMatrix?.(); } catch {}
      });
      box.dispose();
    } catch (e) {
      console.warn("[AvatarStage] Camera adjust:", e);
    }
  }

  async _initLipsync() {
    const head = /** @type {TalkingHead} */ (this.head);
    if (!head.audioCtx) { console.warn("[AvatarStage] No AudioContext"); return; }
    try { await head.audioCtx.audioWorklet.addModule(HEADAUDIO_WORKLET_URL); } catch (e) { console.warn("[AvatarStage] Worklet:", e); return; }

    const headaudio = new HeadAudio(head.audioCtx);
    await headaudio.loadModel(HEADAUDIO_MODEL_URL);
    if (!head.audioSpeechGainNode) { console.warn("[AvatarStage] No audioSpeechGainNode"); return; }
    head.audioSpeechGainNode.connect(headaudio);

    headaudio.onvalue = (key, value) => {
      const mt = head.mtAvatar?.[key];
      if (mt) Object.assign(mt, { newvalue: value, needsUpdate: true });
    };
    head.opt.update = headaudio.update.bind(headaudio);
    headaudio.onended = () => { this._lastSpeechEnded = Date.now(); };
    headaudio.onstarted = () => {
      if (Date.now() - this._lastSpeechEnded > 150) { head.lookAtCamera(500); head.speakWithHands(); }
    };
    this._headaudio = headaudio;
  }

  get audioCtx() { return this.head?.audioCtx ?? null; }
  get voiceSink() { return this.head?.audioAnalyzerNode ?? null; }
  get lastError() { return this._lastError; }
  get isStaticModel() { return this._staticModel; }

  resume() {
    this.head?.start();
    if (this.head && this.head.audioCtx.state === "suspended") {
      return this.head.audioCtx.resume().catch(() => {});
    }
    return Promise.resolve();
  }

  setConversationState(status) {
    const h = this.head;
    if (!h) return;
    switch (status) {
      case "user-speaking": h.isSpeaking = false; h.lookAtCamera(800); break;
      case "ai-speaking": h.isSpeaking = true; break;
      case "processing": h.isSpeaking = false; break;
      case "closed": case "error": case "idle": h.isSpeaking = false; h.stopGesture(300); break;
      default: h.isSpeaking = false;
    }
  }

  runTool(name, args) {
    const h = this.head;
    if (!h) return null;
    if (this._staticModel) {
      if (name === "set_mood") return "Mood unavailable (static model)";
      if (name === "make_hand_gesture") return "Gestures unavailable (static model)";
      if (name === "make_facial_expression") return "Expressions unavailable (static model)";
    }
    if (name === "set_mood") {
      if (!AVATAR_MOODS.includes(args.mood)) return `Unknown mood: ${args.mood}`;
      h.setMood(args.mood); return `Mood set to ${args.mood}.`;
    }
    if (name === "make_hand_gesture") {
      if (!AVATAR_GESTURES.includes(args.gesture)) return `Unknown gesture: ${args.gesture}`;
      h.playGesture(args.gesture, 3); return `Playing gesture ${args.gesture}.`;
    }
    if (name === "make_facial_expression") {
      if (!args.emoji) return "No emoji given.";
      h.speakEmoji(args.emoji); return `Expressing ${args.emoji}.`;
    }
    return null;
  }
}