import { BufferGeometry, Group, Mesh, Vector3 } from 'three';
import { STLLoader } from 'three/examples/jsm/loaders/STLLoader.js';
import { OBJLoader } from 'three/examples/jsm/loaders/OBJLoader.js';
import { mergeGeometries, mergeVertices } from 'three/examples/jsm/utils/BufferGeometryUtils.js';
import {
  buildPCGrid3D,
  CAD_HEATMAP_SHELL_OFFSET_FRAC,
  displayDistanceRange,
  displayErrorRangeSmoothHeatmap,
  distanceToRgb,
  errorToRgb,
  medianNearestNeighbor3D,
  nearestPointInGrid,
  normalizePointsLikeViewer,
  type HeatmapPoint,
} from '../shared/cadHeatmapMath.js';
import type { SampledPoint } from './pointcloudSample.js';

/**
 * 服务端平滑热力：默认「先求快、可接受稍粗网格」（与前端实时预览的 shared 常量可不同）。
 * - CAD_HEATMAP_MAX_VERTICES：焊接后顶点上限（默认 60000；要更细可设 120000）。
 * - CAD_HEATMAP_MEDIAN_NN_SAMPLES：尺度估计采样（默认 4096；要更稳可设 8192）。
 */
function parseHeatmapInt(envVal: string | undefined, fallback: number, min: number, max: number): number {
  if (envVal == null || envVal === '') return fallback;
  const n = Number(envVal);
  if (!Number.isFinite(n)) return fallback;
  return Math.min(max, Math.max(min, Math.floor(n)));
}

/** 默认偏快：约半顶点上限，merge 更早停 */
const SERVER_FAST_DEFAULT_MAX_VERTICES = 60_000;
const SERVER_FAST_DEFAULT_MEDIAN_NN = 4_096;

const SERVER_MAX_VERTICES = parseHeatmapInt(
  process.env.CAD_HEATMAP_MAX_VERTICES,
  SERVER_FAST_DEFAULT_MAX_VERTICES,
  8_000,
  500_000,
);
const SERVER_MEDIAN_NN_SAMPLES = parseHeatmapInt(
  process.env.CAD_HEATMAP_MEDIAN_NN_SAMPLES,
  SERVER_FAST_DEFAULT_MEDIAN_NN,
  512,
  32_768,
);

function normalizeCadGeometryLikeViewer(geo: BufferGeometry): void {
  geo.computeBoundingBox();
  geo.center();
  const box = geo.boundingBox;
  if (box) {
    const size = new Vector3();
    box.getSize(size);
    const maxDim = Math.max(size.x, size.y, size.z);
    if (maxDim > 0) {
      const s = 10 / maxDim;
      geo.scale(s, s, s);
    }
  }
}

function buildHeatmapPoints(sampledPoints: SampledPoint[]): HeatmapPoint[] {
  const pts: HeatmapPoint[] = sampledPoints.map((p) => ({
    x: p.x,
    y: p.y,
    z: p.z,
    error: p.error,
  }));
  normalizePointsLikeViewer(pts);
  return pts;
}

function parseObjBufferToGeometry(buffer: Buffer): BufferGeometry | null {
  const text = buffer.toString('utf8');
  if (!text.trim()) return null;
  const loader = new OBJLoader();
  const group = loader.parse(text) as Group;
  const geometries: BufferGeometry[] = [];
  group.traverse((child) => {
    if (child instanceof Mesh && child.geometry) {
      geometries.push(child.geometry);
    }
  });
  if (geometries.length === 0) return null;
  if (geometries.length === 1) return geometries[0].clone();
  try {
    const merged = mergeGeometries(geometries, false);
    if (merged) return merged;
  } catch {
    /* fall through */
  }
  return geometries[0].clone();
}

/**
 * 对已解析的 BufferGeometry（尚未 normalizeCad）计算顶点色。
 */
