import { Suspense, useLayoutEffect } from 'react';
import { useLoader } from '@react-three/fiber';
import { useTexture } from '@react-three/drei';
import type { TTSObject } from '@tts/shared';
import * as THREE from 'three';
import type { Object3D } from 'three';
import { assetUrl } from '@tts/http';
import { FlexibleModelLoader } from './flexibleModelLoader';
import { objectTint, tintedColor } from './sharedResources';

/**
 * The mesh content for a custom model, exported so the full-setup view can
 * compose it into a shared scene. Renders the model from `CustomMesh.MeshURL`
 * (or a neutral box placeholder when absent).
 */
export function CustomModelMesh({ object }: { object: TTSObject }) {
  const meshUrl = object.CustomMesh?.MeshURL;
  const tint = objectTint(object);
  if (!meshUrl) {
    return (
      <mesh>
        <boxGeometry args={[1, 1, 1]} />
        <meshStandardMaterial color={tintedColor(new THREE.Color('#52525b'), tint)} />
      </mesh>
    );
  }

  return (
    <Suspense fallback={null}>
      <Model
        meshUrl={meshUrl}
        diffuseUrl={object.CustomMesh?.DiffuseURL}
        tint={tint}
      />
    </Suspense>
  );
}

function Model({
  meshUrl,
  diffuseUrl,
  tint,
}: {
  meshUrl: string;
  diffuseUrl?: string;
  tint: THREE.Color;
}) {
  const root = useLoader(FlexibleModelLoader, assetUrl(meshUrl));

  return (
    <>
      <primitive object={root} scale={0.5} />
      {diffuseUrl && <DiffuseTexture root={root} url={diffuseUrl} />}
      <Tint root={root} tint={tint} />
    </>
  );
}

// Apply the object's tint to every material on the loaded model, multiplying
// the existing color. Rendered after the model so it runs once the materials
// exist.
function Tint({ root, tint }: { root: Object3D; tint: THREE.Color }) {
  useLayoutEffect(() => {
    root.traverse((child) => {
      const mesh = child as THREE.Mesh;
      if (mesh.isMesh) {
        const material = Array.isArray(mesh.material)
          ? mesh.material[0]
          : mesh.material;
        if (material && 'color' in material) {
          (material as THREE.MeshStandardMaterial).color.multiply(tint);
          material.needsUpdate = true;
        }
      }
    });
  }, [root, tint]);

  return null;
}

// Rendered inside the Canvas so `useTexture` can access the R3F store. Only
// mounted when a diffuse URL exists, so the hook count stays consistent.
function DiffuseTexture({ root, url }: { root: Object3D; url: string }) {
  const texture = useTexture(assetUrl(url));

  // Diffuse art is sRGB-encoded; `TextureLoader` leaves `colorSpace` as
  // `NoColorSpace`, which double-decodes it in the shader and washes out
  // contrast. Mark it sRGB so the GPU decodes it once, correctly.
  texture.colorSpace = THREE.SRGBColorSpace;

  // Apply the diffuse texture to every mesh material on the loaded model.
  useLayoutEffect(() => {
    root.traverse((child) => {
      const mesh = child as THREE.Mesh;
      if (mesh.isMesh) {
        const material = Array.isArray(mesh.material)
          ? mesh.material[0]
          : mesh.material;
        if (material && 'map' in material) {
          material.map = texture;
          material.needsUpdate = true;
        }
      }
    });
  }, [root, texture]);

  return null;
}