import { memo, useEffect, useMemo, useRef, type MutableRefObject } from "react";
import { useFrame } from "@react-three/fiber";
import {
  AdditiveBlending, BufferAttribute, BufferGeometry, ClampToEdgeWrapping, DoubleSide,
  Group, Matrix4, MeshBasicMaterial, MultiplyBlending, NormalBlending, Quaternion,
  RepeatWrapping, Texture, Vector3,
} from "three";
import { sampleOriginalCombatMeshTrack } from "./originalCombatMeshRuntime";
import { updateWowBoneMatrices, wowModelFrame, type OriginalWowModelGraph, type WowRibbon } from "./originalWowModelRuntime";

const blend = (mode: number) => [3, 4, 7].includes(mode) ? AdditiveBlending
  : [5, 6].includes(mode) ? MultiplyBlending : NormalBlending;

export const OriginalWowModel = memo(function OriginalWowModel({ graph, textures, active, playbackRef, loop, opacity = 1 }: {
  readonly graph: OriginalWowModelGraph;
  readonly textures: ReadonlyMap<number, Texture>;
  readonly active: boolean;
  readonly playbackRef: MutableRefObject<{ elapsedMs: number; decaying: boolean }>;
  readonly loop: boolean;
  readonly opacity?: number;
}) {
  const root = useRef<Group>(null);
  const scratch = useMemo(() => ({
    bones: graph.bones.map(() => new Matrix4()), position: new Vector3(), vertex: new Vector3(),
    normal: new Vector3(), normalSum: new Vector3(), camera: new Quaternion(), world: new Quaternion(),
    billboardAxis: new Quaternion().setFromAxisAngle(new Vector3(0, 1, 0), -Math.PI / 2),
  }), [graph]);
  const geometry = useMemo(() => {
    if (!graph.mesh) return null;
    const geometry = new BufferGeometry();
    geometry.setAttribute("position", new BufferAttribute(new Float32Array(graph.mesh.positions), 3));
    geometry.setAttribute("normal", new BufferAttribute(new Float32Array(graph.mesh.normals), 3));
    geometry.setAttribute("uv", new BufferAttribute(new Float32Array(graph.mesh.uvs), 2));
    geometry.setAttribute("uv1", new BufferAttribute(new Float32Array(graph.mesh.uvs2), 2));
    geometry.setIndex([...graph.mesh.indices]);
    graph.mesh.layers.forEach((layer, index) => geometry.addGroup(layer.indexStart, layer.indexCount, index));
    return geometry;
  }, [graph]);
  const materials = useMemo(() => graph.mesh?.layers.map((layer) => {
    const source = graph.materials[layer.materialIndex];
    const maps = layer.textureIndices.map((index, slot) => {
      const texture = textures.get(index)!.clone();
      texture.wrapS = graph.textures[index].flags & 1 ? RepeatWrapping : ClampToEdgeWrapping;
      texture.wrapT = graph.textures[index].flags & 2 ? RepeatWrapping : ClampToEdgeWrapping;
      texture.channel = Math.max(0, layer.textureCoordinateIndices[slot] ?? 0);
      texture.needsUpdate = true;
      return texture;
    });
    const material = new MeshBasicMaterial({
      map: maps[0], blending: blend(source.blendMode), transparent: source.blendMode !== 0,
      alphaTest: source.blendMode === 1 ? 0.5 : 0, depthWrite: !(source.flags & 16) && source.blendMode < 2,
      depthTest: !(source.flags & 8), side: DoubleSide, toneMapped: false,
    });
    // The WotLK spell batches use one or two source textures. Preserve the
    // second color/alpha layer and its independent UV transform in the shader.
    if (maps[1]) {
      material.onBeforeCompile = (shader) => {
        shader.uniforms.wowSecondMap = { value: maps[1] };
        shader.uniforms.wowSecondMatrix = { value: maps[1].matrix };
        shader.vertexShader = shader.vertexShader.replace("#include <common>",
          "#include <common>\nuniform mat3 wowSecondMatrix;\nvarying vec2 vWowSecondUv;"
          + (maps[1].channel === 1 ? "\n#ifndef USE_UV1\nattribute vec2 uv1;\n#endif" : ""));
        shader.vertexShader = shader.vertexShader.replace("#include <uv_vertex>",
          `#include <uv_vertex>\nvWowSecondUv = (wowSecondMatrix * vec3(${maps[1].channel === 1 ? "uv1" : "uv"}, 1.0)).xy;`);
        shader.fragmentShader = shader.fragmentShader.replace("#include <common>",
          "#include <common>\nuniform sampler2D wowSecondMap;\nvarying vec2 vWowSecondUv;");
        shader.fragmentShader = shader.fragmentShader.replace("#include <map_fragment>",
          "#include <map_fragment>\ndiffuseColor *= texture2D(wowSecondMap, vWowSecondUv);");
      };
      material.customProgramCacheKey = () => `wow-dual-${maps[1].channel}`;
    }
    return { material, maps };
  }) ?? [], [graph, textures]);
  useEffect(() => () => {
    geometry?.dispose();
    materials.forEach(({ material, maps }) => { material.dispose(); maps.forEach((texture) => texture.dispose()); });
  }, [geometry, materials]);
  useFrame((state) => {
    if (!active) return;
    const frame = wowModelFrame(graph, playbackRef.current.elapsedMs, loop, playbackRef.current.decaying);
    root.current?.getWorldQuaternion(scratch.world);
    state.camera.getWorldQuaternion(scratch.camera).premultiply(scratch.world.invert());
    // Authored billboard planes face source +X, whereas the camera faces +Z.
    scratch.camera.multiply(scratch.billboardAxis);
    updateWowBoneMatrices(graph, frame, scratch.bones, scratch.camera);
    if (!geometry || !graph.mesh) return;
    const mesh = graph.mesh;
    const positions = geometry.getAttribute("position") as BufferAttribute;
    const normals = geometry.getAttribute("normal") as BufferAttribute;
    for (let index = 0; index < mesh.positions.length / 3; index++) {
      scratch.position.set(0, 0, 0); scratch.normalSum.set(0, 0, 0);
      let total = 0;
      for (let slot = 0; slot < 4; slot++) {
        const weight = mesh.boneWeights[index * 4 + slot];
        if (!weight) continue;
        const matrix = scratch.bones[mesh.boneIndices[index * 4 + slot]];
        scratch.vertex.fromArray(mesh.positions, index * 3).applyMatrix4(matrix);
        scratch.normal.fromArray(mesh.normals, index * 3).transformDirection(matrix);
        scratch.position.addScaledVector(scratch.vertex, weight);
        scratch.normalSum.addScaledVector(scratch.normal, weight);
        total += weight;
      }
      if (!total) { scratch.position.fromArray(mesh.positions, index * 3); scratch.normalSum.fromArray(mesh.normals, index * 3); }
      positions.setXYZ(index, ...scratch.position.toArray());
      normals.setXYZ(index, ...scratch.normalSum.normalize().toArray());
    }
    positions.needsUpdate = true; normals.needsUpdate = true;
    mesh.layers.forEach((layer, index) => {
      const { material, maps } = materials[index];
      const color = mesh.colors[layer.colorIndex];
      const weight = mesh.textureWeights[layer.textureWeightIndex];
      const rgb = color ? sampleOriginalCombatMeshTrack(graph, color.color, frame) : undefined;
      material.color.setRGB(...(rgb ?? [1, 1, 1]));
      material.opacity = Math.max(0, Math.min(1, (color ? sampleOriginalCombatMeshTrack(graph, color.alpha, frame) ?? 32767 : 32767)
        * (weight ? sampleOriginalCombatMeshTrack(graph, weight, frame) ?? 32767 : 32767) / (32767 * 32767))) * opacity;
      maps.forEach((texture, slot) => {
        const transform = graph.textureTransforms[layer.textureTransformIndices[slot]];
        if (!transform) return;
        const t = sampleOriginalCombatMeshTrack(graph, transform.translation, frame) ?? [0, 0, 0];
        const s = sampleOriginalCombatMeshTrack(graph, transform.scale, frame) ?? [1, 1, 1];
        const q = sampleOriginalCombatMeshTrack(graph, transform.rotation, frame) ?? [0, 0, 0, 1];
        texture.offset.set(t[0], t[1]); texture.repeat.set(s[0], s[1]); texture.center.set(0.5, 0.5);
        texture.rotation = 2 * Math.atan2(q[2], q[3]); texture.updateMatrix();
      });
    });
  });
  return <group ref={root} rotation={[-Math.PI / 2, 0, 0]}>
    {geometry && <mesh name="original-wow-animated-mesh" geometry={geometry}
      material={materials.map((entry) => entry.material)} frustumCulled={false} />}
    {graph.ribbons.map((ribbon, index) => <OriginalWowRibbon key={index} graph={graph} ribbon={ribbon}
      textures={textures} active={active} bones={scratch.bones} playbackRef={playbackRef} loop={loop} opacity={opacity} />)}
  </group>;
});

