import { ShaderMaterial, Texture } from "three";
import { GroundedSkybox as GroundProjection } from 'three/examples/jsm/objects/GroundedSkybox.js';

import { Gizmos } from "../engine/engine_gizmos.js";
import { serializable } from "../engine/engine_serialization_decorator.js";
import { getBoundingBox, getTempVector, getWorldScale, Graphics, setVisibleInCustomShadowRendering, setWorldPosition } from "../engine/engine_three_utils.js";
import { delayForFrames, getParam, Watch as Watch } from "../engine/engine_utils.js";
import { Behaviour } from "./Component.js";

const debug = getParam("debuggroundprojection");

/**
 * GroundProjectedEnv creates a ground projection of the current environment map.  
 * @category Rendering
 * @group Components
 */
export class GroundProjectedEnv extends Behaviour {

    /**
     * If true the projection will be created on awake and onEnable
     * @default false
     */
    @serializable()
    applyOnAwake: boolean = false;

    /**
     * When enabled the position of the projected environment will be adjusted to be centered in the scene (and ground level).
     * @default true
     */
    @serializable()
    autoFit: boolean = true;

    /**
     * Radius of the projection sphere. Set it large enough so the camera stays inside (make sure the far plane is also large enough)
     * @default 50
     */
    @serializable()
    set radius(val: number) {
        this._radius = val;
        if (this._projection) this.updateProjection();
    }
    get radius(): number { return this._radius; }
    private _radius: number = 50;

    /**
     * How far the camera that took the photo was above the ground. A larger value will magnify the downward part of the image.
     * @default 3
     */
    @serializable()
    set height(val: number) {
        this._height = val;
        if (this._projection) this.updateProjection();
    }
    get height(): number { return this._height; }
    private _height: number = 3;

    /**
     * Blending factor for the AR projection being blended with the scene background.    
     * 0 = not visible in AR - 1 = blended with real world background.    
     * Values between 0 and 1 control the smoothness of the blend while lower values result in smoother blending.
     * @default 0
     */
    @serializable()
    set arBlending(val: number) {
        this._arblending = val;
        this._needsTextureUpdate = true;
    }
    get arBlending(): number { return this._arblending; }
    private _arblending = 0;

    private _lastBackground?: Texture;
    private _lastRadius?: number;
    private _lastHeight?: number;
    private _projection?: GroundProjection;
    private _watcher?: Watch;


    /** @internal */
    awake() {
        if (this.applyOnAwake)
            this.updateAndCreate();
    }
    /** @internal */
    onEnable() {
        // TODO: if we do this in the first frame we can not disable it again. Something buggy with the watch?!
        if (this.context.time.frameCount > 0) {
            if (this.applyOnAwake)
                this.updateAndCreate();
        }
        if (!this._watcher) {
            this._watcher = new Watch(this.context.scene, "background");
            this._watcher.subscribeWrite(_ => {
                if (debug) console.log("Background changed", this.context.scene.background);
                this._needsTextureUpdate = true;
            });
        }
    }
    /** @internal */
    onDisable() {
        this._watcher?.revoke();
        this._projection?.removeFromParent();
    }
    /** @internal */
    onEnterXR(): void {
        if (!this.activeAndEnabled) return;
        this._needsTextureUpdate = true;
        this.updateProjection();
    }
    /** @internal */
    async onLeaveXR() {
        if (!this.activeAndEnabled) return;
        await delayForFrames(1);
        this.updateProjection();
    }
    /** @internal */
    onBeforeRender(): void {
        if (this._projection && this.scene.backgroundRotation) {
            this._projection.rotation.copy(this.scene.backgroundRotation);
        }

        const blurrinessChanged = this.context.scene.backgroundBlurriness !== undefined && this._lastBlurriness != this.context.scene.backgroundBlurriness && this.context.scene.backgroundBlurriness > 0.001;
        if (blurrinessChanged) {
            this.updateProjection();
        }
        else if (this._needsTextureUpdate && this.context.scene.background instanceof Texture) {
            this.updateBlurriness(this.context.scene.background, this.context.scene.backgroundBlurriness);
        }
    }

