import type { CircuitJson } from "circuit-json"
import Color from "color"
import { jscadToStep } from "jscad-to-step"
import {
  ApplicationContext,
  ApplicationProtocolDefinition,
  AdvancedBrepShapeRepresentation,
  Axis2Placement3D,
  CartesianPoint,
  ContextDependentShapeRepresentation,
  Direction,
  ItemDefinedTransformation,
  NextAssemblyUsageOccurrence,
  Product,
  ProductContext,
  ProductDefinition,
  ProductDefinitionContext,
  ProductDefinitionFormation,
  ProductDefinitionShape,
  Repository,
  MechanicalDesignGeometricPresentationRepresentation,
  ShapeDefinitionRepresentation,
  ShapeRepresentation,
  Unknown,
  parseRepository,
  type Entity,
  type Ref,
} from "stepts"
import {
  importStepModelGeometry,
  mergeExternalStepModels,
} from "./step-model-merger"
import { generateComponentMeshes } from "./mesh-generation"
import { generatePanelSolids } from "./panel-geometry"
import type { GeneratedSceneSolid } from "./scene-geometry"
import { readStepFile } from "./step-model-merger/read-step-file"
import { transformStepGeometry } from "./transform-step-geometry"
import {
  PCB_SUBSTRATE_COLOR_RGB,
  createStyleCache,
  createStyledItems,
} from "./step-style-utils"
import {
  asVector3,
  toRadians,
  transformDirection,
} from "./step-model-merger/vector-utils"
import { VERSION } from "./version"

