import { Color, Vector3 } from 'three';

export const APPEARANCE_REGIONS = Object.freeze(['body', 'hair', 'iris']);
const editorStates = new WeakMap();
const colorStates = new WeakMap();

export function getAppearanceRegion(material) {
  const region = material?.userData?.falukantAppearance?.region;
  if (APPEARANCE_REGIONS.includes(region)) return region;
  const name = String(material?.name || '').toLowerCase();
  if (['hair', 'falukant_hair'].includes(name)) return 'hair';
  if (['iris', 'falukant_iris'].includes(name)) return 'iris';
  return 'body';
}

function setupColor(material) {
  if (colorStates.has(material)) return colorStates.get(material);
  const uniforms = { target: { value: new Color() }, enabled: { value: 0 } };
  const previousCompile = material.onBeforeCompile;
  material.onBeforeCompile = function (shader, renderer) {
    previousCompile.call(this, shader, renderer);
    shader.uniforms.falukantTargetColor = uniforms.target;
    shader.uniforms.falukantColorEnabled = uniforms.enabled;
    shader.fragmentShader = `uniform vec3 falukantTargetColor;
uniform float falukantColorEnabled;\n${shader.fragmentShader}`;
    shader.fragmentShader = shader.fragmentShader.replace('#include <map_fragment>', `
#include <map_fragment>
float falukantDetail = clamp(dot(diffuseColor.rgb, vec3(0.2126, 0.7152, 0.0722)) * 1.8 + 0.18, 0.18, 1.0);
diffuseColor.rgb = mix(diffuseColor.rgb, falukantTargetColor * falukantDetail, falukantColorEnabled);
`);
    // Some Meshy models also bake the original color into an emissive map.
    shader.fragmentShader = shader.fragmentShader.replace('#include <emissivemap_fragment>', `
#include <emissivemap_fragment>
float falukantEmission = dot(totalEmissiveRadiance, vec3(0.2126, 0.7152, 0.0722));
totalEmissiveRadiance = mix(totalEmissiveRadiance, falukantTargetColor * falukantEmission, falukantColorEnabled);
`);
  };
  material.customProgramCacheKey = () => 'falukant-appearance-v1';
  material.needsUpdate = true;
  colorStates.set(material, uniforms);
  return uniforms;
}

export function applyCharacterAppearance(model, appearance = {}) {
  model.traverse(object => {
    if (!object.isMesh) return;
    for (const material of [].concat(object.material || [])) {
      const region = getAppearanceRegion(material);
      if (region === 'body') continue;
      const color = region === 'hair' ? appearance.hairColor : appearance.eyeColor;
      const state = setupColor(material);
      const valid = typeof color === 'string' && /^#[0-9a-f]{6}$/i.test(color);
      state.enabled.value = valid ? 1 : 0;
      if (valid) state.target.value.set(color);
    }
  });
}

// Keep texture references, UVs, normals and skinning attributes. Only split
// triangle indices into standard glTF material groups. No texture guessing.
export function prepareAppearanceEditor(model) {
  model.traverse(mesh => {
    if (!mesh.isMesh || editorStates.has(mesh)) return;
    const original = mesh.geometry;
    if (!original.getAttribute('position')) return;
    mesh.geometry = original.clone();
    const geometry = mesh.geometry;
    if (!geometry.index) {
      geometry.setIndex(Array.from({ length: geometry.getAttribute('position').count }, (_, i) => i));
    }
    const oldMaterials = [].concat(mesh.material);
    const sources = new Map();
    const materialSources = [];
    const materials = [];
    oldMaterials.forEach((material, oldIndex) => {
      const key = material.userData?.falukantAppearance?.source ?? `material-${oldIndex}`;
      if (!sources.has(key)) {
        sources.set(key, sources.size);
        for (const region of APPEARANCE_REGIONS) {
          const clone = material.clone();
          clone.name = `Falukant_${region}_${sources.get(key)}`;
          clone.userData.falukantAppearance = { version: 1, region, source: key };
          materials.push(clone);
        }
      }
      materialSources[oldIndex] = sources.get(key);
    });
    const count = geometry.index.count / 3;
    const regions = new Uint8Array(count);
    const bases = new Uint16Array(count);
    const groups = geometry.groups.length ? geometry.groups : [{ start: 0, count: geometry.index.count, materialIndex: 0 }];
    for (const group of groups) {
      const materialIndex = group.materialIndex || 0;
      const region = APPEARANCE_REGIONS.indexOf(getAppearanceRegion(oldMaterials[materialIndex]));
      for (let face = group.start / 3; face < (group.start + group.count) / 3; face++) {
        regions[face] = region;
        bases[face] = materialSources[materialIndex];
      }
    }
    mesh.material = materials;
    const state = { regions, bases };
    editorStates.set(mesh, state);
    rebuildGroups(mesh, state);
  });
  model.userData.falukantAppearance = { version: 1, regions: APPEARANCE_REGIONS.slice() };
}

