// SPDX-License-Identifier: MIT
// Copyright 2026 David Ronai

import { Geometry, Material, Mesh, Object3D, Texture, type Renderer } from 'mini3'
import {
  DomText,
  type VanillaDecorationInstances,
  type VanillaDomTextOptions,
  type VanillaGlyphInstances,
  type VanillaRectInstances,
  type VanillaSdfAtlas,
} from 'ultra-text/vanilla'

const GLYPH_WGSL = /* wgsl */ `
struct Camera { projectionViewMatrix: mat4x4f }
struct Model { modelMatrix: mat4x4f, prevModelMatrix: mat4x4f, color: vec4f, flags: vec4f }

@group(0) @binding(0) var<uniform> camera: Camera;
@group(0) @binding(1) var<uniform> model: Model;
@group(0) @binding(2) var samp: sampler;
@group(0) @binding(3) var atlas: texture_2d<f32>;

struct VSIn {
  @location(0) position: vec3f,
  @location(1) uv: vec2f,
  @location(2) color: vec4f,
  @location(3) params: vec4f,
}

struct VSOut {
  @builtin(position) pos: vec4f,
  @location(0) uv: vec2f,
  @location(1) color: vec4f,
  @location(2) params: vec4f,
}

@vertex fn vs(input: VSIn) -> VSOut {
  var out: VSOut;
  let world = model.modelMatrix * vec4f(input.position, 1.0);
  out.pos = camera.projectionViewMatrix * world;
  out.uv = input.uv;
  out.color = input.color;
  out.params = input.params;
  return out;
}

@fragment fn fs(input: VSOut) -> @location(0) vec4f {
  let sd = textureSample(atlas, samp, input.uv).r - 0.5;
  let unitRange = input.params.xy;
  let screenTexSize = vec2f(1.0) / max(fwidth(input.uv), vec2f(1.0e-5));
  let screenPxRange = max(0.5 * dot(unitRange, screenTexSize), 1.0);
  let screenPxDistance = screenPxRange * sd + input.params.z;
  let alpha = smoothstep(-0.5, 0.5, screenPxDistance) * input.color.a * model.color.a;
  if (alpha < 0.001) { discard; }
  return vec4f(input.color.rgb * model.color.rgb, alpha);
}
`

const RECT_WGSL = /* wgsl */ `
struct Camera { projectionViewMatrix: mat4x4f }
struct Model { modelMatrix: mat4x4f, prevModelMatrix: mat4x4f, color: vec4f, flags: vec4f }

@group(0) @binding(0) var<uniform> camera: Camera;
@group(0) @binding(1) var<uniform> model: Model;

struct VSIn {
  @location(0) position: vec3f,
  @location(1) localUv: vec2f,
  @location(2) fill: vec4f,
  @location(3) sizeRadius: vec4f,
}

struct VSOut {
  @builtin(position) pos: vec4f,
  @location(0) localPx: vec2f,
  @location(1) fill: vec4f,
  @location(2) sizeRadius: vec4f,
}

@vertex fn vs(input: VSIn) -> VSOut {
  var out: VSOut;
  let world = model.modelMatrix * vec4f(input.position, 1.0);
  out.pos = camera.projectionViewMatrix * world;
  out.localPx = (input.localUv - vec2f(0.5)) * input.sizeRadius.xy;
  out.fill = input.fill;
  out.sizeRadius = input.sizeRadius;
  return out;
}

@fragment fn fs(input: VSOut) -> @location(0) vec4f {
  let halfSize = 0.5 * input.sizeRadius.xy;
  let radius = min(input.sizeRadius.z, min(halfSize.x, halfSize.y));
  let q = abs(input.localPx) - (halfSize - vec2f(radius));
  let dist = length(max(q, vec2f(0.0))) + min(max(q.x, q.y), 0.0) - radius;
  let aa = max(fwidth(dist), 1.0e-4);
  let outerAlpha = 1.0 - smoothstep(-0.5 * aa, 0.5 * aa, dist);
  let alpha = outerAlpha * input.fill.a * model.color.a;
  if (alpha < 0.001) { discard; }
  return vec4f(input.fill.rgb * model.color.rgb, alpha);
}
`

