import * as THREE from "three";
import { GLTFLoader } from "three/examples/jsm/loaders/GLTFLoader.js";
import { parseMachineUrdf, type MachineStructure, type Vec3 } from "./machine-urdf";
import type { ImportedMachine } from "./imported-machine";

export type MachinePartRole = "base" | "x" | "y" | "z" | "tool" | "tilt" | "c" | "work";
export type MachinePartSource = { id: string; filename: string; hash: string; scene: THREE.Group };
export type MachineNodeTransform = { translationMm: Vec3; rotationDeg: Vec3 };
export type MachinePartDraft = MachinePartSource & {
  role: MachinePartRole | "";
  parentId: string | null;
  visible: boolean;
  showOrigin: boolean;
  originMm: Vec3;
  translationMm: Vec3;
  rotationDeg: Vec3;
  nodeParents: Record<string, string | null>;
  nodeRoles: Record<string, MachinePartRole | "">;
  nodeTransforms: Record<string, MachineNodeTransform>;
};
export type MachineTreeNode = { path: string; name: string; type: "group" | "mesh"; originMm: Vec3; worldOriginMm: Vec3; rotationDeg: Vec3; children: MachineTreeNode[] };
export type PhysicalZero = { x: number; y: number; z: number; a: number; b: number; c: number };
export type ManagedMachineGeometry = {
  schema: "linuxcnc.machine-manager/1";
  coordinates: "XYZAC" | "XYZBC";
  cOriginMm: Vec3;
  tiltOffsetMm: Vec3;
  physicalZero: PhysicalZero;
  parts: Array<{ id: string; filename: string; hash: string; role: MachinePartRole; parentId: string | null; originMm: Vec3; translationMm: Vec3; rotationDeg: Vec3; visible: boolean; nodeParents: Record<string, string | null>; nodeRoles: Record<string, MachinePartRole | "">; nodeTransforms: Record<string, MachineNodeTransform> }>;
};

export const MACHINE_PART_ROLES: MachinePartRole[] = ["base", "x", "y", "z", "tool", "tilt", "c", "work"];
export const roleLabel = (role: MachinePartRole, coordinates: string) => ({ base: "床身底座", x: "X 轴部件", y: "Y 轴部件", z: "Z 轴部件", tool: "主轴/刀具", tilt: `${coordinates === "XYZBC" ? "B" : "A"} 轴部件`, c: "C 轴部件", work: "工件/工作台" }[role]);
export const expectedParentRole = (role: MachinePartRole): MachinePartRole | null => ({ base: null, x: "base", y: "x", z: "y", tool: "z", tilt: "base", c: "tilt", work: "c" }[role] as MachinePartRole | null);


const finite3 = (value: Vec3, label: string) => {
  if (value.length !== 3 || value.some((number) => !Number.isFinite(number))) throw new Error(`${label}需要三个有限数值`);
};
const cloneMaterial = (material: THREE.Material) => {
  const cloned = material.clone();
  for (const [key, value] of Object.entries(cloned)) if (value instanceof THREE.Texture) (cloned as unknown as Record<string, unknown>)[key] = value.clone();
  return cloned;
};
export function cloneMachinePartScene(source: THREE.Object3D) {
  const clone = source.clone(true);
  clone.traverse((node) => {
    const mesh = node as THREE.Mesh;
    if (mesh.geometry) mesh.geometry = mesh.geometry.clone();
    if (mesh.material) mesh.material = Array.isArray(mesh.material) ? mesh.material.map(cloneMaterial) : cloneMaterial(mesh.material);
  });
  return clone;
}

function scenePathMap(scene: THREE.Object3D) {
  const nodes = new Map<string, THREE.Object3D>(), paths = new WeakMap<THREE.Object3D, string>();
  const visit = (node: THREE.Object3D, path: string) => { nodes.set(path, node); paths.set(node, path); node.children.forEach((child, index) => visit(child, `${path}/${index}`)); };
  visit(scene, "scene"); return { nodes, paths };
}

export function applyMachineNodeParents(scene: THREE.Object3D, overrides: Record<string, string | null>) {
  const { nodes } = scenePathMap(scene);
  for (const [sourcePath, targetPath] of Object.entries(overrides ?? {})) {
    const source = nodes.get(sourcePath), target = targetPath === null ? scene : nodes.get(targetPath);
    if (!source || source === scene || !target || source === target) throw new Error("GLB 层级拖放目标无效");
    let cursor: THREE.Object3D | null = target;
    while (cursor) { if (cursor === source) throw new Error("GLB 层级不能形成循环"); cursor = cursor.parent; }
    scene.updateMatrixWorld(true); target.attach(source); scene.updateMatrixWorld(true);
  }
  return scene;
}

