import { BufferGeometry, Float32BufferAttribute, Vector3 } from 'three';
import { STLLoader } from 'three/examples/jsm/loaders/STLLoader.js';
import { MeshBVH } from 'three-mesh-bvh';
import { buildPCGrid3D, nearestPointInGrid, type HeatmapPoint } from '../shared/cadHeatmapMath.js';
import type { SampledPoint } from './pointcloudSample.js';
import type { StepManifestPart } from './stepParse.js';

export type PartMeshBVH = {
  part: StepManifestPart;
  geometry: BufferGeometry;
  bvh: MeshBVH;
};

export function loadStlToBVH(buf: Buffer): { geometry: BufferGeometry; bvh: MeshBVH } {
  const loader = new STLLoader();
  const u8 = new Uint8Array(buf);
  const ab = u8.buffer.slice(u8.byteOffset, u8.byteOffset + u8.byteLength) as ArrayBuffer;
  const geometry = loader.parse(ab) as BufferGeometry;
  geometry.computeVertexNormals();
  const bvh = new MeshBVH(geometry);
  return { geometry, bvh };
}

function distancePointToMesh(bvh: MeshBVH, p: Vector3): number {
  const hit = bvh.closestPointToPoint(p);
  return hit ? hit.distance : Infinity;
}

/** 零件 AABB 的外接球（心 + 半对角线）；点到网格距离 ≥ max(0, d(点,心)−r) */
export type PartSphere = { cx: number; cy: number; cz: number; r: number };

export function computePartSphere(geometry: BufferGeometry): PartSphere {
  if (!geometry.boundingBox) geometry.computeBoundingBox();
  const bb = geometry.boundingBox!;
  const cx = (bb.min.x + bb.max.x) * 0.5;
  const cy = (bb.min.y + bb.max.y) * 0.5;
  const cz = (bb.min.z + bb.max.z) * 0.5;
  const dx = bb.max.x - bb.min.x;
  const dy = bb.max.y - bb.min.y;
  const dz = bb.max.z - bb.min.z;
  const r = 0.5 * Math.sqrt(dx * dx + dy * dy + dz * dz) + 1e-9;
  return { cx, cy, cz, r };
}

function sphereDistLowerBound(px: number, py: number, pz: number, s: PartSphere): number {
  const dx = px - s.cx;
  const dy = py - s.cy;
  const dz = pz - s.cz;
  const d = Math.sqrt(dx * dx + dy * dy + dz * dz);
  return Math.max(0, d - s.r);
}

export type PartStats = {
  partId: string;
  name: string;
  min: number;
  max: number;
  mean: number;
  stdDev: number;
  passRate: number;
  pointCount: number;
};

const ASSIGN_PROGRESS_EVERY = 12_000;

export type AssignNearestOptions = {
  /** 每处理若干点上报一次，避免长时间卡在 0.55 无反馈 */
  onProgress?: (done: number, total: number) => void | Promise<void>;
};

/**
 * 对每个测点找最近零件网格距离：朴素 O(点数×零件数) BVH 在零件多时极慢。
 * 用外接球下界排序 + 下界≥当前最优时跳过 BVH，大幅减少查询次数。
 */
export async function assignNearestPartAndStats(
  parts: PartMeshBVH[],
  points: SampledPoint[],
  toleranceMm: number,
  options: AssignNearestOptions = {},
): Promise<{
  assignments: number[];
  perPartDists: number[][];
  pointDistances: number[];
  stats: PartStats[];
}> {
  const nP = parts.length;
  const nPt = points.length;
  const assignments = new Array<number>(nPt);
  const pointDistances = new Array<number>(nPt);
  const perPartDists: number[][] = Array.from({ length: nP }, () => []);

  const spheres: PartSphere[] = parts.map((p) => computePartSphere(p.geometry));
  const lbs = new Float64Array(nP);
  const order: number[] = new Array(nP);

  const tmp = new Vector3();
  for (let pi = 0; pi < nPt; pi++) {
    const x = points[pi].x;
    const y = points[pi].y;
    const z = points[pi].z;
    tmp.set(x, y, z);

    for (let j = 0; j < nP; j++) {
      lbs[j] = sphereDistLowerBound(x, y, z, spheres[j]!);
    }
    for (let j = 0; j < nP; j++) order[j] = j;
    order.sort((a, b) => lbs[a]! - lbs[b]!);

    let best = Infinity;
    let bestIdx = 0;
    for (let k = 0; k < nP; k++) {
      const j = order[k]!;
      const lb = lbs[j]!;
      if (lb >= best) break;
      const d = distancePointToMesh(parts[j].bvh, tmp);
      if (d < best) {
        best = d;
        bestIdx = j;
      }
    }

    assignments[pi] = bestIdx;
    pointDistances[pi] = best;
    perPartDists[bestIdx].push(best);

    if (options.onProgress && (pi + 1) % ASSIGN_PROGRESS_EVERY === 0) {
      await options.onProgress(pi + 1, nPt);
    }
  }

  if (options.onProgress && nPt > 0) {
    await options.onProgress(nPt, nPt);
  }

  const stats = computePartStatsFromPerPartDists(
    parts.map((p) => p.part),
    perPartDists,
    toleranceMm,
  );

  return { assignments, perPartDists, pointDistances, stats };
}