function computeCadShellHeatmapCore(
  geo: BufferGeometry,
  pts: HeatmapPoint[],
): { colors: Float32Array; dispMin: number; dispMax: number; vertexCount: number } | null {
  normalizeCadGeometryLikeViewer(geo);

  try {
    mergeVertices(geo);
  } catch {
    /* ignore */
  }
  geo.computeVertexNormals();

  let weldTol = 2e-4;
  while (
    geo.attributes.position &&
    geo.attributes.position.count > SERVER_MAX_VERTICES &&
    weldTol < 0.12
  ) {
    mergeVertices(geo, weldTol);
    weldTol *= 1.85;
    geo.computeVertexNormals();
  }

  const posAttr = geo.attributes.position;
  const norAttr = geo.attributes.normal;
  if (!posAttr || posAttr.count === 0) return null;

  const nV = posAttr.count;
  const pos = posAttr.array as Float32Array;
  const nor = norAttr?.array as Float32Array | undefined;

  const grid = buildPCGrid3D(pts);
  const nnMed = medianNearestNeighbor3D(pts, grid, SERVER_MEDIAN_NN_SAMPLES);
  const { minError: dispMin, maxError: dispMax } = displayErrorRangeSmoothHeatmap(pts);

  geo.computeBoundingBox();
  const bb = geo.boundingBox;
  const diag = bb ? Math.hypot(bb.max.x - bb.min.x, bb.max.y - bb.min.y, bb.max.z - bb.min.z) : 1;
  const maxBind = Math.max(nnMed * 14, diag * 0.045);
  const shellOffset = Math.max(diag * CAD_HEATMAP_SHELL_OFFSET_FRAC, nnMed * 0.35);

  const colors = new Float32Array(nV * 3);
  const noData = [0.32, 0.34, 0.38] as const;

  for (let i = 0; i < nV; i++) {
    const i3 = i * 3;
    const x = pos[i3];
    const y = pos[i3 + 1];
    const z = pos[i3 + 2];

    const hit = nearestPointInGrid(grid, x, y, z);
    if (!hit || hit.dist > maxBind) {
      colors[i3] = noData[0];
      colors[i3 + 1] = noData[1];
      colors[i3 + 2] = noData[2];
    } else {
      const err = pts[hit.idx].error ?? 0;
      const [r, g, b] = errorToRgb(err, dispMin, dispMax);
      colors[i3] = r;
      colors[i3 + 1] = g;
      colors[i3 + 2] = b;
    }

    if (nor && nor.length >= nV * 3) {
      pos[i3] = x + nor[i3] * shellOffset;
      pos[i3 + 1] = y + nor[i3 + 1] * shellOffset;
      pos[i3 + 2] = z + nor[i3 + 2] * shellOffset;
    }
  }

  posAttr.needsUpdate = true;

  return {
    colors,
    dispMin,
    dispMax,
    vertexCount: nV,
  };
}

/**
 * 在服务端对 STL 与（已采样的）测点做与前端一致的 CAD 壳热力，返回每顶点 RGB。
 */
export function computeCadShellHeatmapFromStl(
  stlBuffer: Buffer,
  sampledPoints: SampledPoint[],
): { colors: Float32Array; dispMin: number; dispMax: number; vertexCount: number } | null {
  if (sampledPoints.length < 1) return null;
  const pts = buildHeatmapPoints(sampledPoints);

  const u8 = new Uint8Array(stlBuffer);
  const ab = u8.buffer.slice(u8.byteOffset, u8.byteOffset + u8.byteLength) as ArrayBuffer;
  const loader = new STLLoader();
  const geo = loader.parse(ab);
  return computeCadShellHeatmapCore(geo, pts);
}

/** OBJ（多 mesh 会尝试合并） */
export function computeCadShellHeatmapFromObj(
  objBuffer: Buffer,
  sampledPoints: SampledPoint[],
): { colors: Float32Array; dispMin: number; dispMax: number; vertexCount: number } | null {
  if (sampledPoints.length < 1) return null;
  const pts = buildHeatmapPoints(sampledPoints);
  const geo = parseObjBufferToGeometry(objBuffer);
  if (!geo) return null;
  return computeCadShellHeatmapCore(geo, pts);
}

// ─── 偏差热力：用顶点到最近点的距离着色 ────────────────────────────────────────

