import * as THREE from 'three';
import {FullScreenQuad} from 'three/addons/postprocessing/Pass.js';

import {XRDeviceCamera} from '../../camera/XRDeviceCamera.js';

// Use a small canvas since a full size canvas can consume a lot of memory and
// cause toDataUrl to be slow.
const DEFAULT_CANVAS_WIDTH = 640;

function flipBufferVertically(
  buffer: Uint8Array,
  width: number,
  height: number
) {
  const bytesPerRow = width * 4;
  const tempRow = new Uint8Array(bytesPerRow);
  for (let y = 0; y < height / 2; y++) {
    const topRowY = y;
    const bottomRowY = height - 1 - y;
    const topRowOffset = topRowY * bytesPerRow;
    const bottomRowOffset = bottomRowY * bytesPerRow;
    tempRow.set(buffer.subarray(topRowOffset, topRowOffset + bytesPerRow));
    buffer.set(
      buffer.subarray(bottomRowOffset, bottomRowOffset + bytesPerRow),
      topRowOffset
    );
    buffer.set(tempRow, bottomRowOffset);
  }
}

class PendingScreenshotRequest {
  constructor(
    public resolve: (value: string) => void,
    public reject: (reason?: Error) => void,
    public overlayOnCamera: boolean
  ) {}
}

export class ScreenshotSynthesizer {
  private pendingScreenshotRequests: PendingScreenshotRequest[] = [];
  private virtualCanvas?: HTMLCanvasElement;
  private virtualBuffer = new Uint8Array();
  // Smaller resolution render target than the main render target.
  private virtualRenderTarget?: THREE.WebGLRenderTarget;
  private virtualRealCanvas?: HTMLCanvasElement;
  private virtualRealBuffer = new Uint8Array();
  private virtualRealRenderTarget?: THREE.WebGLRenderTarget;
  private fullScreenQuad?: FullScreenQuad;
  private renderTargetWidth = DEFAULT_CANVAS_WIDTH;

  async onAfterRender(
    renderer: THREE.WebGLRenderer,
    renderSceneFn: () => void,
    deviceCamera?: XRDeviceCamera
  ) {
    if (this.pendingScreenshotRequests.length == 0) {
      return;
    }

    const haveVirtualOnlyRequests = this.pendingScreenshotRequests.every(
      (request) => !request.overlayOnCamera
    );
    if (haveVirtualOnlyRequests) {
      this.createVirtualImageDataURL(renderer, renderSceneFn).then(
        (virtualImageDataUrl) => {
          this.resolveVirtualOnlyRequests(virtualImageDataUrl);
        }
      );
    }

    const haveVirtualAndRealReqeusts = this.pendingScreenshotRequests.some(
      (request) => request.overlayOnCamera
    );
    if (haveVirtualAndRealReqeusts && deviceCamera) {
      this.createVirtualRealImageDataURL(
        renderer,
        renderSceneFn,
        deviceCamera
      ).then((virtualRealImageDataUrl) => {
        if (virtualRealImageDataUrl) {
          this.resolveVirtualRealRequests(virtualRealImageDataUrl);
        }
      });
    } else if (haveVirtualAndRealReqeusts) {
      throw new Error('No device camera provided');
    }
  }

  private async createVirtualImageDataURL(
    renderer: THREE.WebGLRenderer,
    renderSceneFn: () => void
  ) {
    const mainRenderTarget = renderer.getRenderTarget()!;
    const isRenderingStereo =
      renderer.xr.isPresenting && renderer.xr.getCamera().cameras.length == 2;
    const mainRenderTargetSingleViewWidth = isRenderingStereo
      ? mainRenderTarget.width / 2
      : mainRenderTarget.width;
    const scaledHeight = Math.round(
      mainRenderTarget.height *
        (this.renderTargetWidth / mainRenderTargetSingleViewWidth)
    );
    if (
      !this.virtualRenderTarget ||
      this.virtualRenderTarget.width != this.renderTargetWidth
    ) {
      this.virtualRenderTarget?.dispose();
      this.virtualRenderTarget = new THREE.WebGLRenderTarget(
        this.renderTargetWidth,
        scaledHeight,
        {colorSpace: THREE.SRGBColorSpace}
      );
    }
    const xrIsPresenting = renderer.xr.isPresenting;
    renderer.xr.isPresenting = false;
    const virtualRenderTarget = this.virtualRenderTarget;
    renderer.setRenderTarget(virtualRenderTarget);
    renderer.clearColor();
    renderer.clearDepth();
    renderSceneFn();
    renderer.setRenderTarget(mainRenderTarget);
    renderer.xr.isPresenting = xrIsPresenting;

    const expectedBufferLength =
      virtualRenderTarget.width * virtualRenderTarget.height * 4;
    if (this.virtualBuffer.length != expectedBufferLength) {
      this.virtualBuffer = new Uint8Array(expectedBufferLength);
    }
    const buffer = this.virtualBuffer;
    await renderer.readRenderTargetPixelsAsync(
      virtualRenderTarget,
      0,
      0,
      virtualRenderTarget.width,
      virtualRenderTarget.height,
      buffer
    );

    flipBufferVertically(
      buffer,
      virtualRenderTarget.width,
      virtualRenderTarget.height
    );
    const canvas =
      this.virtualCanvas ||
      (this.virtualCanvas = document.createElement('canvas'));
    canvas.width = virtualRenderTarget.width;
    canvas.height = virtualRenderTarget.height;
    const context = canvas.getContext('2d');
    if (!context) {
      throw new Error('Failed to get 2D context');
    }
    const imageData = new ImageData(
      new Uint8ClampedArray(buffer),
      virtualRenderTarget.width,
      virtualRenderTarget.height
    );
    context.putImageData(imageData, 0, 0);
    return canvas.toDataURL();
  }

