import type { DepsgraphEvaluationIR } from "./depsgraph";
import { analyzeNonMeshExport, mapBinaryNonMeshForExport, mapEvaluatedNonMeshForExport, type NonMeshExportMapping } from "./nonmesh-export";
import type { MeshSummaryIR, SceneNodeIR, SceneSnapshotIR } from "./scene-ir";
import type { MeshGeometryBuffer } from "./web-engine";
import type { NonMeshGeometryChunk, ReassembledNonMeshAttribute, ReassembledNonMeshGeometry } from "./nonmesh-binary";

export type USDSerializationErrorCode = "USD_GEOMETRY_BUFFER_MISSING" | "USD_GEOMETRY_INVALID" | "USD_EMPTY_SCENE";

export interface USDSerializationError {
  code: USDSerializationErrorCode;
  message: string;
  id?: string;
}

export interface USDExportReport {
  schemaVersion: 1;
  canExport: boolean;
  meshTargets: Array<{ meshId: string; target: "UsdGeomMesh" | "UsdGeomBasisCurves" | "UsdGeomPoints"; status: "MAPPED" | "BLOCKED"; errorCode?: "SUMMARY_ONLY_MESH" }>;
  nonMeshTargets: NonMeshExportMapping[];
  errors?: USDSerializationError[];
}

export interface USDExportResult {
  report: USDExportReport;
  usda?: Uint8Array;
}

export function analyzeUSDExport(snapshot: SceneSnapshotIR, depsgraph?: DepsgraphEvaluationIR): USDExportReport {
  const meshTarget = (mesh: MeshSummaryIR): "UsdGeomMesh" | "UsdGeomBasisCurves" | "UsdGeomPoints" => mesh.topology === "lines" ? "UsdGeomBasisCurves" : mesh.topology === "points" ? "UsdGeomPoints" : "UsdGeomMesh";
  const meshTargets = snapshot.meshes.map((mesh) => mesh.geometryStatus === "summary-only"
    ? { meshId: mesh.id, target: meshTarget(mesh), status: "BLOCKED" as const, errorCode: "SUMMARY_ONLY_MESH" as const }
    : { meshId: mesh.id, target: meshTarget(mesh), status: "MAPPED" as const });
  const nonMesh = analyzeNonMeshExport(snapshot, depsgraph);
  return {
    schemaVersion: 1,
    canExport: meshTargets.every((mapping) => mapping.status === "MAPPED") && nonMesh.canExportUSD,
    meshTargets,
    nonMeshTargets: nonMesh.mappings,
  };
}

function values<T extends Float32Array | Uint32Array>(
  payload: MeshGeometryBuffer | undefined,
  summary: MeshSummaryIR,
  key: "positions" | "indices" | "edgeVertexIndices",
): T | undefined {
  const inline = summary[key];
  if (inline) return (key === "positions" ? Float32Array.from(inline) : Uint32Array.from(inline)) as T;
  const buffer = payload?.[key];
  if (!buffer) return undefined;
  return new (key === "positions" ? Float32Array : Uint32Array)(buffer) as T;
}

function usdString(value: string): string {
  return value.replaceAll("\\", "\\\\").replaceAll("\"", "\\\"").replaceAll("\n", "\\n").replaceAll("\r", "\\r");
}

function primIdentifier(value: string, used: Set<string>): string {
  const base = value.normalize("NFKD").replace(/[^A-Za-z0-9_]/g, "_").replace(/^[^A-Za-z_]/, "_$&") || "Object";
  let result = base;
  for (let suffix = 2; used.has(result); suffix++) result = `${base}_${suffix}`;
  used.add(result);
  return result;
}

function number(value: number): string {
  if (!Number.isFinite(value)) throw new Error("USD geometry contains a non-finite number");
  const normalized = Object.is(value, -0) ? 0 : value;
  return Number.isInteger(normalized) ? String(normalized) : normalized.toPrecision(9).replace(/(?:\.0+|(?:(\.\d*?)0+))(?=e|$)/, "$1");
}