    private updateAndCreate() {
        this.updateProjection();
        this._watcher?.apply();
    }


    private _needsTextureUpdate = false;

    /**
     * Updates the ground projection. This is called automatically when the environment or settings change.
     */
    updateProjection() {
        if (!this.context.scene.background) {
            this._projection?.removeFromParent();
            return;
        }
        const backgroundTexture = this.context.scene.background;
        if (!(backgroundTexture instanceof Texture)) {
            this._projection?.removeFromParent();
            return;
        }

        if (this.context.xr?.isPassThrough || this.context.xr?.isAR) {
            if (this.arBlending === 0) {
                this._projection?.removeFromParent();
                return;
            }
        }
        // If this is called from a setter during initialization
        if (!this.gameObject || this.destroyed) {
            return;
        }

        let needsNewAutoFit = true;
        // offset here must be zero (and not .01) because the plane occlusion (when mesh tracking is active) is otherwise not correct
        const offset = 0;

        const hasChanged = backgroundTexture !== this._lastBackground || this._height !== this._lastHeight || this._radius !== this._lastRadius;

        if (!this._projection || hasChanged) {
            if (debug) console.log("Create/Update Ground Projection", backgroundTexture.name);

            this._projection?.removeFromParent();

            try {
                this._projection = new GroundProjection(backgroundTexture, this._height, this._radius, 64);
            }
            catch (e) {
                console.error("Error creating three GroundProjection", e);
                return;
            }
            this._projection.position.y = this._height - offset;
            this._projection.name = "GroundProjection";
            setVisibleInCustomShadowRendering(this._projection, false);
        }
        else {
            needsNewAutoFit = false;
        }

        if (!this._projection.parent)
            this.gameObject.add(this._projection);

        if (this.autoFit && needsNewAutoFit) {
            // TODO: should also update the radius (?)
            this._projection.updateWorldMatrix(true, true);
            const box = getBoundingBox(this.context.scene.children, [this._projection]);

            const floor_y = box.min.y;
            if (floor_y < Infinity) {
                const wp = getTempVector();
                wp.x = box.min.x + (box.max.x - box.min.x) * .5;
                const scale = getWorldScale(this.gameObject).x;
                wp.y = floor_y + (this._height * scale) - offset;
                wp.z = box.min.z + (box.max.z - box.min.z) * .5;
                setWorldPosition(this._projection, wp);
            }

            if (debug) Gizmos.DrawWireBox3(box, 0x00ff00, 5);
        }

        /* TODO realtime adjustments aren't possible anymore with GroundedSkybox (mesh generation)
        this.env.scale.setScalar(this._scale);
        this.env.radius = this._radius;
        this.env.height = this._height;
        */

        if (this.context.scene.backgroundBlurriness > 0.001 && this._needsTextureUpdate) {
            this.updateBlurriness(backgroundTexture, this.context.scene.backgroundBlurriness);
        }

        this._lastBackground = backgroundTexture;
        this._lastHeight = this._height;
        this._lastRadius = this._radius;
        this._needsTextureUpdate = false;
    }

    private _blurrynessShader: ShaderMaterial | null = null;
    private _lastBlurriness: number = -1;

