import regl from 'regl'
import {Flamechart} from './flamechart'
import {RectangleBatch} from './rectangle-batch-renderer'
import {CanvasContext} from './canvas-context'
import {Vec2, Rect, AffineTransform} from './math'
import {LRUCache} from './lru-cache'
import {Color} from './color'
import {getOrInsert} from './utils'

const MAX_BATCH_SIZE = 10000

class RowAtlas<K> {
  texture: regl.Texture
  private framebuffer: regl.Framebuffer
  private renderToFramebuffer: regl.Command<{}>
  private rowCache: LRUCache<K, number>
  private clearLineBatch: RectangleBatch

  constructor(private canvasContext: CanvasContext) {
    this.texture = canvasContext.gl.texture({
      width: Math.min(canvasContext.getMaxTextureSize(), 4096),
      height: Math.min(canvasContext.getMaxTextureSize(), 1024),
      wrapS: 'clamp',
      wrapT: 'clamp',
    })
    this.framebuffer = canvasContext.gl.framebuffer({color: [this.texture]})
    this.rowCache = new LRUCache(this.texture.height)
    this.renderToFramebuffer = canvasContext.gl({
      framebuffer: this.framebuffer,
    })
    this.clearLineBatch = canvasContext.createRectangleBatch()
    this.clearLineBatch.addRect(Rect.unit, new Color(0, 0, 0, 0))
  }

  has(key: K) {
    return this.rowCache.has(key)
  }
  getResolution() {
    return this.texture.width
  }
  getCapacity() {
    return this.texture.height
  }

  private allocateLine(key: K): number {
    if (this.rowCache.getSize() < this.rowCache.getCapacity()) {
      // Not in cache, but cache isn't full
      const row = this.rowCache.getSize()
      this.rowCache.insert(key, row)
      return row
    } else {
      // Not in cache, and cache is full. Evict something.
      const [, row] = this.rowCache.removeLRU()!
      this.rowCache.insert(key, row)
      return row
    }
  }

  writeToAtlasIfNeeded(keys: K[], render: (textureDstRect: Rect, key: K) => void) {
    this.renderToFramebuffer((context: regl.Context) => {
      for (let key of keys) {
        let row = this.rowCache.get(key)
        if (row != null) {
          // Already cached!
          continue
        }
        // Not cached -- we'll have to actually render
        row = this.allocateLine(key)

        const textureRect = new Rect(new Vec2(0, row), new Vec2(this.texture.width, 1))
        this.canvasContext.drawRectangleBatch({
          batch: this.clearLineBatch,
          configSpaceSrcRect: Rect.unit,
          physicalSpaceDstRect: textureRect,
        })
        render(textureRect, key)
      }
    })
  }

  renderViaAtlas(key: K, dstRect: Rect): boolean {
    let row = this.rowCache.get(key)
    if (row == null) {
      return false
    }

    const textureRect = new Rect(new Vec2(0, row), new Vec2(this.texture.width, 1))

    // At this point, we have the row in cache, and we can
    // paint directly from it into the framebuffer.
    this.canvasContext.drawTexture({
      texture: this.texture,
      srcRect: textureRect,
      dstRect: dstRect,
    })
    return true
  }
}

interface RangeTreeNode {
  getBounds(): Rect
  getRectCount(): number
  getChildren(): RangeTreeNode[]
  forEachLeafNodeWithinBounds(configSpaceBounds: Rect, cb: (leaf: RangeTreeLeafNode) => void): void
}

class RangeTreeLeafNode implements RangeTreeNode {
  private children: RangeTreeNode[] = []

  constructor(
    private batch: RectangleBatch,
    private bounds: Rect,
    private numPrecedingRectanglesInRow: number,
  ) {
    batch.uploadToGPU()
  }

  getBatch() {
    return this.batch
  }
  getBounds() {
    return this.bounds
  }
  getRectCount() {
    return this.batch.getRectCount()
  }
  getChildren() {
    return this.children
  }
  getParity() {
    return this.numPrecedingRectanglesInRow % 2
  }
  forEachLeafNodeWithinBounds(configSpaceBounds: Rect, cb: (leaf: RangeTreeLeafNode) => void) {
    if (!this.bounds.hasIntersectionWith(configSpaceBounds)) return
    cb(this)
  }
}

class RangeTreeInteriorNode implements RangeTreeNode {
  private rectCount: number = 0
  private bounds: Rect
  constructor(private children: RangeTreeNode[]) {
    if (children.length === 0) {
      throw new Error('Empty interior node')
    }
    let minLeft = Infinity
    let maxRight = -Infinity
    let minTop = Infinity
    let maxBottom = -Infinity
    for (let child of children) {
      this.rectCount += child.getRectCount()
      const bounds = child.getBounds()
      minLeft = Math.min(minLeft, bounds.left())
      maxRight = Math.max(maxRight, bounds.right())
      minTop = Math.min(minTop, bounds.top())
      maxBottom = Math.max(maxBottom, bounds.bottom())
    }
    this.bounds = new Rect(
      new Vec2(minLeft, minTop),
      new Vec2(maxRight - minLeft, maxBottom - minTop),
    )
  }