function tuples(values: ArrayLike<number>, width: number): string {
  const result: string[] = [];
  for (let index = 0; index < values.length; index += width) {
    result.push(`(${Array.from({ length: width }, (_, component) => number(values[index + component])).join(", ")})`);
  }
  return `[${result.join(", ")}]`;
}

function integers(values: ArrayLike<number>): string {
  return `[${Array.from(values, number).join(", ")}]`;
}

function scalars(values: ArrayLike<number>, booleanValues = false): string {
  return `[${Array.from(values, (value) => booleanValues ? value ? "true" : "false" : number(value)).join(", ")}]`;
}

function primvarName(value: string): string {
  return value.replace(/[^A-Za-z0-9_]/g, "_").replace(/^[^A-Za-z_]/, "_$&") || "attribute";
}

function attributeValues(attribute: ReassembledNonMeshAttribute, pointOrder?: readonly number[]): ArrayLike<number> {
  if (attribute.domain !== "POINT" || !pointOrder) return attribute.values;
  const result = attribute.storage === "FLOAT32" ? new Float32Array(pointOrder.length * attribute.components) :
    attribute.storage === "INT32" ? new Int32Array(pointOrder.length * attribute.components) : new Uint8Array(pointOrder.length * attribute.components);
  for (const [target, source] of pointOrder.entries()) for (let component = 0; component < attribute.components; component++) {
    result[target * attribute.components + component] = attribute.values[source * attribute.components + component];
  }
  return result;
}

function primvar(attribute: ReassembledNonMeshAttribute, pointOrder?: readonly number[]): string {
  const source = attributeValues(attribute, pointOrder);
  const interpolation = attribute.domain === "POINT" ? "vertex" : attribute.domain === "CURVE" ? "uniform" : "constant";
  const name = primvarName(attribute.name);
  if (attribute.dataType === "BOOL") return `    bool[] primvars:${name} = ${scalars(source, true)} (interpolation = "${interpolation}")`;
  if (attribute.dataType === "INT") return `    int[] primvars:${name} = ${integers(source)} (interpolation = "${interpolation}")`;
  if (attribute.dataType === "BYTE_COLOR") {
    const normalized = Float32Array.from(source, (value) => value / 255);
    return `    color4f[] primvars:${name} = ${tuples(normalized, 4)} (interpolation = "${interpolation}")`;
  }
  const type = attribute.dataType === "FLOAT" ? "float" : attribute.dataType === "FLOAT2" ? "float2" : attribute.dataType === "FLOAT3" ? "vector3f" : "color4f";
  const encoded = attribute.components === 1 ? scalars(source) : tuples(source, attribute.components);
  return `    ${type}[] primvars:${name} = ${encoded} (interpolation = "${interpolation}")`;
}

function matrix(value: readonly number[]): string {
  if (value.length !== 16) throw new Error("USD object matrix must contain 16 values");
  return `(${[0, 4, 8, 12].map((offset) => `(${[0, 1, 2, 3].map((component) => number(value[offset + component])).join(", ")})`).join(", ")})`;
}

function edgeChains(edges: Uint32Array, vertexCount: number): number[][] {
  const adjacency = Array.from({ length: vertexCount }, () => [] as Array<{ edge: number; vertex: number }>);
  for (let edge = 0; edge < edges.length / 2; edge++) {
    const left = edges[edge * 2];
    const right = edges[edge * 2 + 1];
    if (left >= vertexCount || right >= vertexCount || left === right) throw new Error("USD line geometry contains an invalid edge");
    adjacency[left].push({ edge, vertex: right });
    adjacency[right].push({ edge, vertex: left });
  }
  const visited = new Uint8Array(edges.length / 2);
  const chains: number[][] = [];
  const starts = Array.from({ length: vertexCount }, (_, vertex) => vertex).filter((vertex) => adjacency[vertex].length !== 2);
  const walk = (start: number, firstEdge?: number): void => {
    const chain = [start];
    let current = start;
    let edge = firstEdge ?? adjacency[current].find((candidate) => !visited[candidate.edge])?.edge;
    while (edge !== undefined && !visited[edge]) {
      visited[edge] = 1;
      const left = edges[edge * 2];
      const right = edges[edge * 2 + 1];
      current = current === left ? right : left;
      chain.push(current);
      edge = adjacency[current].find((candidate) => !visited[candidate.edge])?.edge;
    }
    if (chain.length > 1) chains.push(chain);
  };
  for (const start of starts) for (const candidate of adjacency[start]) if (!visited[candidate.edge]) walk(start, candidate.edge);
  for (let edge = 0; edge < visited.length; edge++) if (!visited[edge]) walk(edges[edge * 2], edge);
  return chains;
}