const OriginalWowRibbon = memo(function OriginalWowRibbon({ graph, ribbon, textures, active, bones, playbackRef, loop, opacity }: {
  readonly graph: OriginalWowModelGraph; readonly ribbon: WowRibbon; readonly textures: ReadonlyMap<number, Texture>;
  readonly active: boolean; readonly bones: readonly Matrix4[];
  readonly playbackRef: MutableRefObject<{ elapsedMs: number; decaying: boolean }>; readonly loop: boolean;
  readonly opacity: number;
}) {
  const root = useRef<Group>(null);
  const capacity = Math.min(128, Math.max(4, Math.ceil(ribbon.edgesPerSecond * ribbon.edgeLifetime) + 2));
  const geometry = useMemo(() => {
    const g = new BufferGeometry();
    g.setAttribute("position", new BufferAttribute(new Float32Array(capacity * 6), 3));
    g.setAttribute("uv", new BufferAttribute(new Float32Array(capacity * 4), 2));
    g.setIndex(Array.from({ length: (capacity - 1) * 6 }, (_, i) => {
      const segment = Math.floor(i / 6) * 2;
      return segment + [0, 1, 2, 2, 1, 3][i % 6];
    }));
    g.setDrawRange(0, 0); return g;
  }, [capacity]);
  const material = useMemo(() => new MeshBasicMaterial({
    map: textures.get(ribbon.textureIndices[0]), transparent: true, depthWrite: false,
    side: DoubleSide, blending: blend(graph.materials[ribbon.materialIndices[0]]?.blendMode ?? 4), toneMapped: false,
  }), [graph, ribbon, textures]);
  const history = useMemo(() => ({
    edges: [] as { time: number; above: Vector3; below: Vector3 }[], inverse: new Matrix4(),
    top: new Vector3(), bottom: new Vector3(), local: new Vector3(), last: -Infinity,
  }), []);
  useEffect(() => () => { geometry.dispose(); material.dispose(); }, [geometry, material]);
  useFrame(() => {
    if (!active || !root.current) return;
    const elapsed = playbackRef.current.elapsedMs;
    const frame = wowModelFrame(graph, elapsed, loop, playbackRef.current.decaying);
    const sample = (key: "alpha" | "heightAbove" | "heightBelow" | "textureSlot" | "visibility") =>
      sampleOriginalCombatMeshTrack(graph, ribbon.tracks[key], frame);
    const lifetimeMs = Math.max(20, ribbon.edgeLifetime * 1000);
    if (elapsed < history.last) { history.edges.length = 0; history.last = -Infinity; }
    while (history.edges.length && elapsed - history.edges[0].time > lifetimeMs) history.edges.shift();
    root.current.updateWorldMatrix(true, false);
    if (elapsed - history.last >= 1000 / Math.max(1, ribbon.edgesPerSecond) && (sample("visibility") ?? 1) > 0) {
      const bone = bones[ribbon.parentBone];
      history.top.fromArray(ribbon.position); history.bottom.copy(history.top);
      history.top.z += sample("heightAbove") ?? 0.1;
      history.bottom.z -= sample("heightBelow") ?? 0.1;
      history.top.applyMatrix4(bone).applyMatrix4(root.current.matrixWorld);
      history.bottom.applyMatrix4(bone).applyMatrix4(root.current.matrixWorld);
      history.edges.push({ time: elapsed, above: history.top.clone(), below: history.bottom.clone() });
      if (history.edges.length > capacity) history.edges.shift();
      history.last = elapsed;
    }
    history.inverse.copy(root.current.matrixWorld).invert();
    const positions = geometry.getAttribute("position") as BufferAttribute;
    const uv = geometry.getAttribute("uv") as BufferAttribute;
    const columns = Math.max(1, ribbon.textureColumns), rows = Math.max(1, ribbon.textureRows);
    const slot = Math.max(0, Math.floor(sample("textureSlot") ?? 0)) % (columns * rows);
    history.edges.forEach((edge, index) => {
      for (const side of [0, 1]) {
        history.local.copy(side ? edge.below : edge.above);
        history.local.y -= ribbon.gravity * ((elapsed - edge.time) / 1000) ** 2 * 0.5;
        history.local.applyMatrix4(history.inverse);
        positions.setXYZ(index * 2 + side, ...history.local.toArray());
        uv.setXY(index * 2 + side, (slot % columns + index / Math.max(1, history.edges.length - 1)) / columns,
          (Math.floor(slot / columns) + side) / rows);
      }
    });
    material.color.setRGB(...(sampleOriginalCombatMeshTrack(graph, ribbon.tracks.color, frame) ?? [1, 1, 1]));
    material.opacity = Math.max(0, Math.min(1, (sample("alpha") ?? 32767) / 32767)) * opacity;
    positions.needsUpdate = true; uv.needsUpdate = true;
    geometry.setDrawRange(0, Math.max(0, history.edges.length - 1) * 6);
  });
  return <group ref={root}><mesh name="original-wow-ribbon" geometry={geometry} material={material} frustumCulled={false} /></group>;
});