/** Export all CAD parts and the PCB as distinct named children of one STEP assembly. */
export async function circuitJsonToAssemblyStep(
  circuitJson: CircuitJson,
  {
    productName = "Combined Assembly",
    pcbSteps = [],
    fsMap,
    boardThickness = 1.6,
    includeExternalMeshes = true,
  }: {
    productName?: string
    pcbSteps?: { id: string; name: string; stepText: string }[]
    fsMap?: Record<string, string>
    boardThickness?: number
    includeExternalMeshes?: boolean
  } = {},
): Promise<string> {
  const parts = circuitJson.filter((item) => item.type === "source_component")
  const models = circuitJson.filter((item) => item.type === "cad_component")
  const panel = circuitJson.find((item) => item.type === "pcb_panel")
  const pcbComponents = circuitJson.filter(
    (item) => item.type === "pcb_component",
  )
  if (
    pcbSteps.length === 0 &&
    !panel &&
    models.length === 0 &&
    pcbComponents.length === 0 &&
    parts.length === 0
  ) {
    throw new Error("No assembly geometry found to export")
  }
  for (const part of parts) {
    if (
      ["printedpart", "part"].includes(part.ftype) &&
      !models.some(
        (model) => model.source_component_id === part.source_component_id,
      )
    ) {
      throw new Error(`Part "${part.name}" needs CAD geometry for STEP export`)
    }
  }

  const repo = new Repository()
  const appContext = repo.add(new ApplicationContext("automotive design"))
  repo.add(
    new ApplicationProtocolDefinition(
      "international standard",
      "automotive_design",
      2010,
      appContext,
    ),
  )
  const productContext = repo.add(
    new ProductContext("", appContext, "mechanical"),
  )
  const definitionContext = repo.add(
    new ProductDefinitionContext("part definition", appContext, "design"),
  )
  const createProduct = (id: string, name: string) => {
    const product = repo.add(
      new Product(id, name, `Generated by circuit-json-to-step v${VERSION}`, [
        productContext,
      ]),
    )
    const formation = repo.add(new ProductDefinitionFormation("", "", product))
    const definition = repo.add(
      new ProductDefinition("design", "", formation, definitionContext),
    )
    const shape = repo.add(new ProductDefinitionShape("", "", definition))
    return { definition, shape }
  }

  const lengthUnit = repo.add(
    new Unknown("", [
      "( LENGTH_UNIT() NAMED_UNIT(*) SI_UNIT(.MILLI.,.METRE.) )",
    ]),
  )
  const angleUnit = repo.add(
    new Unknown("", [
      "( NAMED_UNIT(*) PLANE_ANGLE_UNIT() SI_UNIT($,.RADIAN.) )",
    ]),
  )
  const solidAngleUnit = repo.add(
    new Unknown("", [
      "( NAMED_UNIT(*) SI_UNIT($,.STERADIAN.) SOLID_ANGLE_UNIT() )",
    ]),
  )
  const uncertainty = repo.add(
    new Unknown("UNCERTAINTY_MEASURE_WITH_UNIT", [
      "LENGTH_MEASURE(1.E-07)",
      `${lengthUnit}`,
      "'distance_accuracy_value'",
      "'Maximum Tolerance'",
    ]),
  )
  const context = repo.add(
    new Unknown("", [
      `( GEOMETRIC_REPRESENTATION_CONTEXT(3) GLOBAL_UNCERTAINTY_ASSIGNED_CONTEXT((${uncertainty})) GLOBAL_UNIT_ASSIGNED_CONTEXT((${lengthUnit},${angleUnit},${solidAngleUnit})) REPRESENTATION_CONTEXT('','3D') )`,
    ]),
  )
  const createPlacement = (
    position = { x: 0, y: 0, z: 0 },
    rotation = { x: 0, y: 0, z: 0 },
  ) => {
    const radians = toRadians(rotation)
    const axis = repo.add(
      new Direction("", ...transformDirection([0, 0, 1], radians)),
    )
    const referenceDirection = repo.add(
      new Direction("", ...transformDirection([1, 0, 0], radians)),
    )
    const point = repo.add(
      new CartesianPoint("", position.x, position.y, position.z),
    )
    return repo.add(new Axis2Placement3D("", point, axis, referenceDirection))
  }

  const assembly = createProduct(productName, productName)
  const origin = createPlacement()
  const representation = repo.add(
    new ShapeRepresentation(productName, [origin], context),
  )
  repo.add(new ShapeDefinitionRepresentation(assembly.shape, representation))

  const addChild = (
    id: string,
    name: string,
    childRepresentation: Ref<Entity>,
    childOrigin: Ref<Axis2Placement3D>,
    position = { x: 0, y: 0, z: 0 },
    rotation = { x: 0, y: 0, z: 0 },
  ) => {
    const child = createProduct(id, name)
    repo.add(
      new ShapeDefinitionRepresentation(child.shape, childRepresentation),
    )
    const placement = createPlacement(position, rotation)
    representation.resolve(repo).items.push(placement)
    const usage = repo.add(
      new NextAssemblyUsageOccurrence(
        id,
        name,
        "",
        assembly.definition,
        child.definition,
        name,
      ),
    )
    const transformation = repo.add(
      new ItemDefinedTransformation("", "", childOrigin, placement),
    )
    const relationship = repo.add(
      new Unknown("", [
        `( REPRESENTATION_RELATIONSHIP('','',${childRepresentation},${representation}) REPRESENTATION_RELATIONSHIP_WITH_TRANSFORMATION(${transformation}) SHAPE_REPRESENTATION_RELATIONSHIP() )`,
      ]),
    )
    const usageShape = repo.add(
      new Unknown("PRODUCT_DEFINITION_SHAPE", ["''", "''", `${usage}`]),
    )
    repo.add(new ContextDependentShapeRepresentation(relationship, usageShape))
  }
  const addStepChild = (
    id: string,
    name: string,
    geometry: string,
    position?: { x: number; y: number; z: number },
    rotation?: { x: number; y: number; z: number },
    color?: [number, number, number],
  ) => {
    const imported = importStepModelGeometry(repo, geometry, name, { context })
    if (imported.solids.length === 0) {
      throw new Error(`Part "${name}" has no solid geometry to export`)
    }
    imported.representation.resolve(repo).items.unshift(imported.origin)
    if (color)
      repo.add(
        new MechanicalDesignGeometricPresentationRepresentation(
          "",
          createStyledItems(repo, {
            itemRefs: imported.solids,
            rgb: color,
            styleCache: createStyleCache(),
          }),
          context,
        ),
      )
    addChild(
      id,
      name,
      imported.representation,
      imported.origin,
      position,
      rotation,
    )
  }
  for (const pcb of pcbSteps)
    addStepChild(
      pcb.id,
      pcb.name,
      pcb.stepText,
      undefined,
      undefined,
      PCB_SUBSTRATE_COLOR_RGB,
    )

  if (panel) {
    addGeneratedChild(
      generatePanelSolids({ repo, circuitJson, boardThickness }),
      panel.pcb_panel_id,
      "PCB Panel",
      PCB_SUBSTRATE_COLOR_RGB,
    )
  }

  const stepContents = new Map(Object.entries(fsMap ?? {}))
  const validatedStepModels = new Set<string>()
  const importedStepModels = new Map<
    string,
    Pick<
      ReturnType<typeof importStepModelGeometry>,
      "origin" | "representation"
    >
  >()
  for (const model of models) {
    const part = parts.find(
      (item) => item.source_component_id === model.source_component_id,
    )
    const name = part?.name ?? model.cad_component_id
    if (model.model_jscad) {
      const modelOrigin = asVector3(model.model_origin_position)
      const scale = model.model_unit_to_mm_scale_factor ?? 1
      // Apply the same CAD color override as the assembly viewer.
      const shape = model.color
        ? {
            type: "colorize",
            color: Color(model.color).rgb().unitArray().slice(0, 3),
            shape: model.model_jscad!,
          }
        : model.model_jscad!
      // Keep geometry local; rigid placement belongs to the assembly occurrence.
      const geometry = jscadToStep({
        type: "scale",
        factors: [scale, scale, scale],
        shape: {
          type: "translate",
          vector: [-modelOrigin.x, -modelOrigin.y, -modelOrigin.z],
          shape,
        },
      })
      addStepChild(
        model.cad_component_id,
        name,
        geometry,
        asVector3(model.position),
        asVector3(model.rotation),
      )
      continue
    }
    const componentJson = circuitJson.filter((item) =>
      item.type === "cad_component"
        ? item.cad_component_id === model.cad_component_id
        : item.type === "pcb_component"
          ? item.pcb_component_id === model.pcb_component_id
          : true,
    )
    if (model.model_step_url) {
      let geometry = stepContents.get(model.model_step_url)
      if (geometry === undefined) {
        geometry = await readStepFile(model.model_step_url)
        stepContents.set(model.model_step_url, geometry)
      }
      if (
        !model.pcb_component_id ||
        model.model_origin_position ||
        model.model_unit_to_mm_scale_factor !== undefined
      ) {
        const modelOrigin = asVector3(model.model_origin_position)
        const scale = model.model_unit_to_mm_scale_factor ?? 1
        const modelKey = JSON.stringify([
          model.model_step_url,
          modelOrigin,
          scale,
          model.color,
        ])
        let imported = importedStepModels.get(modelKey)
        if (!imported) {
          const geometryModel = importStepModelGeometry(repo, geometry, name, {
            context,
            rejectAssemblies: true,
            color: model.color,
          })
          if (geometryModel.solids.length === 0)
            throw new Error(`Part "${name}" has no solid geometry to export`)
          geometryModel.representation
            .resolve(repo)
            .items.unshift(geometryModel.origin)
          transformStepGeometry(geometryModel.entries, modelOrigin, scale)
          imported = {
            origin: geometryModel.origin,
            representation: geometryModel.representation,
          }
          importedStepModels.set(modelKey, imported)
        }
        addChild(
          model.cad_component_id,
          name,
          imported.representation,
          imported.origin,
          asVector3(model.position),
          asVector3(model.rotation),
        )
        continue
      }
      if (!validatedStepModels.has(model.model_step_url)) {
        if (
          parseRepository(geometry)
            .entries()
            .some(
              ([, entity]) => entity.type === "NEXT_ASSEMBLY_USAGE_OCCURRENCE",
            )
        ) {
          throw new Error(
            `Part "${name}" uses a nested STEP assembly, which is not supported for assembly export`,
          )
        }
        validatedStepModels.add(model.model_step_url)
      }
      const result = await mergeExternalStepModels({
        repo,
        circuitJson: componentJson,
        boardThickness: thicknessForComponents(componentJson),
        fsMap: { [model.model_step_url]: geometry },
      })
      if (result.solids.length === 0) {
        throw new Error(`Part "${name}" could not load its STEP geometry`)
      }
      // Keep mapped instances in the representation; extracting only BREPs would lose their placements.
      const childRepresentation = repo.add(
        new ShapeRepresentation(name, [origin, ...result.solids], context),
      )
      addChild(model.cad_component_id, name, childRepresentation, origin)
      continue
    }
    const hasExternalMesh = Boolean(
      model.model_glb_url ||
        model.model_gltf_url ||
        model.model_stl_url ||
        model.model_obj_url ||
        model.model_3mf_url,
    )
    if (hasExternalMesh && !includeExternalMeshes)
      throw new Error(`External mesh export is disabled for part "${name}"`)
    await addMeshChildren(
      componentJson,
      model.cad_component_id,
      name,
      hasExternalMesh || part?.ftype === "printedpart",
    )
  }
  // Footprints without cad_components (e.g. silkscreen logos and mounting
  // holes) have no body in the 3D viewer. Their PCB geometry is already above.
  if (representation.resolve(repo).items.length === 1) {
    throw new Error("No assembly geometry found to export")
  }

  return repo.toPartFile({ name: productName })

  async function addMeshChildren(
    componentJson: CircuitJson,
    id: string,
    name: string,
    requireMeshGeometry = false,
  ) {
    const generated = await generateComponentMeshes({
      repo,
      circuitJson: componentJson,
      boardThickness: thicknessForComponents(componentJson),
      includeExternalMeshes,
      requireMeshGeometry,
    })
    addGeneratedChild(generated, id, name)
  }

  function addGeneratedChild(
    generated: GeneratedSceneSolid[],
    id: string,
    name: string,
    color: [number, number, number] = [0.75, 0.75, 0.75],
  ) {
    if (generated.length === 0)
      throw new Error(`Part "${name}" has no supported geometry to export`)
    for (const solid of generated) solid.solid.resolve(repo).name = name
    const styleCache = createStyleCache()
    const styles = generated.flatMap((solid) => [
      ...solid.styledItems,
      ...(!solid.usesIntrinsicFaceStyles
        ? createStyledItems(repo, {
            itemRefs: solid.styleTargets,
            rgb: color,
            styleCache,
          })
        : []),
    ])
    if (styles.length > 0)
      repo.add(
        new MechanicalDesignGeometricPresentationRepresentation(
          "",
          styles,
          context,
        ),
      )
    const childRepresentation = repo.add(
      new AdvancedBrepShapeRepresentation(
        name,
        [origin, ...generated.map((solid) => solid.solid)],
        context,
      ),
    )
    addChild(id, name, childRepresentation, origin)
  }

  function thicknessForComponents(componentJson: CircuitJson) {
    const boards = circuitJson.filter((item) => item.type === "pcb_board")
    if (boards.length <= 1) return boardThickness
    const component = componentJson.find(
      (item) => item.type === "pcb_component",
    )
    if (!component) return boardThickness
    const record = component as unknown as Record<string, unknown>
    const matches = boards.filter((board) => {
      const boardRecord = board as unknown as Record<string, unknown>
      return (
        record.subcircuit_id &&
        record.subcircuit_id === boardRecord.subcircuit_id
      )
    })
    if (matches.length === 1) return matches[0]!.thickness ?? boardThickness
    if (new Set(boards.map((board) => board.thickness)).size === 1)
      return boards[0]!.thickness ?? boardThickness
    throw new Error(
      "Cannot determine PCB thickness for a component in a multi-board assembly",
    )
  }
}