import type { WebGLRenderTarget } from "three";
import * as THREE from "three";
import type { SceneCollection } from "./App";
import ColonyStats from "./ColonyStats";
import Config from "./Config";
import {
    MAT_AIR,
    MAT_DIRT,
    MAT_FOOD,
    MAT_HOME,
    MAT_ROCK,
    MAT_SAND,
} from "./constants";
import {
    generateColorData,
    generateLookupData,
    MaterialRegistry,
} from "./materials";
import {
    BEHAVIOR_GAS,
    BEHAVIOR_LIQUID,
    BEHAVIOR_POWDER,
    BEHAVIOR_SOLID,
} from "./materials/types";
import { getBlockOffset } from "./sand/margolus";

interface Resources {
    worldRenderTarget: THREE.WebGLRenderTarget;
    worldRenderTargetCopy: THREE.WebGLRenderTarget;
    worldBlurredRenderTarget: THREE.WebGLRenderTarget;
    sandPhysicsRenderTarget: THREE.WebGLRenderTarget;
    antsComputeTarget0: THREE.WebGLRenderTarget;
    antsComputeTarget1: THREE.WebGLRenderTarget;
    antsDiscreteRenderTarget: THREE.WebGLRenderTarget;
    antsPresenceRenderTarget: THREE.WebGLRenderTarget;
}

export default class Renderer {
    private renderer: THREE.WebGLRenderer;
    public resources!: Resources;
    private frameCounter = 0;
    private colonyStats = new ColonyStats();
    public readonly materialRegistry = new MaterialRegistry();
    public readonly materialPropsTexture!: THREE.DataTexture;
    public readonly materialColorTexture!: THREE.DataTexture;

    constructor(public canvas: HTMLCanvasElement) {
        this.renderer = new THREE.WebGLRenderer({ canvas });

        this.initResources();

        const propsData = generateLookupData(this.materialRegistry);
        this.materialPropsTexture = new THREE.DataTexture(
            propsData,
            256,
            1,
            THREE.RGBAFormat,
            THREE.FloatType,
        );
        this.materialPropsTexture.needsUpdate = true;

        const colorData = generateColorData(this.materialRegistry);
        this.materialColorTexture = new THREE.DataTexture(
            colorData,
            256,
            1,
            THREE.RGBAFormat,
            THREE.FloatType,
        );
        this.materialColorTexture.needsUpdate = true;
    }

    private initResources() {
        const antTextureSize = Math.round(Math.sqrt(2 ** Config.antsCount));

        this.resources = {
            worldRenderTarget: new THREE.WebGLRenderTarget(
                Config.worldSize,
                Config.worldSize,
                {
                    format: THREE.RGBAFormat,
                    type: THREE.FloatType,
                    depthBuffer: false,
                    magFilter: THREE.NearestFilter,
                    minFilter: THREE.NearestFilter,
                },
            ),
            worldRenderTargetCopy: new THREE.WebGLRenderTarget(
                Config.worldSize,
                Config.worldSize,
                {
                    format: THREE.RGBAFormat,
                    type: THREE.FloatType,
                    depthBuffer: false,
                    magFilter: THREE.NearestFilter,
                    minFilter: THREE.NearestFilter,
                },
            ),
            worldBlurredRenderTarget: new THREE.WebGLRenderTarget(
                Config.worldSize,
                Config.worldSize,
                {
                    format: THREE.RGBAFormat,
                    type: THREE.FloatType,
                    depthBuffer: false,
                    magFilter: THREE.NearestFilter,
                    minFilter: THREE.NearestFilter,
                },
            ),
            sandPhysicsRenderTarget: new THREE.WebGLRenderTarget(
                Config.worldSize,
                Config.worldSize,
                {
                    format: THREE.RGBAFormat,
                    type: THREE.FloatType,
                    depthBuffer: false,
                    magFilter: THREE.NearestFilter,
                    minFilter: THREE.NearestFilter,
                },
            ),
            antsComputeTarget0: Renderer.makeAntsMRT(antTextureSize),
            antsComputeTarget1: Renderer.makeAntsMRT(antTextureSize),
            antsDiscreteRenderTarget: new THREE.WebGLRenderTarget(
                Config.worldSize,
                Config.worldSize,
                {
                    format: THREE.RGBAFormat,
                    type: THREE.UnsignedByteType,
                    depthBuffer: false,
                    magFilter: THREE.NearestFilter,
                    minFilter: THREE.NearestFilter,
                },
            ),
            antsPresenceRenderTarget: new THREE.WebGLRenderTarget(
                Config.worldSize,
                Config.worldSize,
                {
                    format: THREE.RGBAFormat,
                    type: THREE.FloatType,
                    depthBuffer: false,
                    magFilter: THREE.NearestFilter,
                    minFilter: THREE.NearestFilter,
                },
            ),
        };
    }

