import { Vector3 } from "three";
import { assert } from "../../../../../core/assert.js";
import { BinaryDataType } from "../../../../../core/binary/type/BinaryDataType.js";
import { AABB3 } from "../../../../../core/geom/3d/aabb/AABB3.js";
import { aabb3_from_v3_array } from "../../../../../core/geom/3d/aabb/aabb3_from_v3_array.js";
import { AttributeSpec } from "../../../geometry/AttributeSpec.js";
import { composeCompile } from "../../../material/composeCompile.js";
import { AbstractMaterialTransformer } from "../../../material/manager/AbstractMaterialTransformer.js";
import { isLitMaterial } from "../../../render/forward_plus/plugin/isLitMaterial.js";
import { AttributeDataTexture } from "../../../texture/AttributeDataTexture.js";
import { octahedral_depth_to_atlas, vector_buffer_to_attribute_texture } from "../../lpv/probe_volume_to_textures.js";
import { tetrahedral_mesh_to_texture } from "../tetrahedral_mesh_to_texture.js";
import { makeLookupTexture3D } from "./makeLookupTexture3D.js";
import { makeOctahedralDepthAtlas } from "./makeOctahedralDepthAtlas.js";
import { space_fragment_transform_shader } from "./space_fragment/space_fragment_transform_shader.js";
import { space_vertex_transform_shader } from "./space_vertex/space_vertex_transform_shader.js";

const PROPERTY_TRANSFORMER_MARKER = '@light-probe-volume-material-transformer';

/**
 * Where should probes be resolved?
 * @enum {number}
 */
export const ProbeResolutionStage = {
    /**
     * Cheaper, but less precise
     */
    Vertex: 0,
    /**
     * More expensive, necessary for large geometries with few vertices
     */
    Fragment: 1
};


export class MaterialTransformer extends AbstractMaterialTransformer {

    /**
     *
     * @type {LightProbeVolume|null}
     */
    #volume = null;

    #mesh_lookup = makeLookupTexture3D(32);
    #mesh_bounds = new AABB3();

    /**
     *
     * @type {ProbeResolutionStage|number}
     */
    #sampling_stage = ProbeResolutionStage.Fragment