    private updateBlurriness(texture: Texture, blurriness: number) {
        if (!this._projection) {
            return;
        }
        else if (!texture) {
            return;
        }

        this._needsTextureUpdate = false;
        if (debug) console.log("Update Blurriness", blurriness);
        this._blurrynessShader ??= new ShaderMaterial({
            name: "GroundProjectionBlurriness",
            uniforms: {
                map: { value: texture },
                blurriness: { value: blurriness },
                blending: { value: 0 },
                alphaFactor: { value: 1 }
            },
            vertexShader: blurVertexShader,
            fragmentShader: blurFragmentShader
        });
        this._blurrynessShader.depthWrite = false;
        this._blurrynessShader.uniforms.map.value = texture;
        this._blurrynessShader.uniforms.blurriness.value = blurriness;
        this._lastBlurriness = blurriness;
        texture.needsUpdate = true;

        const wasTransparent = this._projection.material.transparent;
        this._projection.material.transparent = (this.context.xr?.isAR === true && this.arBlending > 0.000001) ?? false;
        if (this._projection.material.transparent) {
            this._blurrynessShader.uniforms.blending.value = this.arBlending;
        }
        else { this._blurrynessShader.uniforms.blending.value = 0; }

        if (this.context.isInPassThrough) {
            // Make the ground slightly transparent in passthrough mode
            this._blurrynessShader.uniforms.alphaFactor.value = 0.95;
        }
        else {
            this._blurrynessShader.uniforms.alphaFactor.value = 1;
        }

        // Make sure the material is updated if the transparency changed
        if (wasTransparent !== this._projection.material.transparent) {
            this._projection.material.needsUpdate = true;
        }

        // Update the texture
        this._projection.material.map = Graphics.copyTexture(texture, this._blurrynessShader);
        this._projection.material.depthTest = true;
        this._projection.material.depthWrite = false;
    }

}


const blurVertexShader = `
  varying vec2 vUv;

  void main() {
    vUv = uv;
    gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
  }
`;

// Fragment Shader
const blurFragmentShader = `
  uniform sampler2D map;
  uniform float blurriness;
  uniform float alphaFactor;
  uniform float blending;
  varying vec2 vUv;

  const float PI = 3.14159265359;

  // Gaussian function
  float gaussian(float x, float sigma) {
    return exp(-(x * x) / (2.0 * sigma * sigma)) / (sqrt(2.0 * PI) * sigma);
  }

  // Custom smoothstep function for desired falloff
  float customSmoothstep(float edge0, float edge1, float x) {
    float t = clamp((x - edge0) / (edge1 - edge0), 0.0, 1.0);
    return t * t * (3.0 - 2.0 * t);
  }

  void main() {
    vec2 center = vec2(0.0, 0.0);
    vec2 pos = vUv;
    pos.x = 0.0; // Only consider vertical distance
    float distance = length(pos - center);
    
    // Calculate blur amount based on custom falloff
    float blurAmount = customSmoothstep(0.5, 1.0, distance * 2.0);
    blurAmount = clamp(blurAmount, 0.0, 1.0); // Ensure blur amount is within valid range

    // Gaussian blur
    vec2 pixelSize = 1.0 / vec2(textureSize(map, 0));
    vec4 color = vec4(0.0);
    float totalWeight = 0.0;
    int blurSize = int(60.0 * min(1.0, blurriness) * blurAmount); // Adjust blur size based on distance and blurriness
    float lodLevel = log2(float(blurSize)) * 0.5; // Compute LOD level

    for (int x = -blurSize; x <= blurSize; x++) {
        for (int y = -blurSize; y <= blurSize; y++) {
            vec2 offset = vec2(float(x), float(y)) * pixelSize * blurAmount;
            float weight = gaussian(length(vec2(float(x), float(y))), 1000.0 * blurAmount); // Use a fixed sigma value
            color += textureLod(map, vUv + offset, lodLevel) * weight;
            totalWeight += weight;
        }
    }

    color = totalWeight > 0.0 ? color / totalWeight : texture2D(map, vUv);

    gl_FragColor = color;

    float brightness = dot(gl_FragColor.rgb, vec3(0.299, 0.587, 0.114));
    float stepFactor = blending - brightness * .1;
    gl_FragColor.a = pow(1.0 - blending * customSmoothstep(0.35 * stepFactor, 0.45 * stepFactor, distance), 5.);
    gl_FragColor.a *= alphaFactor;
    // gl_FragColor.rgb = vec3(1.0);

    // #include <tonemapping_fragment>
    // #include <colorspace_fragment>
    
    // Uncomment to visualize blur amount
    // gl_FragColor = vec4(blurAmount, 0.0, 0.0, 1.0);
  }
`;