  getBounds() {
    return this.bounds
  }
  getRectCount() {
    return this.rectCount
  }
  getChildren() {
    return this.children
  }

  forEachLeafNodeWithinBounds(configSpaceBounds: Rect, cb: (leaf: RangeTreeLeafNode) => void) {
    if (!this.bounds.hasIntersectionWith(configSpaceBounds)) return
    for (let child of this.children) {
      child.forEachLeafNodeWithinBounds(configSpaceBounds, cb)
    }
  }
}

export interface FlamechartRendererProps {
  configSpaceSrcRect: Rect
  physicalSpaceDstRect: Rect
  renderOutlines: boolean
}

interface FlamechartRowAtlasKey {
  stackDepth: number
  zoomLevel: number
  index: number
}

export class FlamechartRenderer {
  private layers: RangeTreeNode[] = []
  private rowAtlas: RowAtlas<FlamechartRowAtlasKey>
  private rectInfoTexture: regl.Texture
  private framebuffer: regl.Framebuffer

  constructor(private canvasContext: CanvasContext, private flamechart: Flamechart) {
    const nLayers = flamechart.getLayers().length
    this.rowAtlas = new RowAtlas(canvasContext)
    for (let stackDepth = 0; stackDepth < nLayers; stackDepth++) {
      const leafNodes: RangeTreeLeafNode[] = []
      const y = stackDepth

      let minLeft = Infinity
      let maxRight = -Infinity
      let batch = canvasContext.createRectangleBatch()

      let rectCount = 0

      const layer = flamechart.getLayers()[stackDepth]

      for (let i = 0; i < layer.length; i++) {
        const frame = layer[i]
        if (batch.getRectCount() >= MAX_BATCH_SIZE) {
          leafNodes.push(
            new RangeTreeLeafNode(
              batch,
              new Rect(new Vec2(minLeft, stackDepth), new Vec2(maxRight - minLeft, 1)),
              rectCount,
            ),
          )
          minLeft = Infinity
          maxRight = -Infinity
          batch = canvasContext.createRectangleBatch()
        }
        const configSpaceBounds = new Rect(
          new Vec2(frame.start, y),
          new Vec2(frame.end - frame.start, 1),
        )
        minLeft = Math.min(minLeft, configSpaceBounds.left())
        maxRight = Math.max(maxRight, configSpaceBounds.right())

        // We'll use the red channel to indicate the index to allow
        // us to separate adjacent rectangles within a row from one another,
        // the green channel to indicate the row,
        // and the blue channel to indicate the color bucket to render.
        // We add one to each so we have zero reserved for the background color.
        const color = new Color(
          (1 + i % 255) / 256,
          (1 + stackDepth % 255) / 256,
          (1 + this.flamechart.getColorBucketForFrame(frame.node.frame)) / 256,
        )
        batch.addRect(configSpaceBounds, color)
        rectCount++
      }

      if (batch.getRectCount() > 0) {
        leafNodes.push(
          new RangeTreeLeafNode(
            batch,
            new Rect(new Vec2(minLeft, stackDepth), new Vec2(maxRight - minLeft, 1)),
            rectCount,
          ),
        )
      }

      // TODO(jlfwong): Making this into a binary tree
      // range than a tree of always-height-two might make this run faster
      this.layers.push(new RangeTreeInteriorNode(leafNodes))
    }
    this.rectInfoTexture = this.canvasContext.gl.texture({width: 1, height: 1})
    this.framebuffer = this.canvasContext.gl.framebuffer({
      color: [this.rectInfoTexture],
    })
  }

  private atlasKeys = new Map<string, FlamechartRowAtlasKey>()
  getOrInsertKey(key: FlamechartRowAtlasKey): FlamechartRowAtlasKey {
    const hash = `${key.stackDepth}_${key.index}_${key.zoomLevel}`
    return getOrInsert(this.atlasKeys, hash, () => key)
  }

  configSpaceBoundsForKey(key: FlamechartRowAtlasKey): Rect {
    const {stackDepth, zoomLevel, index} = key
    const configSpaceContentWidth = this.flamechart.getTotalWeight()

    const width = configSpaceContentWidth / Math.pow(2, zoomLevel)

    return new Rect(new Vec2(width * index, stackDepth), new Vec2(width, 1))
  }

