import * as THREE from 'three'
import {
  createExtrusionMaterial,
  createLineMaterial,
  createSharedUniforms,
  MAX_TOOLS
} from './ExtrusionMaterial.js'
import {
  ALL_TYPES_MASK,
  EXTRUSION_TYPES_MASK,
  FEATURE_INFO,
  LINE_TYPES_MASK,
  MAX_FEATURE_TYPES
} from './featureTypes.js'
import { TOOL_PALETTE } from './colorModes.js'

// GPU representation of a SegmentStore. Every segment is one instance; the
// box geometry renders extrusions and the two-vertex line geometry renders
// travels (and everything in "lines" mode). Both geometries share the same
// attribute objects, so each field lives in GPU memory exactly once.

function createBoxBase() {
  const box = new THREE.BoxGeometry(1, 1, 1)
  box.translate(0.5, 0, 0)
  return box
}

function createLineBase() {
  const geometry = new THREE.BufferGeometry()
  geometry.setAttribute(
    'position',
    new THREE.Float32BufferAttribute([0, 0, 0, 1, 0, 0], 3)
  )
  return geometry
}

const INSTANCE_FIELDS = [
  ['aType', 'type'],
  ['aTool', 'tool'],
  ['aFan', 'fan'],
  ['aLayer', 'layer'],
  ['aTemp', 'temp'],
  ['aWidth', 'width'],
  ['aHeight', 'height'],
  ['aFeed', 'feed'],
  ['aAccel', 'accel'],
  ['aActualSpeed', 'actualSpeed'],
  ['aLayerTime', 'layerTime']
]

export default class SegmentMesh {
  constructor() {
    this.group = new THREE.Group()
    this.shared = createSharedUniforms()
    this.boxMaterial = createExtrusionMaterial(this.shared)
    this.lineMaterial = createLineMaterial(this.shared)
    this.boxBase = createBoxBase()
    this.lineBase = createLineBase()
    // World-space light directions; converted to view space per frame.
    this.keyLightWorld = new THREE.Vector3(-0.45, -0.35, 1).normalize()
    this.fillLightWorld = new THREE.Vector3(0.7, 0.6, 0.25).normalize()

    this.boxGeometry = null
    this.lineGeometry = null
    this.mesh = null
    this.lines = null
    this.positionBuffer = null
    this.attributes = {}

    this.store = null
    this.generation = -1
    this.uploaded = 0
    this.plannerUploaded = false

    this.renderMode = 'extrusion'
    this.visibleMask = ALL_TYPES_MASK & ~(1 << 1) // travel hidden by default
    this.applyMasks()
    this.setTypeColors()
    this.setToolColors()
  }

  setStore(store) {
    this.store = store
    this.generation = -1
    this.uploaded = 0
    this.plannerUploaded = false
    this.shared.uVisibleEnd.value = 0
    this.sync()
  }

  // Upload anything new since the last sync. Cheap when nothing changed.
  sync() {
    const store = this.store
    if (!store) {
      this.disposeGeometry()
      return
    }
    if (store.generation !== this.generation) {
      this.rebuild()
    } else if (store.count > this.uploaded) {
      this.uploadRange(this.uploaded, store.count)
    }
    if (store.plannerReady && !this.plannerUploaded) {
      this.uploadPlanner()
    }
    this.uploaded = store.count
    if (this.boxGeometry) this.boxGeometry.instanceCount = store.count
    if (this.lineGeometry) this.lineGeometry.instanceCount = store.count
  }

  rebuild() {
    this.disposeGeometry()
    const store = this.store
    this.generation = store.generation

    const positionBuffer = new THREE.InstancedInterleavedBuffer(
      store.positions,
      3,
      1
    )
    positionBuffer.setUsage(THREE.DynamicDrawUsage)
    this.positionBuffer = positionBuffer
    const instanceStart = new THREE.InterleavedBufferAttribute(
      positionBuffer,
      3,
      0
    )
    const instanceEnd = new THREE.InterleavedBufferAttribute(
      positionBuffer,
      3,
      3
    )

    this.attributes = {}
    for (const [name, field] of INSTANCE_FIELDS) {
      const attribute = new THREE.InstancedBufferAttribute(store[field], 1)
      attribute.setUsage(THREE.DynamicDrawUsage)
      this.attributes[name] = attribute
    }

    const applyInstanceAttributes = (geometry) => {
      geometry.setAttribute('instanceStart', instanceStart)
      geometry.setAttribute('instanceEnd', instanceEnd)
      for (const name of Object.keys(this.attributes)) {
        geometry.setAttribute(name, this.attributes[name])
      }
      geometry.instanceCount = store.count
      // The bounding volume would otherwise be computed from the base
      // geometry only; culling is skipped since the model fills the view.
      geometry.boundingSphere = new THREE.Sphere(new THREE.Vector3(), 1e6)
    }

    const boxGeometry = new THREE.InstancedBufferGeometry()
    boxGeometry.index = this.boxBase.index
    boxGeometry.setAttribute('position', this.boxBase.attributes.position)
    boxGeometry.setAttribute('normal', this.boxBase.attributes.normal)
    applyInstanceAttributes(boxGeometry)

    const lineGeometry = new THREE.InstancedBufferGeometry()
    lineGeometry.setAttribute('position', this.lineBase.attributes.position)
    applyInstanceAttributes(lineGeometry)

    this.boxGeometry = boxGeometry
    this.lineGeometry = lineGeometry

    this.mesh = new THREE.Mesh(boxGeometry, this.boxMaterial)
    this.mesh.frustumCulled = false
    this.mesh.visible = this.renderMode === 'extrusion'
    this.lines = new THREE.LineSegments(lineGeometry, this.lineMaterial)
    this.lines.frustumCulled = false
    this.group.add(this.mesh, this.lines)

    this.uploaded = store.count
    this.plannerUploaded = false
  }