function geometryBlock(
  primName: string,
  blenderId: string,
  transform: readonly number[],
  summary: MeshSummaryIR,
  payload: MeshGeometryBuffer | undefined,
  rich?: ReassembledNonMeshGeometry,
): string {
  const positions = values<Float32Array>(payload, summary, "positions");
  if (!positions || positions.length !== summary.vertexCount * 3) throw new Error(`USD geometry ${summary.id} has no valid position buffer`);
  const header = `    customData = { string blenderId = "${usdString(blenderId)}" }`;
  const xform = `    matrix4d xformOp:transform = ${matrix(transform)}\n    uniform token[] xformOpOrder = ["xformOp:transform"]`;
  if (summary.topology === "points") {
    const widths = rich?.radii ? Float32Array.from(rich.radii, (radius) => radius * 2) : new Float32Array([0.01]);
    const widthInterpolation = rich?.radii ? "vertex" : "constant";
    const attributes = rich?.attributes.map((attribute) => primvar(attribute)).join("\n") ?? "";
    return `  def Points "${primName}" (\n${header}\n  ) {\n${xform}\n    point3f[] points = ${tuples(positions, 3)}\n    float[] widths = ${scalars(widths)} (interpolation = "${widthInterpolation}")${attributes ? `\n${attributes}` : ""}\n  }`;
  }
  if (summary.topology === "lines") {
    const edges = values<Uint32Array>(payload, summary, "edgeVertexIndices");
    if (!edges || edges.length !== summary.edgeCount * 2 || edges.length === 0) throw new Error(`USD line geometry ${summary.id} has no valid edge buffer`);
    const chains = rich?.curveOffsets ? Array.from({ length: rich.curveOffsets.length - 1 }, (_, curve) =>
      Array.from({ length: rich.curveOffsets![curve + 1] - rich.curveOffsets![curve] }, (__, point) => rich.curveOffsets![curve] + point)) : edgeChains(edges, summary.vertexCount);
    const points = new Float32Array(chains.reduce((total, chain) => total + chain.length, 0) * 3);
    const pointOrder: number[] = [];
    let cursor = 0;
    for (const chain of chains) for (const vertex of chain) {
      points.set(positions.subarray(vertex * 3, vertex * 3 + 3), cursor);
      pointOrder.push(vertex);
      cursor += 3;
    }
    const widths = rich?.radii ? Float32Array.from(pointOrder, (point) => rich.radii![point] * 2) : new Float32Array([0.01]);
    const widthInterpolation = rich?.radii ? "vertex" : "constant";
    const attributes = rich?.attributes.map((attribute) => primvar(attribute, pointOrder)).join("\n") ?? "";
    return `  def BasisCurves "${primName}" (\n${header}\n  ) {\n${xform}\n    uniform token type = "linear"\n    uniform token wrap = "nonperiodic"\n    int[] curveVertexCounts = ${integers(chains.map((chain) => chain.length))}\n    point3f[] points = ${tuples(points, 3)}\n    float[] widths = ${scalars(widths)} (interpolation = "${widthInterpolation}")${attributes ? `\n${attributes}` : ""}\n  }`;
  }
  const indices = values<Uint32Array>(payload, summary, "indices");
  if (!indices || indices.length !== (summary.triangleCount ?? summary.faceCount) * 3 || indices.some((index) => index >= summary.vertexCount)) {
    throw new Error(`USD mesh geometry ${summary.id} has no valid triangle buffer`);
  }
  const faceCounts = new Uint32Array(indices.length / 3).fill(3);
  return `  def Mesh "${primName}" (\n${header}\n  ) {\n${xform}\n    uniform token subdivisionScheme = "none"\n    uniform token orientation = "rightHanded"\n    point3f[] points = ${tuples(positions, 3)}\n    int[] faceVertexCounts = ${integers(faceCounts)}\n    int[] faceVertexIndices = ${integers(indices)}\n  }`;
}