// Index source paths before reparenting so edits and axis bindings keep their identity.
export function configureMachinePartScene(scene: THREE.Object3D, part: Pick<MachinePartDraft, "nodeParents" | "nodeTransforms">) {
  const { nodes, paths } = scenePathMap(scene);
  applyMachineNodeParents(scene, part.nodeParents ?? {});
  for (const [path, transform] of Object.entries(part.nodeTransforms ?? {})) {
    const node = nodes.get(path);
    if (!node || node === scene) throw new Error("部件变换节点无效");
    finite3(transform.translationMm, "节点相对父级位置");
    finite3(transform.rotationDeg, "节点相对父级姿态");
    node.position.fromArray(transform.translationMm).multiplyScalar(.001);
    node.rotation.set(...transform.rotationDeg.map(THREE.MathUtils.degToRad) as Vec3, "XYZ");
    node.updateMatrix();
  }
  scene.updateMatrixWorld(true);
  return { scene, nodes, paths };
}

export function reparentMachineNode(part: MachinePartDraft, sourcePath: string, targetPath: string) {
  const { scene, nodes } = configureMachinePartScene(part.scene.clone(true), part);
  const source = nodes.get(sourcePath), target = nodes.get(targetPath);
  if (!source || source === scene || !target || source === target) throw new Error("GLB 层级拖放目标无效");
  for (let cursor: THREE.Object3D | null = target; cursor; cursor = cursor.parent) {
    if (cursor === source) throw new Error("GLB 层级不能形成循环");
  }
  target.attach(source);
  // Capture edited local poses after attach: changing parent preserves the current
  // assembly pose, and subsequent numeric edits use the new parent's frame.
  const nodeTransforms = Object.fromEntries([...nodes].filter(([path]) => path !== "scene").map(([path, node]) => [path, {
    translationMm: node.position.toArray().map((value) => value * 1000) as Vec3,
    rotationDeg: new THREE.Euler().setFromQuaternion(node.quaternion, "XYZ").toArray().slice(0, 3).map((value) => THREE.MathUtils.radToDeg(Number(value))) as Vec3,
  }]));
  return { nodeParents: { ...part.nodeParents, [sourcePath]: targetPath === "scene" ? null : targetPath }, nodeTransforms };
}

export function machineSceneNode(scene: THREE.Object3D, path: string) {
  return scenePathMap(scene).nodes.get(path) ?? null;
}

export function machinePartTree(part: Pick<MachinePartDraft, "filename" | "scene" | "nodeParents" | "nodeTransforms">): MachineTreeNode {
  const { scene, paths } = configureMachinePartScene(part.scene.clone(true), part);
  scene.updateMatrixWorld(true);
  const build = (node: THREE.Object3D): MachineTreeNode => ({
    path: paths.get(node)!,
    name: node.name || (paths.get(node) === "scene" ? part.filename : `${(node as THREE.Mesh).isMesh ? "Mesh" : "Node"} ${paths.get(node)!.split("/").at(-1)}`),
    type: (node as THREE.Mesh).isMesh ? "mesh" : "group",
    originMm: node.position.toArray().map((value) => value * 1000) as Vec3,
    worldOriginMm: node.getWorldPosition(new THREE.Vector3()).toArray().map((value) => value * 1000) as Vec3,
    rotationDeg: new THREE.Euler().setFromQuaternion(node.quaternion, "XYZ").toArray().slice(0, 3).map((value) => THREE.MathUtils.radToDeg(Number(value))) as Vec3,
    children: node.children.map(build),
  });
  return build(scene);
}

export async function loadMachinePart(bytes: ArrayBuffer, filename: string): Promise<MachinePartSource> {
  if (!bytes.byteLength) throw new Error("GLB 部件为空");
  const manager = new THREE.LoadingManager();
  manager.setURLModifier((url) => { if (!url.startsWith("blob:")) throw new Error("GLB 部件不能加载外部资源"); return url; });
  const gltf = await new GLTFLoader(manager).parseAsync(bytes, "");
  const hash = [...new Uint8Array(await crypto.subtle.digest("SHA-256", bytes))].map((number) => number.toString(16).padStart(2, "0")).join("");
  gltf.scene.traverse((node) => { if ((node as THREE.Mesh).isMesh) { node.castShadow = true; node.receiveShadow = true; } });
  return { id: crypto.randomUUID(), filename, hash, scene: gltf.scene };
}