const DECORATION_WGSL = /* wgsl */ `
struct Camera { projectionViewMatrix: mat4x4f }
struct Model { modelMatrix: mat4x4f, prevModelMatrix: mat4x4f, color: vec4f, flags: vec4f }

@group(0) @binding(0) var<uniform> camera: Camera;
@group(0) @binding(1) var<uniform> model: Model;

struct VSIn {
  @location(0) position: vec3f,
  @location(1) color: vec4f,
}

struct VSOut {
  @builtin(position) pos: vec4f,
  @location(0) color: vec4f,
}

@vertex fn vs(input: VSIn) -> VSOut {
  var out: VSOut;
  let world = model.modelMatrix * vec4f(input.position, 1.0);
  out.pos = camera.projectionViewMatrix * world;
  out.color = input.color;
  return out;
}

@fragment fn fs(input: VSOut) -> @location(0) vec4f {
  let alpha = input.color.a * model.color.a;
  if (alpha < 0.001) { discard; }
  return vec4f(input.color.rgb * model.color.rgb, alpha);
}
`

function buildQuadIndices(quadCount: number): Uint16Array | Uint32Array {
  const IndexArray = quadCount * 4 <= 65535 ? Uint16Array : Uint32Array
  const indices = new IndexArray(quadCount * 6)
  for (let i = 0; i < quadCount; i++) {
    const v = i * 4
    const at = i * 6
    indices[at] = v
    indices[at + 1] = v + 1
    indices[at + 2] = v + 2
    indices[at + 3] = v
    indices[at + 4] = v + 2
    indices[at + 5] = v + 3
  }
  return indices
}

function buildGlyphGeometry(
  glyphs: VanillaGlyphInstances,
  atlas: VanillaSdfAtlas,
  toneBoost: number,
  offsetX: number,
  offsetY: number,
): Geometry {
  const count = glyphs.count
  const positions = new Float32Array(count * 12)
  const uvs = new Float32Array(count * 8)
  const colors = new Float32Array(count * 16)
  const params = new Float32Array(count * 16)
  const unitX = atlas.pxRange / Math.max(1, atlas.width)
  const unitY = atlas.pxRange / Math.max(1, atlas.height)
  const invU16 = 1 / 65535
  const invU8 = 1 / 255

  for (let i = 0; i < count; i++) {
    const r4 = i * 4
    const x0 = glyphs.rects[r4]! + offsetX
    const y0 = glyphs.rects[r4 + 1]! + offsetY
    const x1 = glyphs.rects[r4 + 2]! + offsetX
    const y1 = glyphs.rects[r4 + 3]! + offsetY
    const z = glyphs.metadata[i * glyphs.metaItemSize + 1]!

    // uvRects are packed as unorm16 [u0, v1, u1, v0]
    const u0 = glyphs.uvRects[r4]! * invU16
    const v1 = glyphs.uvRects[r4 + 1]! * invU16
    const u1 = glyphs.uvRects[r4 + 2]! * invU16
    const v0 = glyphs.uvRects[r4 + 3]! * invU16

    const pBase = i * 12
    positions[pBase] = x0
    positions[pBase + 1] = y0
    positions[pBase + 2] = z
    positions[pBase + 3] = x1
    positions[pBase + 4] = y0
    positions[pBase + 5] = z
    positions[pBase + 6] = x1
    positions[pBase + 7] = y1
    positions[pBase + 8] = z
    positions[pBase + 9] = x0
    positions[pBase + 10] = y1
    positions[pBase + 11] = z

    const uvBase = i * 8
    uvs[uvBase] = u0
    uvs[uvBase + 1] = v1
    uvs[uvBase + 2] = u1
    uvs[uvBase + 3] = v1
    uvs[uvBase + 4] = u1
    uvs[uvBase + 5] = v0
    uvs[uvBase + 6] = u0
    uvs[uvBase + 7] = v0

    const cr = glyphs.colors[r4]! * invU8
    const cg = glyphs.colors[r4 + 1]! * invU8
    const cb = glyphs.colors[r4 + 2]! * invU8
    const ca = glyphs.colors[r4 + 3]! * invU8
    const cBase = i * 16
    for (let corner = 0; corner < 4; corner++) {
      const at = cBase + corner * 4
      colors[at] = cr
      colors[at + 1] = cg
      colors[at + 2] = cb
      colors[at + 3] = ca
      params[at] = unitX
      params[at + 1] = unitY
      params[at + 2] = toneBoost
      params[at + 3] = 0
    }
  }

  const geometry = new Geometry()
  geometry.setAttribute('position', positions, 3)
  geometry.setAttribute('uv', uvs, 2)
  geometry.setAttribute('color', colors, 4)
  geometry.setAttribute('params', params, 4)
  geometry.setIndex(buildQuadIndices(count))
  return geometry
}

