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

export class XRPass extends Pass {
  render(
    _renderer: THREE.WebGLRenderer,
    _writeBuffer: THREE.WebGLRenderTarget,
    _readBuffer: THREE.WebGLRenderTarget,
    _deltaTime: number,
    _maskActive: boolean,
    _viewId: number = 0
  ) {}
}

/**
 * XREffects manages the XR rendering pipeline.
 * Use core.effects
 * It handles multiple passes and render targets for applying effects to XR
 * scenes.
 */
export class XREffects {
  passes: XRPass[] = [];
  renderTargets: THREE.WebGLRenderTarget[] = [];
  dimensions = new THREE.Vector2();

  constructor(
    private renderer: THREE.WebGLRenderer,
    private scene: THREE.Scene,
    private timer: THREE.Timer
  ) {}

  /**
   * Adds a pass to the effect pipeline.
   */
  addPass(pass: XRPass) {
    pass.renderToScreen = false;
    this.passes.push(pass);
  }

  /**
   * Sets up render targets for the effect pipeline.
   */
  setupRenderTargets(dimensions: THREE.Vector2) {
    const defaultTarget = this.renderer.getRenderTarget();
    if (defaultTarget == null) {
      return;
    }
    const neededRenderTargets = this.renderer.xr.isPresenting ? 4 : 2;
    for (let i = 0; i < neededRenderTargets; i++) {
      if (
        i >= this.renderTargets.length ||
        this.renderTargets[i].width != dimensions.x ||
        this.renderTargets[i].height != dimensions.y
      ) {
        this.renderTargets[i]?.depthTexture?.dispose();
        this.renderTargets[i]?.dispose();
        this.renderTargets[i] = defaultTarget.clone();
        this.renderTargets[i].depthTexture = new THREE.DepthTexture(
          dimensions.x,
          dimensions.y
        );
      }
    }
    for (let i = neededRenderTargets; i < this.renderTargets.length; i++) {
      this.renderTargets[i].depthTexture?.dispose();
      this.renderTargets[i].dispose();
    }
  }

  /**
   * Renders the XR effects.
   */
  render(camera: THREE.Camera) {
    this.renderer.getDrawingBufferSize(this.dimensions);
    this.setupRenderTargets(this.dimensions);
    this.renderer.xr.cameraAutoUpdate = false;
    const defaultTarget = this.renderer.getRenderTarget();
    if (!defaultTarget) {
      return;
    }
    if (this.renderer.xr.isPresenting) {
      this.renderXr();
    } else {
      this.renderSimulator(camera);
    }
  }

  private renderXr() {
    const defaultTarget = this.renderer.getRenderTarget()!;
    const renderer = this.renderer;
    const xrEnabled = renderer.xr.enabled;
    const xrIsPresenting = renderer.xr.isPresenting;
    const renderTargets = this.renderTargets;
    renderer.xr.cameraAutoUpdate = false;
    renderer.xr.enabled = false;
    const deltaTime = this.timer.getDelta();
    const numCameras = renderer.xr.getCamera().cameras.length;
    if (numCameras > 0) {
      for (let camIndex = 0; camIndex < numCameras; ++camIndex) {
        const cam = renderer.xr.getCamera().cameras[camIndex];
        renderer.setViewport(cam.viewport);
        renderer.setRenderTarget(renderTargets[camIndex]);
        renderer.clear();
        renderer.xr.isPresenting = true;
        renderer.render(this.scene, cam);
      }
      renderer.setRenderTarget(defaultTarget);
      renderer.clear();
      renderer.xr.isPresenting = false;
      renderer.autoClearColor = false;
      for (let eye = 0; eye < numCameras; eye++) {
        for (let i = 0; i < this.passes.length - 1; ++i) {
          const lastRenderTargetIndex = i % 2;
          const nextRenderTargetIndex = (i + 1) % 2;
          defaultTarget.viewport.set(
            (eye * this.dimensions.x) / numCameras,
            0,
            this.dimensions.x / numCameras,
            this.dimensions.y
          );
          this.passes[i].render(
            renderer,
            this.renderTargets[2 * nextRenderTargetIndex + eye],
            this.renderTargets[2 * lastRenderTargetIndex + eye],
            deltaTime,
            /*maskActive=*/ false,
            /*viewId=*/ eye
          );
        }
        if (this.passes.length > 0) {
          const lastRenderTargetIndex = (this.passes.length - 1) % 2;
          defaultTarget.viewport.set(
            (eye * this.dimensions.x) / numCameras,
            0,
            this.dimensions.x / numCameras,
            this.dimensions.y
          );
          this.passes[this.passes.length - 1].render(
            renderer,
            defaultTarget,
            this.renderTargets[2 * lastRenderTargetIndex + eye],
            deltaTime,
            /*maskActive=*/ false,
            /*viewId=*/ eye
          );
        }
      }
      renderer.xr.enabled = xrEnabled;
      renderer.xr.isPresenting = xrIsPresenting;
    }
  }

  private renderSimulator(camera: THREE.Camera) {
    const defaultTarget = this.renderer.getRenderTarget()!;
    const renderer = this.renderer;
    const xrEnabled = renderer.xr.enabled;
    const xrIsPresenting = renderer.xr.isPresenting;
    renderer.xr.cameraAutoUpdate = false;
    renderer.xr.enabled = false;
    const deltaTime = this.timer.getDelta();
    if (this.passes.length === 0) {
      renderer.setRenderTarget(defaultTarget);
      renderer.render(this.scene, camera);
      renderer.xr.enabled = xrEnabled;
      renderer.xr.isPresenting = xrIsPresenting;
      return;
    }
    renderer.setRenderTarget(this.renderTargets[0]);
    renderer.clear();
    renderer.render(this.scene, camera);
    renderer.setRenderTarget(defaultTarget);
    renderer.clear();
    renderer.xr.isPresenting = false;
    renderer.autoClearColor = false;
    for (let i = 0; i < this.passes.length - 1; ++i) {
      const lastRenderTargetIndex = i % 2;
      const nextRenderTargetIndex = (i + 1) % 2;
      this.passes[i].render(
        renderer,
        this.renderTargets[nextRenderTargetIndex],
        this.renderTargets[lastRenderTargetIndex],
        deltaTime,
        /*maskActive=*/ false,
        /*viewId=*/ 0
      );
    }
    if (this.passes.length > 0) {
      const lastRenderTargetIndex = (this.passes.length - 1) % 2;
      this.passes[this.passes.length - 1].render(
        renderer,
        defaultTarget,
        this.renderTargets[lastRenderTargetIndex],
        deltaTime,
        /*maskActive=*/ false,
        /*viewId=*/ 0
      );
    }
    renderer.xr.enabled = xrEnabled;
    renderer.xr.isPresenting = xrIsPresenting;
  }

  dispose() {
    let firstError: unknown;
    for (const target of this.renderTargets) {
      for (const dispose of [
        () => target.depthTexture?.dispose(),
        () => target.dispose(),
      ]) {
        try {
          dispose();
        } catch (error: unknown) {
          firstError ??= error;
        }
      }
    }
    this.renderTargets.length = 0;
    for (const pass of this.passes) {
      try {
        pass.dispose();
      } catch (error: unknown) {
        firstError ??= error;
      }
    }
    this.passes.length = 0;
    if (firstError !== undefined) throw firstError;
  }
}