    private static makeAntsMRT(size: number): THREE.WebGLRenderTarget {
        const mrt = new THREE.WebGLRenderTarget(size, size, {
            count: 2,
            type: THREE.FloatType,
            magFilter: THREE.NearestFilter,
            minFilter: THREE.NearestFilter,
            depthBuffer: false,
        });
        // both attachments get the same texture params
        for (const texture of mrt.textures) {
            texture.type = THREE.FloatType;
            texture.magFilter = THREE.NearestFilter;
            texture.minFilter = THREE.NearestFilter;
        }
        return mrt;
    }

    public renderSimulation(scenes: SceneCollection) {
        const [antsComputeSource, antsComputeTarget] =
            scenes.ants.getRenderTargets();

        // sand physics pass
        const blockOffset = getBlockOffset(this.frameCounter);
        this.setViewportFromRT(this.resources.sandPhysicsRenderTarget);
        this.renderer.setRenderTarget(this.resources.sandPhysicsRenderTarget);
        scenes.sandPhysics.material.uniforms.uWorld.value =
            this.resources.worldRenderTarget.texture;
        scenes.sandPhysics.material.uniforms.uMaterialProps.value =
            this.materialPropsTexture;
        scenes.sandPhysics.material.uniforms.uBlockOffset.value.set(
            blockOffset.x,
            blockOffset.y,
        );
        scenes.sandPhysics.material.uniforms.uFrame.value = this.frameCounter;
        this.renderer.render(scenes.sandPhysics, scenes.sandPhysics.camera);
        this.frameCounter++;

        this.setViewportFromRT(this.resources.worldBlurredRenderTarget);
        this.renderer.setRenderTarget(this.resources.worldBlurredRenderTarget);
        scenes.worldBlur.material.uniforms.tWorld.value =
            this.resources.sandPhysicsRenderTarget.texture;
        this.renderer.render(scenes.worldBlur, scenes.worldBlur.camera);

        this.renderer.setRenderTarget(this.resources.antsPresenceRenderTarget);
        this.renderer.clear();

        this.setViewportFromRT(antsComputeTarget);
        this.renderer.setRenderTarget(antsComputeTarget);
        scenes.ants.material.uniforms.tLastState.value =
            antsComputeSource.textures[0];
        scenes.ants.material.uniforms.tLastExtState.value =
            antsComputeSource.textures[1];
        scenes.ants.material.uniforms.tWorld.value =
            this.resources.worldBlurredRenderTarget.texture;
        scenes.ants.material.uniforms.tPresence.value =
            this.resources.antsPresenceRenderTarget.texture;
        this.renderer.render(scenes.ants, scenes.ants.camera);

        this.setViewportFromRT(this.resources.antsDiscreteRenderTarget);
        this.renderer.setRenderTarget(this.resources.antsDiscreteRenderTarget);
        scenes.discretize.material.uniforms.tDataCurrent.value =
            antsComputeTarget.textures[0];
        scenes.discretize.material.uniforms.tDataLast.value =
            antsComputeSource.textures[0];
        scenes.discretize.material.uniforms.tDataExtCurrent.value =
            antsComputeTarget.textures[1];
        this.renderer.render(scenes.discretize, scenes.discretize.camera);

        this.setViewportFromRT(this.resources.worldRenderTarget);
        this.renderer.setRenderTarget(this.resources.worldRenderTarget);
        scenes.world.material.uniforms.tLastState.value =
            this.resources.worldBlurredRenderTarget.texture;
        scenes.world.material.uniforms.tDiscreteAnts.value =
            this.resources.antsDiscreteRenderTarget.texture;
        this.renderer.render(scenes.world, scenes.world.camera);

        scenes.screen.material.uniforms.tData.value =
            antsComputeTarget.textures[0];
        scenes.screen.groundMaterial.uniforms.map.value =
            this.resources.worldRenderTargetCopy.texture;

        // colony stats readback — read ant state texture, compute aggregate stats on CPU
        this.colonyStats.update(this.renderer, antsComputeTarget);
        scenes.ants.material.uniforms.uForagerRatio.value =
            this.colonyStats.data.foragerRatio;
    }

    private setViewportFromRT(rt: WebGLRenderTarget) {
        this.renderer.setViewport(0, 0, rt.width, rt.height);
    }

    public renderToScreen(scenes: SceneCollection) {
        this.setViewportFromRT(this.resources.worldRenderTargetCopy);
        this.renderer.setRenderTarget(this.resources.worldRenderTargetCopy);
        scenes.draw.material.uniforms.tWorld.value =
            this.resources.worldRenderTarget.texture;
        scenes.draw.material.uniforms.pointerPosition.value =
            scenes.screen.pointerPosition;
        scenes.draw.material.uniforms.drawMode.value = scenes.screen.drawMode;
        scenes.draw.material.uniforms.brushRadius.value = Config.brushRadius;
        this.renderer.render(scenes.draw, scenes.draw.camera);
        this.renderer.copyFramebufferToTexture(
            this.resources.worldRenderTarget.texture,
            new THREE.Vector2(),
        );

        this.renderer.setViewport(
            0,
            0,
            scenes.screen.renderWidth,
            scenes.screen.renderHeight,
        );
        this.renderer.setRenderTarget(null);
        this.renderer.render(scenes.screen, scenes.screen.camera);
    }