function buildRectGeometry(
  rects: VanillaRectInstances,
  offsetX: number,
  offsetY: number,
): Geometry {
  const count = rects.count
  const positions = new Float32Array(count * 12)
  const localUvs = new Float32Array(count * 8)
  const fills = new Float32Array(count * 16)
  const sizeRadius = new Float32Array(count * 16)
  const invU8 = 1 / 255

  for (let i = 0; i < count; i++) {
    const r4 = i * 4
    const x0 = rects.rects[r4]! + offsetX
    const y0 = rects.rects[r4 + 1]! + offsetY
    const x1 = rects.rects[r4 + 2]! + offsetX
    const y1 = rects.rects[r4 + 3]! + offsetY
    const w = x1 - x0
    const h = y1 - y0
    const z = rects.depths[i]!

    const pBase = i * 12
    positions[pBase] = x0
    positions[pBase + 1] = y0
    positions[pBase + 2] = z
    positions[pBase + 3] = x1
    positions[pBase + 4] = y0
    positions[pBase + 5] = z
    positions[pBase + 6] = x1
    positions[pBase + 7] = y1
    positions[pBase + 8] = z
    positions[pBase + 9] = x0
    positions[pBase + 10] = y1
    positions[pBase + 11] = z

    const uvBase = i * 8
    localUvs[uvBase] = 0
    localUvs[uvBase + 1] = 0
    localUvs[uvBase + 2] = 1
    localUvs[uvBase + 3] = 0
    localUvs[uvBase + 4] = 1
    localUvs[uvBase + 5] = 1
    localUvs[uvBase + 6] = 0
    localUvs[uvBase + 7] = 1

    const cr = rects.colors[r4]! * invU8
    const cg = rects.colors[r4 + 1]! * invU8
    const cb = rects.colors[r4 + 2]! * invU8
    const ca = rects.colors[r4 + 3]! * invU8
    const radius = rects.radii[i]!
    const cBase = i * 16
    for (let corner = 0; corner < 4; corner++) {
      const at = cBase + corner * 4
      fills[at] = cr
      fills[at + 1] = cg
      fills[at + 2] = cb
      fills[at + 3] = ca

      sizeRadius[at] = w
      sizeRadius[at + 1] = h
      sizeRadius[at + 2] = radius
      sizeRadius[at + 3] = 0
    }
  }

  const geometry = new Geometry()
  geometry.setAttribute('position', positions, 3)
  geometry.setAttribute('localUv', localUvs, 2)
  geometry.setAttribute('fill', fills, 4)
  geometry.setAttribute('sizeRadius', sizeRadius, 4)
  geometry.setIndex(buildQuadIndices(count))
  return geometry
}