function identity(): number[] {
  return [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1];
}

export function exportUSD(
  snapshot: SceneSnapshotIR,
  geometryBuffers: readonly MeshGeometryBuffer[] = [],
  depsgraph?: DepsgraphEvaluationIR,
  nonMeshGeometryBuffers: readonly NonMeshGeometryChunk[] = [],
): USDExportResult {
  const analysis = analyzeUSDExport(snapshot, depsgraph);
  if (!analysis.canExport) return { report: analysis };
  const evaluated = depsgraph ? mapEvaluatedNonMeshForExport(snapshot, geometryBuffers, depsgraph) : { snapshot, geometryBuffers: [...geometryBuffers] };
  let mapped: ReturnType<typeof mapBinaryNonMeshForExport>;
  try {
    mapped = mapBinaryNonMeshForExport(evaluated.snapshot, evaluated.geometryBuffers, nonMeshGeometryBuffers);
  }
  catch (error) {
    return { report: { ...analysis, canExport: false, errors: [{ code: "USD_GEOMETRY_INVALID", message: error instanceof Error ? error.message : "WNM geometry is invalid" }] } };
  }
  const bufferById = new Map(mapped.geometryBuffers.map((buffer) => [buffer.meshId, buffer]));
  const meshById = new Map(mapped.snapshot.meshes.map((mesh) => [mesh.id, mesh]));
  const used = new Set<string>();
  const blocks: string[] = [];
  const referenced = new Set<string>();
  const errors: USDSerializationError[] = [];
  const append = (summary: MeshSummaryIR, node?: SceneNodeIR): void => {
    const geometryId = summary.geometryBufferId ?? summary.id;
    try {
      blocks.push(geometryBlock(
        primIdentifier(node?.name ?? summary.name, used),
        node?.id ?? summary.id,
        node?.worldMatrix ?? identity(),
        summary,
        bufferById.get(geometryId) ?? bufferById.get(summary.id),
        mapped.geometryByMeshId.get(summary.id),
      ));
    }
    catch (error) {
      errors.push({
        code: error instanceof Error && error.message.includes("buffer") ? "USD_GEOMETRY_BUFFER_MISSING" : "USD_GEOMETRY_INVALID",
        message: error instanceof Error ? error.message : `USD geometry ${summary.id} is invalid`,
        id: summary.id,
      });
    }
  };
  for (const node of mapped.snapshot.nodes) {
    if (!node.dataId) continue;
    const summary = meshById.get(node.dataId);
    if (!summary) continue;
    referenced.add(summary.id);
    append(summary, node);
  }
  for (const summary of mapped.snapshot.meshes) if (!referenced.has(summary.id)) append(summary);
  if (blocks.length === 0 && errors.length === 0) errors.push({ code: "USD_EMPTY_SCENE", message: "The scene contains no serializable geometry" });
  if (errors.length > 0) return { report: { ...analysis, canExport: false, errors } };
  const source = `#usda 1.0\n(\n    defaultPrim = "Scene"\n    metersPerUnit = 1\n    upAxis = "Z"\n)\n\ndef Xform "Scene" {\n${blocks.join("\n\n")}\n}\n`;
  return { report: analysis, usda: new TextEncoder().encode(source) };
}