function rebuildGroups(mesh, state) {
  const geometry = mesh.geometry;
  const oldIndex = geometry.index.array;
  const buckets = Array.from({ length: mesh.material.length }, () => []);
  for (let face = 0; face < state.regions.length; face++) {
    buckets[state.bases[face] * 3 + state.regions[face]].push(face);
  }
  const index = new oldIndex.constructor(oldIndex.length);
  const regions = new Uint8Array(state.regions.length);
  const bases = new Uint16Array(state.bases.length);
  let offset = 0;
  geometry.clearGroups();
  buckets.forEach((faces, materialIndex) => {
    if (!faces.length) return;
    geometry.addGroup(offset * 3, faces.length * 3, materialIndex);
    for (const face of faces) {
      index.set(oldIndex.subarray(face * 3, face * 3 + 3), offset * 3);
      regions[offset] = state.regions[face];
      bases[offset] = state.bases[face];
      offset++;
    }
  });
  state.regions = regions;
  state.bases = bases;
  // Reuse the GPU index buffer while painting instead of leaking a new one
  // for each brush movement.
  geometry.index.array.set(index);
  geometry.index.needsUpdate = true;
}

export function paintAppearanceRegion(hit, worldRadius, region, cameraDirection) {
  const mesh = hit.object;
  const state = editorStates.get(mesh);
  const label = APPEARANCE_REGIONS.indexOf(region);
  if (!state || label < 0 || !Number.isInteger(hit.faceIndex)) return false;
  const point = mesh.worldToLocal(hit.point.clone());
  const scale = mesh.getWorldScale(new Vector3());
  const radius = worldRadius / Math.max(Math.abs(scale.x), Math.abs(scale.y), Math.abs(scale.z));
  const direction = cameraDirection.clone().transformDirection(mesh.matrixWorld.clone().invert());
  const positions = mesh.geometry.getAttribute('position');
  const index = mesh.geometry.index.array;
  const a = new Vector3(), b = new Vector3(), c = new Vector3();
  const normal = new Vector3(), edge = new Vector3(), center = new Vector3();
  let changed = false;
  for (let face = 0; face < state.regions.length; face++) {
    if (state.regions[face] === label) continue;
    a.fromBufferAttribute(positions, index[face * 3]);
    b.fromBufferAttribute(positions, index[face * 3 + 1]);
    c.fromBufferAttribute(positions, index[face * 3 + 2]);
    center.copy(a).add(b).add(c).multiplyScalar(1 / 3);
    normal.subVectors(b, a).cross(edge.subVectors(c, a));
    if (face === hit.faceIndex || (center.distanceToSquared(point) <= radius * radius && normal.dot(direction) < 0)) {
      state.regions[face] = label;
      changed = true;
    }
  }
  if (changed) rebuildGroups(mesh, state);
  return changed;
}

export function captureAppearanceEditor(model) {
  const snapshot = [];
  model.traverse(mesh => {
    const state = editorStates.get(mesh);
    if (state) snapshot.push({ mesh, index: mesh.geometry.index.clone(), regions: state.regions.slice(), bases: state.bases.slice() });
  });
  return snapshot;
}

export function restoreAppearanceEditor(snapshot) {
  for (const { mesh, index, regions, bases } of snapshot) {
    mesh.geometry.index.array.set(index.array);
    mesh.geometry.index.needsUpdate = true;
    const state = { regions: regions.slice(), bases: bases.slice() };
    editorStates.set(mesh, state);
    rebuildGroups(mesh, state);
  }
}

export function getAppearanceStats(model) {
  const stats = { hair: 0, iris: 0, body: 0 };
  model.traverse(mesh => {
    if (!mesh.isMesh) return;
    const materials = [].concat(mesh.material);
    for (const group of mesh.geometry.groups) {
      stats[getAppearanceRegion(materials[group.materialIndex || 0])] += group.count / 3;
    }
  });
  return stats;
}

export function disposeCharacterModel(model) {
  const geometries = new Set(), materials = new Set(), textures = new Set(), images = new Set();
  model.traverse(object => {
    if (object.geometry) geometries.add(object.geometry);
    for (const material of [].concat(object.material || [])) materials.add(material);
  });
  for (const material of materials) {
    for (const value of Object.values(material)) {
      if (value?.isTexture) {
        textures.add(value);
        if (value.source?.data) images.add(value.source.data);
      }
    }
    material.dispose();
  }
  geometries.forEach(geometry => geometry.dispose());
  textures.forEach(texture => texture.dispose());
  images.forEach(image => image.close?.());
}