function buildDecorationGeometry(
  decorations: VanillaDecorationInstances,
  offsetX: number,
  offsetY: number,
): Geometry {
  const count = decorations.count
  const positions = new Float32Array(count * 12)
  const colors = new Float32Array(count * 16)
  const invU8 = 1 / 255

  for (let i = 0; i < count; i++) {
    const r4 = i * 4
    const x0 = decorations.decorationRects[r4]! + offsetX
    const y0 = decorations.decorationRects[r4 + 1]! + offsetY
    const x1 = decorations.decorationRects[r4 + 2]! + offsetX
    const y1 = decorations.decorationRects[r4 + 3]! + offsetY
    const z = decorations.metadata[i * 2 + 1]!

    const pBase = i * 12
    positions[pBase] = x0
    positions[pBase + 1] = y0
    positions[pBase + 2] = z
    positions[pBase + 3] = x1
    positions[pBase + 4] = y0
    positions[pBase + 5] = z
    positions[pBase + 6] = x1
    positions[pBase + 7] = y1
    positions[pBase + 8] = z
    positions[pBase + 9] = x0
    positions[pBase + 10] = y1
    positions[pBase + 11] = z

    const cr = decorations.colors[r4]! * invU8
    const cg = decorations.colors[r4 + 1]! * invU8
    const cb = decorations.colors[r4 + 2]! * invU8
    const ca = decorations.colors[r4 + 3]! * invU8
    const cBase = i * 16
    for (let corner = 0; corner < 4; corner++) {
      const at = cBase + corner * 4
      colors[at] = cr
      colors[at + 1] = cg
      colors[at + 2] = cb
      colors[at + 3] = ca
    }
  }

  const geometry = new Geometry()
  geometry.setAttribute('position', positions, 3)
  geometry.setAttribute('color', colors, 4)
  geometry.setIndex(buildQuadIndices(count))
  return geometry
}

export interface UltraTextOptions extends VanillaDomTextOptions {
  /** Center the captured layout around the local origin. Defaults to `true`. */
  center?: boolean
}

/**
 * Retained Mini3 scene node (`Object3D`) powered by `ultra-text/vanilla`.
 * Captures live HTML/CSS, builds a WebGPU JFA SDF atlas on `renderer.device`,
 * observes DOM mutations/resizes, and manages child Mini3 meshes and cleanup.
 */
export class UltraText extends Object3D {
  readonly domText: DomText
  readonly center: boolean

  private atlasTexture: Texture | null = null
  private currentAtlas: VanillaSdfAtlas | null = null
  private glyphMesh: Mesh | null = null
  private plainRectMesh: Mesh | null = null
  private litRectMesh: Mesh | null = null
  private decorationMesh: Mesh | null = null
  private unsubscribe: (() => void) | null = null
  private disposed = false

  protected constructor(domText: DomText, center: boolean) {
    super()
    this.domText = domText
    this.center = center
    this.syncMeshes()
    this.unsubscribe = domText.subscribe((_instance, reason) => {
      if (reason === 'dispose') {
        this.disposeMeshes()
        return
      }
      this.syncMeshes()
    })
  }

  static async create(
    renderer: Renderer,
    source: HTMLElement,
    options: UltraTextOptions = {},
  ): Promise<UltraText> {
    options.signal?.throwIfAborted()
    await document.fonts.ready
    options.signal?.throwIfAborted()

    const pixelsPerUnit = options.pixelsPerUnit ?? 100
    const { center = true, ...domOptions } = options
    const domText = await DomText.create(source, renderer, {
      pixelsPerUnit,
      textDepth: 1 / pixelsPerUnit,
      rectDepth: 0,
      panelDepth: -1 / pixelsPerUnit,
      ...domOptions,
      motion: { enabled: false, strength: 18 / pixelsPerUnit, ...domOptions.motion },
      onError: domOptions.onError ?? console.error,
    })

    if (options.signal?.aborted) {
      domText.dispose()
      options.signal.throwIfAborted()
    }

    return new UltraText(domText, center)
  }

  get width(): number {
    return this.domText.width
  }

  get height(): number {
    return this.domText.height
  }