  uploadRange(from, to) {
    if (to <= from) return
    const count = to - from
    // Segment `to - 1` ends at vertex `to`, so include one extra vertex.
    this.positionBuffer.addUpdateRange(from * 3, (count + 1) * 3)
    this.positionBuffer.needsUpdate = true
    for (const [name] of INSTANCE_FIELDS) {
      if (name === 'aActualSpeed' || name === 'aLayerTime') continue
      const attribute = this.attributes[name]
      attribute.addUpdateRange(from, count)
      attribute.needsUpdate = true
    }
  }

  uploadPlanner() {
    for (const name of ['aActualSpeed', 'aLayerTime']) {
      const attribute = this.attributes[name]
      if (!attribute) continue
      attribute.addUpdateRange(0, this.store.count)
      attribute.needsUpdate = true
    }
    this.plannerUploaded = true
  }

  disposeGeometry() {
    if (this.mesh) this.group.remove(this.mesh)
    if (this.lines) this.group.remove(this.lines)
    this.boxGeometry?.dispose()
    this.lineGeometry?.dispose()
    this.boxGeometry = null
    this.lineGeometry = null
    this.mesh = null
    this.lines = null
    this.positionBuffer = null
    this.attributes = {}
  }

  dispose() {
    this.disposeGeometry()
    this.boxBase.dispose()
    this.lineBase.dispose()
    this.boxMaterial.dispose()
    this.lineMaterial.dispose()
  }

  // ---- uniform driven state -------------------------------------------

  // Call before rendering so the fixed world lights follow camera rotation.
  updateLights(camera) {
    const { uKeyLightDir, uFillLightDir } = this.boxMaterial.uniforms
    uKeyLightDir.value
      .copy(this.keyLightWorld)
      .transformDirection(camera.matrixWorldInverse)
    uFillLightDir.value
      .copy(this.fillLightWorld)
      .transformDirection(camera.matrixWorldInverse)
  }

  applyMasks() {
    const mask = this.visibleMask & ALL_TYPES_MASK
    if (this.renderMode === 'lines') {
      this.boxMaterial.uniforms.uTypeMask.value = 0
      this.lineMaterial.uniforms.uTypeMask.value = mask
      this.lineMaterial.uniforms.uOpacity.value = 1
      this.lineMaterial.transparent = false
      this.lineMaterial.depthWrite = true
    } else {
      this.boxMaterial.uniforms.uTypeMask.value = mask & EXTRUSION_TYPES_MASK
      this.lineMaterial.uniforms.uTypeMask.value = mask & LINE_TYPES_MASK
      this.lineMaterial.uniforms.uOpacity.value = 0.7
      this.lineMaterial.transparent = true
      this.lineMaterial.depthWrite = false
    }
    if (this.mesh) this.mesh.visible = this.renderMode === 'extrusion'
  }

  setRenderMode(mode) {
    this.renderMode = mode === 'lines' ? 'lines' : 'extrusion'
    this.applyMasks()
  }

  setVisibleMask(mask) {
    this.visibleMask = mask
    this.applyMasks()
  }

  setColorMode(shaderMode, logScale = false) {
    this.shared.uColorMode.value = shaderMode | 0
    this.shared.uLogScale.value = logScale ? 1 : 0
  }

  setRange(min, max) {
    this.shared.uRangeMin.value = Number.isFinite(min) ? min : 0
    this.shared.uRangeMax.value = Number.isFinite(max) ? max : 1
  }

  setLayerRange(min, max) {
    this.shared.uLayerMin.value = min
    this.shared.uLayerMax.value = max
  }

  // Show segments [0, end) fully and fraction `t` of segment `end - 1`.
  setProgress(end, t = 1) {
    const count = this.store ? this.store.count : 0
    const clamped = Math.max(0, Math.min(count, end | 0))
    this.shared.uVisibleEnd.value = clamped
    this.shared.uProgressT.value = Math.max(0, Math.min(1, t))
  }

  showAll() {
    this.shared.uVisibleEnd.value = this.store ? this.store.count : 0
    this.shared.uProgressT.value = 1
  }

  setTypeColors(overrides = {}) {
    const colors = this.shared.uTypeColors.value
    for (let i = 0; i < MAX_FEATURE_TYPES; i++) {
      const info = FEATURE_INFO[i]
      const hex = overrides[i] || info?.color || '#888888'
      colors[i].set(hex)
    }
  }

  setToolColors(hexColors = []) {
    const colors = this.shared.uToolColors.value
    for (let i = 0; i < MAX_TOOLS; i++) {
      const hex = hexColors[i] || TOOL_PALETTE[i % TOOL_PALETTE.length]
      try {
        colors[i].set(hex)
      } catch {
        colors[i].set(TOOL_PALETTE[i % TOOL_PALETTE.length])
      }
    }
  }
}