/** 由 assignments + distances 填充分桶后算统计（GPU 路径复用） */
export function perPartDistsFromAssignments(
  nParts: number,
  assignments: number[],
  pointDistances: number[],
): number[][] {
  const perPartDists: number[][] = Array.from({ length: nParts }, () => []);
  for (let i = 0; i < assignments.length; i++) {
    perPartDists[assignments[i]!]!.push(pointDistances[i]!);
  }
  return perPartDists;
}

export function computePartStatsFromPerPartDists(
  manifestParts: StepManifestPart[],
  perPartDists: number[][],
  toleranceMm: number,
): PartStats[] {
  return manifestParts.map((part, j) => {
    const ds = perPartDists[j];
    const pointCount = ds?.length ?? 0;
    if (pointCount === 0) {
      return {
        partId: part.id,
        name: part.name,
        min: 0,
        max: 0,
        mean: 0,
        stdDev: 0,
        passRate: 1,
        pointCount: 0,
      };
    }
    let min = Infinity;
    let max = -Infinity;
    let sum = 0;
    for (const d of ds) {
      if (d < min) min = d;
      if (d > max) max = d;
      sum += d;
    }
    const mean = sum / pointCount;
    let vsum = 0;
    for (const d of ds) {
      const t = d - mean;
      vsum += t * t;
    }
    const stdDev = Math.sqrt(vsum / pointCount);
    let pass = 0;
    for (const d of ds) {
      if (d < toleranceMm) pass++;
    }
    const passRate = pass / pointCount;
    return {
      partId: part.id,
      name: part.name,
      min,
      max,
      mean,
      stdDev,
      passRate,
      pointCount,
    };
  });
}

/** 每顶点：到「归属该零件的测点」的最近距离，用于简易热力（可选） */
export function computeVertexDistanceHeatmap(
  geometry: BufferGeometry,
  pointsForPart: SampledPoint[],
): Float32Array | null {
  const pos = geometry.attributes.position;
  if (!pos || pos.count === 0) return null;
  const nV = pos.count;
  const hp: HeatmapPoint[] = pointsForPart.map((p) => ({ x: p.x, y: p.y, z: p.z, error: 0 }));
  if (hp.length === 0) {
    return new Float32Array(nV).fill(0);
  }
  const grid = buildPCGrid3D(hp);
  const arr = pos.array as Float32Array;
  const out = new Float32Array(nV);
  for (let i = 0; i < nV; i++) {
    const x = arr[i * 3];
    const y = arr[i * 3 + 1];
    const z = arr[i * 3 + 2];
    const nn = nearestPointInGrid(grid, x, y, z);
    out[i] = nn ? nn.dist : 0;
  }
  return out;
}

export function float32ToVertexColorRgb(
  values: Float32Array,
  globalMin: number,
  globalMax: number,
): Float32Array {
  const n = values.length;
  const rgb = new Float32Array(n * 3);
  const lo = globalMin;
  const hi = Math.max(lo + 1e-9, globalMax);
  for (let i = 0; i < n; i++) {
    const t = (values[i] - lo) / (hi - lo);
    const g = Math.max(0, Math.min(1, 1 - t));
    const r = Math.max(0, Math.min(1, t));
    rgb[i * 3] = r;
    rgb[i * 3 + 1] = g;
    rgb[i * 3 + 2] = 0.15;
  }
  return rgb;
}