    #textures = {
        vertex: new AttributeDataTexture(
            AttributeSpec.fromJSON({
                name: "vertex",
                type: BinaryDataType.Uint32,
                itemSize: 4,
                normalized: false
            }),
            256, 1
        ),
        neighbour: new AttributeDataTexture(
            AttributeSpec.fromJSON({
                name: "neighbour",
                type: BinaryDataType.Uint32,
                itemSize: 4,
                normalized: false
            }),
            256, 1
        ),
        positions: new AttributeDataTexture(
            AttributeSpec.fromJSON({
                name: "positions",
                type: BinaryDataType.Float32,
                itemSize: 3,
                normalized: false
            }),
            256, 1
        ),
        harmonics: new AttributeDataTexture(
            AttributeSpec.fromJSON({
                name: "harmonics",
                type: BinaryDataType.Float32,
                itemSize: 3,
                normalized: false
            }),
            256, 9
        ),
        depth: makeOctahedralDepthAtlas()
    };

    #uniforms = {
        lpv_t_mesh_vertices: {
            type: "t",
            value: this.#textures.vertex.texture
        },
        lpv_t_mesh_neighbours: {
            type: "t",
            value: this.#textures.neighbour.texture
        },
        lpv_t_probe_positions: {
            type: "t",
            value: this.#textures.positions.texture
        },
        lpv_t_probe_data: {
            type: "t",
            value: this.#textures.harmonics.texture
        },
        lpv_t_probe_depth: {
            type: "t",
            value: this.#textures.depth
        },
        lpv_u_probe_depth_resolution: {
            value: 1
        },
        lpv_t_mesh_lookup: {
            type: "t",
            value: this.#mesh_lookup
        },
        lpv_v3_bounds_min: {
            value: new Vector3(),
        },
        lpv_v3_bounds_max: {
            value: new Vector3(),
        },
        lpv_u_mesh_tet_count: {
            value: 0
        },
        lpv_f_intensity: {
            value: 1
        }
    };

    /**
     *
     * @type {number}
     */
    #uniform_version = -1;

    /**
     *
     * @param {{uniforms, glslVersion, defines, vertexShader, fragmentShader}} shader
     */
    #onBeforeCompile = (shader) => {
        Object.assign(shader.uniforms, this.#uniforms);

        switch (this.#sampling_stage) {
            case ProbeResolutionStage.Vertex:
                space_vertex_transform_shader(shader);
                break;
            case ProbeResolutionStage.Fragment:
                space_fragment_transform_shader(shader);
                break;
            default:
                throw new Error(`Unsupported resolution space ${this.#sampling_stage}`);
        }
    };

    #materials = [];

    /**
     *
     * @return {number}
     */
    get intensity() {
        return this.#uniforms.lpv_f_intensity.value;
    }

    /**
     *
     * @param {number} v
     */
    set intensity(v) {
        this.#uniforms.lpv_f_intensity.value = v;
    }


    /**
     *
     * @param {LightProbeVolume} volume
     * @param {number|ProbeResolutionStage} stage either "vertex" or "fragment". Vertex sampling is faster but less accurate, fragment is more expensive and precise
     */
    constructor({
                    volume,
                    stage = ProbeResolutionStage.Vertex,
                }) {

        super();

        assert.defined(volume, 'volume');
        assert.enum(stage, ProbeResolutionStage, 'stage');

        this.#volume = volume;
        this.#sampling_stage = stage;

    }

    update_bounds() {

        aabb3_from_v3_array(this.#mesh_bounds, this.#volume.points, this.#volume.count * 3);

    }

    update_lookup() {
        this.update_bounds();

        const bounds = this.#mesh_bounds;

        const ex = bounds.getExtentsX();
        const ey = bounds.getExtentsY();
        const ez = bounds.getExtentsZ();

        const lookup = this.#mesh_lookup;

        const image = lookup.image;

        const res_z = image.depth;
        const res_y = image.height;
        const res_x = image.width;

        const volume = this.#volume;
        const points = volume.points;
        const mesh = volume.mesh;

        let last_tet = 0;

        for (let z = 0; z < res_z; z++) {
            for (let y = 0; y < res_y; y++) {
                for (let x = 0; x < res_x; x++) {

                    // convert pixel position to world coordinate
                    const n_x = (x + 0.5) / res_x
                    const n_y = (y + 0.5) / res_y
                    const n_z = (z + 0.5) / res_z

                    const w_x = n_x * ex + bounds.x0;
                    const w_y = n_y * ey + bounds.y0;
                    const w_z = n_z * ez + bounds.z0;

                    const tet = mesh.walkToTetraContainingPoint(w_x, w_y, w_z, points, last_tet);

                    if (tet !== -1) {
                        last_tet = tet;
                    }

                    image.data[z * res_y * res_x + y * res_x + x] = last_tet;

                }
            }
        }

        lookup.needsUpdate = true;
    }

    update() {
        const lpv = this.#volume;

        if (lpv === null) {
            return;
        }

        if (this.#uniform_version === lpv.version) {
            // up to date
            return;
        }

        this.update_lookup();

        const uniforms = this.#uniforms;

        uniforms.lpv_u_mesh_tet_count.value = lpv.mesh.count;

        uniforms.lpv_v3_bounds_min.value.set(
            this.#mesh_bounds.x0, this.#mesh_bounds.y0, this.#mesh_bounds.z0
        );
        uniforms.lpv_v3_bounds_max.value.set(
            this.#mesh_bounds.x1, this.#mesh_bounds.y1, this.#mesh_bounds.z1
        );

        uniforms.lpv_u_probe_depth_resolution.value = lpv.depth_map_resolution;

        this.#uniform_version = lpv.version;

        console.time('prepareTextures');

        vector_buffer_to_attribute_texture(this.#textures.positions, lpv.points, lpv.count, 3);
        vector_buffer_to_attribute_texture(this.#textures.harmonics, lpv.harmonics, lpv.count * 9, 3);
        tetrahedral_mesh_to_texture(
            this.#textures.vertex, this.#textures.neighbour, lpv.mesh
        );
        octahedral_depth_to_atlas(lpv, this.#textures.depth.image.data, this.#textures.depth.image.width);
        this.#textures.depth.needsUpdate = true;

        console.timeEnd('prepareTextures');

        console.log(this.#textures);

        for (let i = 0; i < this.#materials.length; i++) {
            const material = this.#materials[i];

            // material.uniformsNeedUpdate = true;
        }
    }

    transform(source) {
        if (source.hasOwnProperty(PROPERTY_TRANSFORMER_MARKER)) {
            // already transformed

            if (source[PROPERTY_TRANSFORMER_MARKER] !== this) {
                throw new Error('The material is already transformed, but is associated with a different transformer instance');
            } else {
                return source;
            }

        }


        let result = source;

        if (isLitMaterial(source)) {

            result.clone();

            result.onBeforeCompile = composeCompile(source.onBeforeCompile, this.#onBeforeCompile);

            result[PROPERTY_TRANSFORMER_MARKER] = this;

            this.#materials.push(result);
        }

        return result;
    }
}