function computeDeviationHeatmapCore(
  geo: BufferGeometry,
  pts: HeatmapPoint[],
): { colors: Float32Array; dispMin: number; dispMax: number; vertexCount: number } | null {
  normalizeCadGeometryLikeViewer(geo);

  // 不做 mergeVertices —— 保持与客户端 Three.js Loader 一致的非索引顶点布局，
  // 使输出的颜色数组长度 == 客户端 rawCount，触发快速路径而非主线程 merge 卡死。
  geo.computeVertexNormals();

  const posAttr = geo.attributes.position;
  const norAttr = geo.attributes.normal;
  if (!posAttr || posAttr.count === 0) return null;

  const nV = posAttr.count;
  const pos = posAttr.array as Float32Array;
  const nor = norAttr?.array as Float32Array | undefined;

  const grid = buildPCGrid3D(pts);
  const nnMed = medianNearestNeighbor3D(pts, grid, SERVER_MEDIAN_NN_SAMPLES);

  geo.computeBoundingBox();
  const bb = geo.boundingBox;
  const diag = bb ? Math.hypot(bb.max.x - bb.min.x, bb.max.y - bb.min.y, bb.max.z - bb.min.z) : 1;
  const maxBind = Math.max(nnMed * 14, diag * 0.045);
  const shellOffset = Math.max(diag * CAD_HEATMAP_SHELL_OFFSET_FRAC, nnMed * 0.35);

  const distances: number[] = [];
  const vertexDists = new Float64Array(nV);

  for (let i = 0; i < nV; i++) {
    const i3 = i * 3;
    const hit = nearestPointInGrid(grid, pos[i3], pos[i3 + 1], pos[i3 + 2]);
    if (hit && hit.dist <= maxBind) {
      vertexDists[i] = hit.dist;
      distances.push(hit.dist);
    } else {
      vertexDists[i] = -1;
    }
  }

  const { minDist: dispMin, maxDist: dispMax } = displayDistanceRange(distances);

  const colors = new Float32Array(nV * 3);
  const noData = [0.32, 0.34, 0.38] as const;

  for (let i = 0; i < nV; i++) {
    const i3 = i * 3;
    const d = vertexDists[i];
    if (d < 0) {
      colors[i3] = noData[0];
      colors[i3 + 1] = noData[1];
      colors[i3 + 2] = noData[2];
    } else {
      const [r, g, b] = distanceToRgb(d, dispMin, dispMax);
      colors[i3] = r;
      colors[i3 + 1] = g;
      colors[i3 + 2] = b;
    }

    if (nor && nor.length >= nV * 3) {
      pos[i3] = pos[i3] + nor[i3] * shellOffset;
      pos[i3 + 1] = pos[i3 + 1] + nor[i3 + 1] * shellOffset;
      pos[i3 + 2] = pos[i3 + 2] + nor[i3 + 2] * shellOffset;
    }
  }

  posAttr.needsUpdate = true;

  return { colors, dispMin, dispMax, vertexCount: nV };
}

export function computeDeviationHeatmapFromStl(
  stlBuffer: Buffer,
  sampledPoints: SampledPoint[],
): { colors: Float32Array; dispMin: number; dispMax: number; vertexCount: number } | null {
  if (sampledPoints.length < 1) return null;
  const pts = buildHeatmapPoints(sampledPoints);
  const u8 = new Uint8Array(stlBuffer);
  const ab = u8.buffer.slice(u8.byteOffset, u8.byteOffset + u8.byteLength) as ArrayBuffer;
  const loader = new STLLoader();
  const geo = loader.parse(ab);
  return computeDeviationHeatmapCore(geo, pts);
}

export function computeDeviationHeatmapFromObj(
  objBuffer: Buffer,
  sampledPoints: SampledPoint[],
): { colors: Float32Array; dispMin: number; dispMax: number; vertexCount: number } | null {
  if (sampledPoints.length < 1) return null;
  const pts = buildHeatmapPoints(sampledPoints);
  const geo = parseObjBufferToGeometry(objBuffer);
  if (!geo) return null;
  return computeDeviationHeatmapCore(geo, pts);
}