  private resolveVirtualOnlyRequests(virtualImageDataUrl: string) {
    let remainingRequests = 0;
    for (let i = 0; i < this.pendingScreenshotRequests.length; i++) {
      const request = this.pendingScreenshotRequests[i];
      if (!request.overlayOnCamera) {
        request.resolve(virtualImageDataUrl);
      } else {
        this.pendingScreenshotRequests[remainingRequests++] = request;
      }
    }
    this.pendingScreenshotRequests.length = remainingRequests;
  }

  private async createVirtualRealImageDataURL(
    renderer: THREE.WebGLRenderer,
    renderSceneFn: () => void,
    deviceCamera: XRDeviceCamera
  ) {
    if (!deviceCamera.loaded) {
      console.debug('Waiting for device camera to be loaded');
      return null;
    }
    const mainRenderTarget = renderer.getRenderTarget();
    const isRenderingStereo =
      renderer.xr.isPresenting && renderer.xr.getCamera().cameras.length == 2;
    const mainRenderTargetSize = new THREE.Vector2();
    if (mainRenderTarget) {
      mainRenderTargetSize.set(mainRenderTarget.width, mainRenderTarget.height);
    } else {
      renderer.getSize(mainRenderTargetSize);
    }
    const mainRenderTargetSingleViewWidth = isRenderingStereo
      ? mainRenderTargetSize.x / 2
      : mainRenderTargetSize.y;
    const scaledHeight = Math.round(
      mainRenderTargetSize.y *
        (this.renderTargetWidth / mainRenderTargetSingleViewWidth)
    );
    if (
      !this.virtualRealRenderTarget ||
      this.virtualRealRenderTarget.height != scaledHeight
    ) {
      this.virtualRealRenderTarget?.dispose();
      this.virtualRealRenderTarget = new THREE.WebGLRenderTarget(
        this.renderTargetWidth,
        scaledHeight,
        {colorSpace: THREE.SRGBColorSpace}
      );
    }

    const renderTarget = this.virtualRealRenderTarget;
    renderer.setRenderTarget(renderTarget);
    const xrIsPresenting = renderer.xr.isPresenting;
    renderer.xr.isPresenting = false;
    const quad = this.getFullScreenQuad();
    (quad.material as THREE.MeshBasicMaterial).map = deviceCamera.texture;
    quad.render(renderer);
    renderSceneFn();
    renderer.xr.isPresenting = xrIsPresenting;
    renderer.setRenderTarget(mainRenderTarget);

    if (
      this.virtualRealBuffer.length !=
      renderTarget.width * renderTarget.height * 4
    ) {
      this.virtualRealBuffer = new Uint8Array(
        renderTarget.width * renderTarget.height * 4
      );
    }
    const buffer = this.virtualRealBuffer;
    await renderer.readRenderTargetPixelsAsync(
      renderTarget,
      0,
      0,
      renderTarget.width,
      renderTarget.height,
      buffer
    );

    flipBufferVertically(buffer, renderTarget.width, renderTarget.height);
    const canvas =
      this.virtualRealCanvas ||
      (this.virtualRealCanvas = document.createElement('canvas'));
    canvas.width = renderTarget.width;
    canvas.height = renderTarget.height;
    const context = canvas.getContext('2d');
    if (!context) {
      throw new Error('Failed to get 2D context');
    }
    const imageData = new ImageData(
      new Uint8ClampedArray(buffer),
      renderTarget.width,
      renderTarget.height
    );
    context.putImageData(imageData, 0, 0);
    return canvas.toDataURL();
  }

  private resolveVirtualRealRequests(virtualRealImageDataUrl: string) {
    let remainingRequests = 0;
    for (let i = 0; i < this.pendingScreenshotRequests.length; i++) {
      const request = this.pendingScreenshotRequests[i];
      if (request.overlayOnCamera) {
        request.resolve(virtualRealImageDataUrl);
      } else {
        this.pendingScreenshotRequests[remainingRequests++] = request;
      }
    }
    this.pendingScreenshotRequests.length = remainingRequests;
  }

  private getFullScreenQuad() {
    if (!this.fullScreenQuad) {
      this.fullScreenQuad = new FullScreenQuad(
        new THREE.MeshBasicMaterial({transparent: true})
      );
    }
    return this.fullScreenQuad;
  }

  /**
   * Requests a screenshot from the scene as a DataURL.
   * @param overlayOnCamera - If true, overlays the image on a camera image
   *     without any projection or aspect ratio correction.
   * @returns Promise which returns the screenshot as a data uri.
   */
  async getScreenshot(overlayOnCamera = false): Promise<string> {
    return new Promise<string>((resolve, reject) => {
      this.pendingScreenshotRequests.push(
        new PendingScreenshotRequest(resolve, reject, overlayOnCamera)
      );
    });
  }
}