export function validateMachineAssembly(coordinates: "XYZAC" | "XYZBC", parts: MachinePartDraft[]) {
  if (parts.length !== MACHINE_PART_ROLES.length) throw new Error("应用前需要正好八个机床角色部件");
  const byRole = new Map<MachinePartRole, MachinePartDraft>();
  const ids = new Set(parts.map((part) => part.id));
  if (ids.size !== parts.length) throw new Error("部件 ID 重复");
  for (const part of parts) {
    if (!part.role) throw new Error(`${part.filename} 尚未指定机床角色`);
    if (byRole.has(part.role)) throw new Error(`${roleLabel(part.role, coordinates)}角色重复`);
    finite3(part.originMm, `${part.filename}局部原点`); finite3(part.translationMm, `${part.filename}平移`); finite3(part.rotationDeg, `${part.filename}旋转`);
    for (const [path, transform] of Object.entries(part.nodeTransforms ?? {})) { if (path === "scene") throw new Error(`${part.filename} 的 scene 节点不能重复定义部件变换`); finite3(transform.translationMm, `${part.filename}/${path} 相对父级位置`); finite3(transform.rotationDeg, `${part.filename}/${path} 相对父级姿态`); }
    byRole.set(part.role, part);
  }
  for (const role of MACHINE_PART_ROLES) if (!byRole.has(role)) throw new Error(`缺少${roleLabel(role, coordinates)}`);
  for (const [role, part] of byRole) {
    const parentRole = expectedParentRole(role);
    if (parentRole === null) { if (part.parentId !== null) throw new Error("床身底座不能设置父部件"); continue; }
    const expected = byRole.get(parentRole)!;
    if (part.parentId !== expected.id) throw new Error(`${roleLabel(role, coordinates)}必须挂到${roleLabel(parentRole, coordinates)}`);
  }
  return byRole;
}

function machineUrdf(coordinates: "XYZAC" | "XYZBC", cOriginMm: Vec3, tiltOffsetMm: Vec3) {
  finite3(cOriginMm, "C 轴原点"); finite3(tiltOffsetMm, "A/B 轴偏置");
  if (coordinates === "XYZAC" && Math.abs(tiltOffsetMm[0]) > 1e-9) throw new Error("AC 原生 TRT 不支持沿 A 轴的 X 偏置");
  if (coordinates === "XYZBC" && Math.abs(tiltOffsetMm[1]) > 1e-9) throw new Error("BC 原生 TRT 不支持沿 B 轴的 Y 偏置");
  const metres = (values: Vec3) => values.map((value) => value / 1000) as Vec3;
  const p = metres(cOriginMm), d = metres(tiltOffsetMm);
  const t = p.map((value, index) => value + d[index]) as Vec3;
  const u = d.map((value) => -value) as Vec3;
  const work = p.map((value) => -value) as Vec3;
  const tilt = coordinates === "XYZAC" ? "A" : "B";
  const tiltAxis = coordinates === "XYZAC" ? "1 0 0" : "0 1 0";
  const radians = (degrees: number) => degrees * Math.PI / 180;
  const joint = (name: string, type: string, parent: string, child: string, origin: Vec3, axis = "", coordinate = "", limits: number[] = []) => `  <joint name="${name}" type="${type}">\n    <parent link="${parent}"/><child link="${child}"/>\n    <origin xyz="${origin.join(" ")}" rpy="0 0 0"/>${axis ? `\n    <axis xyz="${axis}"/><cnc:axis coordinate="${coordinate}"/>\n    <limit lower="${limits[0]}" upper="${limits[1]}" velocity="${limits[2]}" effort="1000"/>` : ""}\n  </joint>`;
  return `<?xml version="1.0"?>\n<robot name="managed_${coordinates.toLowerCase()}" xmlns:cnc="urn:linuxcnc:web:machine:1">\n  <cnc:machine profile="trt-xyz-head/1" base="base" tool="tool" work="work"/>\n${["base", "x", "y", "z", "tool", "tilt", "c", "work"].map((name) => `  <link name="${name}"/>`).join("\n")}\n${[
    joint("x_joint", "prismatic", "base", "x", [0, 0, 0], "1 0 0", "X", [-.2, .2, .02]),
    joint("y_joint", "prismatic", "x", "y", [0, 0, 0], "0 1 0", "Y", [-.1, .1, .02]),
    joint("z_joint", "prismatic", "y", "z", [0, 0, 0], "0 0 1", "Z", [-.12, .12, .02]),
    joint("tool_mount", "fixed", "z", "tool", [0, 0, 0]),
    joint("tilt_joint", "revolute", "base", "tilt", t, tiltAxis, tilt, [radians(coordinates === "XYZAC" ? -100 : -36000), radians(coordinates === "XYZAC" ? 50 : 36000), radians(30)]),
    joint("c_joint", "revolute", "tilt", "c", u, "0 0 1", "C", [radians(-36000), radians(36000), radians(30)]),
    joint("work_mount", "fixed", "c", "work", work),
  ].join("\n")}\n</robot>\n`;
}