  render(props: FlamechartRendererProps) {
    const {configSpaceSrcRect, physicalSpaceDstRect} = props

    const atlasKeysToRender: {stackDepth: number; zoomLevel: number; index: number}[] = []

    // We want to render the lowest resolution we can while still guaranteeing that the
    // atlas line is higher resolution than its corresponding destination rectangle on
    // the screen.
    const configToPhysical = AffineTransform.betweenRects(configSpaceSrcRect, physicalSpaceDstRect)
    if (configSpaceSrcRect.isEmpty()) {
      // Prevent an infinite loop
      return
    }

    let zoomLevel = 0
    while (true) {
      const configSpaceBounds = this.configSpaceBoundsForKey({stackDepth: 0, zoomLevel, index: 0})
      const physicalBounds = configToPhysical.transformRect(configSpaceBounds)
      if (physicalBounds.width() < this.rowAtlas.getResolution()) {
        break
      }
      zoomLevel++
    }

    const top = Math.max(0, Math.floor(configSpaceSrcRect.top()))
    const bottom = Math.min(this.layers.length, Math.ceil(configSpaceSrcRect.bottom()))

    const configSpaceContentWidth = this.flamechart.getTotalWeight()
    const numAtlasEntriesPerLayer = Math.pow(2, zoomLevel)
    const left = Math.floor(
      numAtlasEntriesPerLayer * configSpaceSrcRect.left() / configSpaceContentWidth,
    )
    const right = Math.ceil(
      numAtlasEntriesPerLayer * configSpaceSrcRect.right() / configSpaceContentWidth,
    )

    for (let stackDepth = top; stackDepth < bottom; stackDepth++) {
      for (let index = left; index <= right; index++) {
        const key = this.getOrInsertKey({stackDepth, zoomLevel, index})
        const configSpaceBounds = this.configSpaceBoundsForKey(key)
        if (!configSpaceBounds.hasIntersectionWith(configSpaceSrcRect)) continue
        atlasKeysToRender.push(key)
      }
    }

    const cacheCapacity = this.rowAtlas.getCapacity()
    const keysToRenderCached = atlasKeysToRender.slice(0, cacheCapacity)
    const keysToRenderUncached = atlasKeysToRender.slice(cacheCapacity)

    // Fill the cache
    this.rowAtlas.writeToAtlasIfNeeded(keysToRenderCached, (textureDstRect, key) => {
      const configSpaceBounds = this.configSpaceBoundsForKey(key)
      this.layers[key.stackDepth].forEachLeafNodeWithinBounds(configSpaceBounds, leaf => {
        this.canvasContext.drawRectangleBatch({
          batch: leaf.getBatch(),
          configSpaceSrcRect: configSpaceBounds,
          physicalSpaceDstRect: textureDstRect,
          parityMin: key.stackDepth % 2 == 0 ? 2 : 0,
          parityOffset: leaf.getParity(),
        })
      })
    })

    this.framebuffer.resize(physicalSpaceDstRect.width(), physicalSpaceDstRect.height())
    this.framebuffer.use(context => {
      this.canvasContext.gl.clear({color: [0, 0, 0, 0]})
      const viewportRect = new Rect(
        Vec2.zero,
        new Vec2(context.viewportWidth, context.viewportHeight),
      )

      const configToViewport = AffineTransform.betweenRects(configSpaceSrcRect, viewportRect)

      // Render from the cache
      for (let key of keysToRenderCached) {
        const configSpaceSrcRect = this.configSpaceBoundsForKey(key)
        this.rowAtlas.renderViaAtlas(key, configToViewport.transformRect(configSpaceSrcRect))
      }

      // Render entries that didn't make it into the cache
      for (let key of keysToRenderUncached) {
        const configSpaceBounds = this.configSpaceBoundsForKey(key)
        const physicalBounds = configToViewport.transformRect(configSpaceBounds)
        this.layers[key.stackDepth].forEachLeafNodeWithinBounds(configSpaceBounds, leaf => {
          this.canvasContext.drawRectangleBatch({
            batch: leaf.getBatch(),
            configSpaceSrcRect,
            physicalSpaceDstRect: physicalBounds,
            parityMin: key.stackDepth % 2 == 0 ? 2 : 0,
            parityOffset: leaf.getParity(),
          })
        })
      }
    })

    this.canvasContext.drawFlamechartColorPass({
      rectInfoTexture: this.rectInfoTexture,
      srcRect: new Rect(
        Vec2.zero,
        new Vec2(this.rectInfoTexture.width, this.rectInfoTexture.height),
      ),
      dstRect: physicalSpaceDstRect,
      renderOutlines: props.renderOutlines,
    })

    // Overlay the atlas on top of the canvas for debugging
    /*
    this.canvasContext.drawTexture({
      texture: this.rowAtlas.texture,
      srcRect: new Rect(Vec2.zero, new Vec2(this.rowAtlas.texture.width, this.rowAtlas.texture.height)),
      dstRect: new Rect(Vec2.zero, new Vec2(800, 800))
    })
    */
  }
}