    public resizeCanvas(width: number, height: number) {
        this.renderer.setSize(width, height, false);
    }

    public getCommonMaterialDefines(): Record<string, string> {
        return {
            WORLD_SIZE: Renderer.convertNumberToFloatString(Config.worldSize),
            SCENT_THRESHOLD: Renderer.convertNumberToFloatString(
                Config.scentThreshold,
            ),
            SCENT_FADE_OUT_FACTOR: Renderer.convertNumberToFloatString(
                Config.scentFadeOutFactor,
            ),
            SCENT_BLUR_RADIUS: Renderer.convertNumberToFloatString(
                Config.scentBlurRadius,
            ),
            SCENT_MAX_STORAGE: Renderer.convertNumberToFloatString(
                Config.scentMaxStorage,
            ),
            SCENT_PER_MARKER: Renderer.convertNumberToFloatString(
                Config.scentPerMarker,
            ),
            SCENT_MAX_PER_CELL: Renderer.convertNumberToFloatString(
                Config.scentMaxPerCell,
            ),
            ANT_SPEED: Renderer.convertNumberToFloatString(Config.antSpeed),
            ANT_ROTATION_ANGLE: Renderer.convertNumberToFloatString(
                Config.antRotationAngle,
            ),
            REPELLENT_FADE_OUT_FACTOR: Renderer.convertNumberToFloatString(
                Config.repellentFadeOutFactor,
            ),
            REPELLENT_BLUR_RADIUS: Renderer.convertNumberToFloatString(
                Config.repellentBlurRadius,
            ),
            REPELLENT_MAX_PER_CELL: Renderer.convertNumberToFloatString(
                Config.repellentMaxPerCell,
            ),
            REPELLENT_THRESHOLD: Renderer.convertNumberToFloatString(
                Config.repellentThreshold,
            ),
            MAT_AIR: String(MAT_AIR),
            MAT_SAND: String(MAT_SAND),
            MAT_DIRT: String(MAT_DIRT),
            MAT_ROCK: String(MAT_ROCK),
            MAT_FOOD: String(MAT_FOOD),
            MAT_HOME: String(MAT_HOME),
            BEHAVIOR_POWDER:
                Renderer.convertNumberToFloatString(BEHAVIOR_POWDER),
            BEHAVIOR_LIQUID:
                Renderer.convertNumberToFloatString(BEHAVIOR_LIQUID),
            BEHAVIOR_GAS: Renderer.convertNumberToFloatString(BEHAVIOR_GAS),
            BEHAVIOR_SOLID: Renderer.convertNumberToFloatString(BEHAVIOR_SOLID),
        };
    }

    public reset(scenes: SceneCollection) {
        const antTextureSize = Math.round(Math.sqrt(2 ** Config.antsCount));

        this.resources.worldRenderTarget.setSize(
            Config.worldSize,
            Config.worldSize,
        );
        this.renderer.setRenderTarget(this.resources.worldRenderTarget);
        this.renderer.clear();

        this.resources.worldRenderTargetCopy.setSize(
            Config.worldSize,
            Config.worldSize,
        );
        this.renderer.setRenderTarget(this.resources.worldRenderTargetCopy);
        this.renderer.clear();

        this.resources.worldBlurredRenderTarget.setSize(
            Config.worldSize,
            Config.worldSize,
        );
        this.renderer.setRenderTarget(this.resources.worldBlurredRenderTarget);
        this.renderer.clear();

        this.resources.sandPhysicsRenderTarget.setSize(
            Config.worldSize,
            Config.worldSize,
        );
        this.renderer.setRenderTarget(this.resources.sandPhysicsRenderTarget);
        this.renderer.clear();

        this.frameCounter = 0;

        this.resources.antsComputeTarget0.setSize(
            antTextureSize,
            antTextureSize,
        );
        this.renderer.setRenderTarget(this.resources.antsComputeTarget0);
        this.renderer.clear();

        this.resources.antsComputeTarget1.setSize(
            antTextureSize,
            antTextureSize,
        );
        this.renderer.setRenderTarget(this.resources.antsComputeTarget1);
        this.renderer.clear();

        this.resources.antsDiscreteRenderTarget.setSize(
            Config.worldSize,
            Config.worldSize,
        );
        this.renderer.setRenderTarget(this.resources.antsDiscreteRenderTarget);
        this.renderer.clear();

        this.resources.antsPresenceRenderTarget.setSize(
            Config.worldSize,
            Config.worldSize,
        );
        this.renderer.setRenderTarget(this.resources.antsPresenceRenderTarget);
        this.renderer.clear();

        for (const scene of Object.values(scenes)) {
            scene.recompileMaterials();
        }
    }

    static convertNumberToFloatString(n: number): string {
        return n.toFixed(8);
    }
}