import {
  CompareFunction,
  CullMode,
  Material,
  Matrix,
  MeshRenderer,
  MeshTopology,
  ModelMesh,
  RenderQueueType,
  Script,
  Shader,
  Texture2D,
  TextureFormat,
  Vector2,
  Vector3,
} from "oasis-engine";

/**
 * 使用时再重写，性能可以得到提升
 */
export class ARBackground extends Script {
  private projectMatrix: Matrix = new Matrix();
  private frameWidth: number = 0;
  private frameHeight: number = 0;
  private yTexture: Texture2D;
  private uvTexture: Texture2D;
  private material: Material;

  /**
   * 更新相机纹理和顶点调整矩阵
   * @param frameWidth
   * @param frameHeight
   * @param curARFrame
   */
  onARUpdate(
    frameWidth: number,
    frameHeight: number,
    curARFrame: WechatMiniprogram.VKFrame
  ): void {
    const { engine } = this;
    const { gl, _activeTextures: activeTextures } = engine._hardwareRenderer;
    // 同步三维空间的背景
    const cameraTexture = curARFrame.getCameraTexture(gl);
    // 问题原因：微信在返回 WebglTexture 的时候修改了底层纹理插槽的绑定关系
    // 而此时 Oasis 的 cpu 层并没有记录这个操作，导致缓存失效，IOS 可能不能
    // 正常渲染背景。
    // 解决办法：在每次微信返回 YUV 纹理时， refresh 下 cpu 层维护的记录。
    activeTextures.length = 0;
    // 调整这个 Plane 的宽高
    if (cameraTexture) {
      if (frameWidth !== this.frameWidth || frameHeight !== this.frameHeight) {
        this.frameWidth = frameWidth;
        this.frameHeight = frameHeight;
        this.yTexture && this.yTexture.destroy();
        this.uvTexture && this.uvTexture.destroy();
        // prettier-ignore
        this.yTexture = new Texture2D(engine, frameWidth, frameHeight, TextureFormat.R8G8B8A8, false);
        // prettier-ignore
        this.uvTexture = new Texture2D(engine, frameWidth, frameHeight, TextureFormat.R8G8B8A8, false);
        // prettier-ignore
        this.material.shaderData.setTexture("y_texture", this.yTexture);
        this.material.shaderData.setTexture("uv_texture", this.uvTexture);
        // @ts-ignore
        this.yTexture._platformTexture._glTexture = cameraTexture.yTexture;
        // @ts-ignore
        this.uvTexture._platformTexture._glTexture = cameraTexture.uvTexture;
      }
      // 更新顶点调整矩阵
      const { elements } = this.projectMatrix;
      const displayTransform = curARFrame.getDisplayTransform();
      elements[0] = displayTransform[0];
      elements[1] = displayTransform[1];
      elements[2] = displayTransform[2];
      elements[4] = displayTransform[3];
      elements[5] = displayTransform[4];
      elements[6] = displayTransform[5];
      elements[8] = displayTransform[6];
      elements[9] = displayTransform[7];
      elements[10] = displayTransform[8];
    }
  }

  /** 初始化背景网格 */
  onAwake() {
    const { entity } = this;
    // MeshRenderer
    const renderer = entity.addComponent(MeshRenderer);
    renderer.mesh = this._createPlane();
    const material = this._getARBackgroundMaterial();
    renderer.setMaterial(material);
    this.material = material;
  }

  /**
   * 获取 AR 背景
   * @returns
   */
  private _getARBackgroundMaterial(): Material {
    // 初始化AR背景
    const yuv_vs = `
        attribute vec3 POSITION;
        attribute vec2 TEXCOORD_0;
        uniform mat4 u_proMatrix;
        varying vec2 v_texCoord;
        
        void main() {
          gl_Position = u_proMatrix * vec4(POSITION.xy, 1.0, 1.0);
          v_texCoord = TEXCOORD_0;
        }
        `;
    const yuv_fs = `
        uniform sampler2D y_texture;
        uniform sampler2D uv_texture;
        varying vec2 v_texCoord;
        void main() {
          vec4 y_color = texture2D(y_texture, v_texCoord);
          vec4 uv_color = texture2D(uv_texture, v_texCoord);

          float Y, U, V;
          float R ,G, B;
          Y = y_color.r;
          U = uv_color.r - 0.5;
          V = uv_color.a - 0.5;
          
          R = Y + 1.402 * V;
          G = Y - 0.344 * U - 0.714 * V;
          B = Y + 1.772 * U;
          gl_FragColor = vec4(R, G, B, 1.0);
        }
      `;
    const material = new Material(
      this.engine,
      Shader.create("ar-bg-android", yuv_vs, yuv_fs)
    );
    material.renderState.depthState.compareFunction = CompareFunction.LessEqual;
    material.renderState.rasterState.cullMode = CullMode.Off;
    material.renderQueueType = RenderQueueType.Opaque + 1;
    material.shaderData.setMatrix("u_proMatrix", this.projectMatrix);
    return material;
  }

  /**
   * 创建一个平面
   * @returns
   */
  private _createPlane(): ModelMesh {
    const mesh = new ModelMesh(this.engine);
    mesh.setPositions([
      new Vector3(1, 1, 1),
      new Vector3(-1, 1, 1),
      new Vector3(1, -1, 1),
      new Vector3(-1, -1, 1),
    ]);
    mesh.setUVs([
      new Vector2(1, 1),
      new Vector2(0, 1),
      new Vector2(1, 0),
      new Vector2(0, 0),
    ]);
    mesh.uploadData(true);
    mesh.addSubMesh(0, 4, MeshTopology.TriangleStrip);
    return mesh;
  }
}