  get stats() {
    return this.domText.stats
  }

  reflow(width: number): this {
    this.domText.reflow(width)
    return this
  }

  async rebuild(): Promise<this> {
    await this.domText.rebuild()
    return this
  }

  dispose(): void {
    if (this.disposed) return
    this.disposed = true
    this.unsubscribe?.()
    this.unsubscribe = null
    this.disposeMeshes()
    this.domText.dispose()
    this.parent?.remove(this)
  }

  private syncMeshes(): void {
    if (this.disposed || this.domText.isDisposed) return
    this.disposeMeshes()

    const offsetX = this.center ? -this.domText.width * 0.5 : 0
    const offsetY = this.center ? -this.domText.height * 0.5 : 0
    const atlas = this.domText.atlas

    this.currentAtlas = atlas
    this.atlasTexture = new Texture({
      gpuTexture: atlas.gpuTexture ?? null,
      width: atlas.width,
      height: atlas.height,
      format: 'r8unorm',
      wrapS: 'clamp',
      wrapT: 'clamp',
      minFilter: 'linear',
      magFilter: 'linear',
      flipY: false,
      mips: false,
    })

    if (this.domText.plainRects && this.domText.plainRects.count > 0) {
      const geo = buildRectGeometry(this.domText.plainRects, offsetX, offsetY)
      const mat = new Material({
        shader: RECT_WGSL,
        side: 'both',
        transparent: true,
        depthWrite: false,
        unlit: true,
      })
      this.plainRectMesh = new Mesh(geo, mat)
      this.plainRectMesh.renderOrder = 0
      this.add(this.plainRectMesh)
    }

    if (this.domText.litRects && this.domText.litRects.count > 0) {
      const geo = buildRectGeometry(this.domText.litRects, offsetX, offsetY)
      const mat = new Material({
        shader: RECT_WGSL,
        side: 'both',
        transparent: true,
        depthWrite: false,
        unlit: true,
      })
      this.litRectMesh = new Mesh(geo, mat)
      this.litRectMesh.receiveShadow = true
      this.litRectMesh.renderOrder = 1
      this.add(this.litRectMesh)
    }

    if (this.domText.glyphs.count > 0) {
      const toneBoost = this.domText.tone ? this.domText.tone.strength * 0.14 : 0
      const geo = buildGlyphGeometry(this.domText.glyphs, atlas, toneBoost, offsetX, offsetY)
      const mat = new Material({
        texture: this.atlasTexture,
        shader: GLYPH_WGSL,
        side: 'both',
        transparent: true,
        depthWrite: false,
        unlit: true,
      })
      this.glyphMesh = new Mesh(geo, mat)
      this.glyphMesh.renderOrder = 2
      this.add(this.glyphMesh)
    }

    if (this.domText.decorations && this.domText.decorations.count > 0) {
      const geo = buildDecorationGeometry(this.domText.decorations, offsetX, offsetY)
      const mat = new Material({
        shader: DECORATION_WGSL,
        side: 'both',
        transparent: true,
        depthWrite: false,
        unlit: true,
      })
      this.decorationMesh = new Mesh(geo, mat)
      this.decorationMesh.renderOrder = 3
      this.add(this.decorationMesh)
    }
  }

  private disposeMeshes(): void {
    for (const mesh of [
      this.glyphMesh,
      this.plainRectMesh,
      this.litRectMesh,
      this.decorationMesh,
    ]) {
      if (!mesh) continue
      this.remove(mesh)
      mesh.geometry?.dispose()
      mesh.material?.dispose()
      mesh.dispose()
    }
    this.glyphMesh = null
    this.plainRectMesh = null
    this.litRectMesh = null
    this.decorationMesh = null
    if (this.atlasTexture) {
      // DomText owns the underlying GPUTexture lifecycle; detach before Texture.dispose().
      this.atlasTexture.gpuTexture = null
      this.atlasTexture.dispose()
      this.atlasTexture = null
    }
    this.currentAtlas = null
  }
}