export async function buildManagedMachine(input: { coordinates: "XYZAC" | "XYZBC"; parts: MachinePartDraft[]; cOriginMm: Vec3; tiltOffsetMm: Vec3; physicalZero: PhysicalZero }): Promise<ImportedMachine> {
  const byRole = validateMachineAssembly(input.coordinates, input.parts);
  for (const [axis, value] of Object.entries(input.physicalZero)) if (!Number.isFinite(value)) throw new Error(`${axis.toUpperCase()} 物理零点无效`);
  const urdf = machineUrdf(input.coordinates, input.cOriginMm, input.tiltOffsetMm);
  const structure: MachineStructure = parseMachineUrdf(urdf);
  const metadata: ManagedMachineGeometry = {
    schema: "linuxcnc.machine-manager/1", coordinates: input.coordinates, cOriginMm: [...input.cOriginMm], tiltOffsetMm: [...input.tiltOffsetMm], physicalZero: { ...input.physicalZero },
    parts: MACHINE_PART_ROLES.map((role) => { const part = byRole.get(role)!; return { id: part.id, filename: part.filename, hash: part.hash, role, parentId: part.parentId, originMm: [...part.originMm], translationMm: [...part.translationMm], rotationDeg: [...part.rotationDeg], visible: part.visible, nodeParents: { ...(part.nodeParents ?? {}) }, nodeRoles: { ...(part.nodeRoles ?? {}) }, nodeTransforms: Object.fromEntries(Object.entries(part.nodeTransforms ?? {}).map(([path, transform]) => [path, { translationMm: [...transform.translationMm] as Vec3, rotationDeg: [...transform.rotationDeg] as Vec3 }])) }; }),
  };
  const scene = new THREE.Group(); scene.name = `Managed_${input.coordinates}`;
  scene.userData.linuxcncMachine = { schema: "linuxcnc.machine/1", urdf };
  scene.userData.linuxcncMachineGeometry = metadata;
  const frame = new THREE.Group(); frame.name = "urdf_frame"; frame.rotation.x = -Math.PI / 2; scene.add(frame);
  const links = Object.fromEntries(MACHINE_PART_ROLES.map((role) => { const node = new THREE.Group(); node.name = role; node.userData.link = role; return [role, node]; })) as Record<MachinePartRole, THREE.Group>;
  frame.add(links.base); links.base.add(links.x, links.tilt); links.x.add(links.y); links.y.add(links.z); links.z.add(links.tool); links.tilt.add(links.c); links.c.add(links.work);
  for (const joint of structure.joints) links[joint.child as MachinePartRole].position.fromArray(joint.origin);
  for (const role of MACHINE_PART_ROLES) {
    const part = byRole.get(role)!;
    const placement = new THREE.Group(); placement.name = `component_${part.id}`; placement.visible = part.visible;
    placement.position.fromArray(part.translationMm).multiplyScalar(.001);
    placement.rotation.set(...part.rotationDeg.map(THREE.MathUtils.degToRad) as Vec3);
    placement.userData.machinePart = { role, filename: part.filename, hash: part.hash, originMm: [...part.originMm] };
    const { scene: visual } = configureMachinePartScene(cloneMachinePartScene(part.scene), part); visual.position.add(new THREE.Vector3(...part.originMm).multiplyScalar(-.001)); placement.add(visual); links[role].add(placement);
  }
  links.base.userData.linuxcncMachineGeometry = metadata;
  const signature = JSON.stringify({ metadata, hashes: MACHINE_PART_ROLES.map((role) => byRole.get(role)!.hash) });
  const hash = [...new Uint8Array(await crypto.subtle.digest("SHA-256", new TextEncoder().encode(signature)))].map((number) => number.toString(16).padStart(2, "0")).join("");
  return { scene, structure, hash, filename: `managed-${input.coordinates.toLowerCase()}.glb`, physicalZero: { ...input.physicalZero } };
}