import { convert_channel_sRGB_to_linear } from "../../../../../core/color/sRGB/sRGB_to_linear.js";
import { renderSamplerFromTexture } from "../../../render/visibility/hiz/buildCanvasViewFromTexture.js";
import { HemiOctahedralUvEncoder } from "../../octahedral/grid/HemiOctahedralUvEncoder.js";
import { OctahedralUvEncoder } from "../../octahedral/grid/OctahedralUvEncoder.js";
import { ImpostorCaptureType } from "../../octahedral/ImpostorCaptureType.js";

/**
 * Layout per output sample (Float32Array, see ATLAS_SAMPLE_STRIDE):
 *
 * | offset | size | content                                    |
 * |--------|------|--------------------------------------------|
 * |   0    |   3  | position.xyz (object space, around offset) |
 * |   3    |   3  | normal.xyz   (object space, unit length)   |
 * |   6    |   3  | albedo.rgb   (linear, no alpha)            |
 * |   9    |   3  | orm (occlusion, roughness, metalness)      |
 *
 * @type {number}
 */
export const ATLAS_SAMPLE_STRIDE = 12;


/**
 * Read back a single MRT attachment of the impostor atlas as flat RGBA8
 * bytes. {@link renderSamplerFromTexture} renders the source texture into
 * a side render target and reads back UNSIGNED_BYTE so this works on every
 * WebGL2 platform, at the cost of 8-bit depth precision (256 steps across
 * the 2*radius diameter).
 *
 * @param {THREE.WebGLRenderer} renderer
 * @param {THREE.Texture} texture
 * @returns {Uint8ClampedArray}
 */
function read_atlas_texture_rgba_u8(renderer, texture) {
    const ctrl = renderSamplerFromTexture({
        renderer,
        texture,
        // We need the bytes in their stored order (no flip) because the
        // bake renderer wrote them in y-up viewport order, and the frame
        // tile index math below assumes that order too.
        flipY: false
    });
    ctrl.render();
    return ctrl.destination.data;
}


/**
 * Decode the bake-time view-space normal back into the impostor's local
 * (object) space. Inverse of the encode side in BakeShaderStandard.
 *
 * The baker placed an OrthographicCamera at
 *
 *     camera_pos = centre + radius * D_frame
 *
 * with three.js's lookAt(centre) using world-up = (0,1,0). That gives a
 * camera basis (bake_right, bake_up, +D_frame) in object space. The
 * baker wrote the surface normal as (normal.x = dot(n, right),
 * normal.y = dot(n, up), normal.z = dot(n, D_frame)) and packed it into
 * the atlas RGB channels as (* 0.5 + 0.5). So to reverse, decode the RGB
 * back to ±1 and recombine through the same basis.
 *
 * @param {number[]} out          length-3 output, object-space normal
 * @param {number} nr             encoded R in [0,1]
 * @param {number} ng             encoded G in [0,1]
 * @param {number} nb             encoded B in [0,1]
 * @param {number[]} D_frame      bake direction in object space (unit)
 */
function decode_normal_to_object_space(out, nr, ng, nb, D_frame) {
    // 1. Decode bake-view-space normal from [0,1] to [-1,1]
    const nx = nr * 2 - 1;
    const ny = ng * 2 - 1;
    const nz = nb * 2 - 1;

    // 2. Mirror three.js Camera.lookAt's polar nudge — when D_frame is
    //    parallel to world-up, lookAt jiggles +Z slightly to avoid a
    //    degenerate cross product. We have to do the same so our basis
    //    matches the baker's.
    let Dx = D_frame[0], Dy = D_frame[1], Dz = D_frame[2];
    if (Math.abs(Dy) > 0.99999) {
        // Normalise (D + (0, 0, 0.0001)) — same nudge the lit shader uses.
        const dz_n = Dz + 0.0001;
        const len = Math.sqrt(Dx * Dx + Dy * Dy + dz_n * dz_n);
        Dx /= len;
        Dy /= len;
        Dz = dz_n / len;
    }

    // 3. bake_right = normalize(cross((0,1,0), D))
    //    Cross product: (0,1,0) × (Dx,Dy,Dz) = (Dz, 0, -Dx)
    let rx = Dz;
    let ry = 0;
    let rz = -Dx;
    {
        const len = Math.sqrt(rx * rx + ry * ry + rz * rz);
        rx /= len; ry /= len; rz /= len;
    }
    // 4. bake_up = cross(D, right) — no normalise needed, D and right
    //    are already orthonormal.
    const ux = Dy * rz - Dz * ry;
    const uy = Dz * rx - Dx * rz;
    const uz = Dx * ry - Dy * rx;

    // 5. Recompose in object space.
    out[0] = rx * nx + ux * ny + Dx * nz;
    out[1] = ry * nx + uy * ny + Dy * nz;
    out[2] = rz * nx + uz * ny + Dz * nz;

    // Re-normalise — interpolation across pixels and 8-bit quantisation
    // can drift the magnitude.
    const out_len = Math.sqrt(out[0] * out[0] + out[1] * out[1] + out[2] * out[2]);
    if (out_len > 1e-6) {
        out[0] /= out_len;
        out[1] /= out_len;
        out[2] /= out_len;
    }
}


/**
 * Reconstruct world (= object-local) position of the surface pixel from
 * the atlas-encoded depth + frame index + frame-local pixel coords.
 *
 * Convention from BakeShaderStandard:
 *   - camera at C = centre + radius * D_frame
 *   - looking back at centre (=> view forward = -D_frame)
 *   - ortho frustum: half-width = radius, half-height = radius, near = 0, far = 2*radius
 *   - depth_atlas = 1 + mvPosition.z / far
 *      - at near (z=0)       => depth = 1
 *      - at centre (z=-r)    => depth = 0.5
 *      - at far  (z=-2r)     => depth = 0
 *
 * So along the view ray (origin = C + u*right + v*up where u,v ∈ [-r, +r],
 * direction = -D_frame), the surface point sits at parametric distance
 * t = (1 - depth) * 2*radius from the ray origin.
 *
 * @param {number[]} out         length-3 output (object space)
 * @param {number} depth01       depth in [0,1] as stored in atlas alpha
 * @param {number} frame_u_local frame-local X in [-1, +1]
 * @param {number} frame_v_local frame-local Y in [-1, +1]
 * @param {number[]} D_frame     bake direction in object space (unit)
 * @param {number} radius        bounding sphere radius
 * @param {number[]} centre      bounding sphere centre (object space)
 */
function reconstruct_position(out, depth01, frame_u_local, frame_v_local, D_frame, radius, centre) {
    // Same polar nudge as the normal decoder, so right/up agree.
    let Dx = D_frame[0], Dy = D_frame[1], Dz = D_frame[2];
    if (Math.abs(Dy) > 0.99999) {
        const dz_n = Dz + 0.0001;
        const len = Math.sqrt(Dx * Dx + Dy * Dy + dz_n * dz_n);
        Dx /= len;
        Dy /= len;
        Dz = dz_n / len;
    }

    // Same basis as decode_normal_to_object_space; see the cross-product
    // derivation there.
    let rx = Dz, ry = 0, rz = -Dx;
    {
        const len = Math.sqrt(rx * rx + ry * ry + rz * rz);
        rx /= len; ry /= len; rz /= len;
    }
    const ux = Dy * rz - Dz * ry;
    const uy = Dz * rx - Dx * rz;
    const uz = Dx * ry - Dy * rx;

    // C = centre + radius * D
    const cx = centre[0] + radius * Dx;
    const cy = centre[1] + radius * Dy;
    const cz = centre[2] + radius * Dz;

    // ray origin = C + frame_u_local*r*right + frame_v_local*r*up
    const ox = cx + frame_u_local * radius * rx + frame_v_local * radius * ux;
    const oy = cy + frame_u_local * radius * ry + frame_v_local * radius * uy;
    const oz = cz + frame_u_local * radius * rz + frame_v_local * radius * uz;

    // surface = origin + t * (-D)
    const t = (1 - depth01) * 2 * radius;
    out[0] = ox - t * Dx;
    out[1] = oy - t * Dy;
    out[2] = oz - t * Dz;
}


/**
 * Read back the baked octahedral atlas (3 MRT textures: diffuse+alpha,
 * normal+depth, orm) and convert every covered pixel into one CPU sample:
 * world-space position, normal, albedo, ORM.
 *
 * This is the noisy point-cloud the voxelisation step bins into a sparse
 * grid. The atlas captures the surface from many directions, so the same
 * surface point shows up multiple times — that's a feature, not a bug:
 * voxel averaging cancels per-pixel noise and per-frame quantisation
 * error, and surfaces visible from many angles (the ones that matter for
 * impostor quality) get the most votes.
 *
 * @param {object} params
 * @param {ImpostorDescription} params.impostor    bake output from ImpostorBaker
 * @param {THREE.WebGLRenderer} params.renderer
 * @returns {{samples: Float32Array, sample_count: number}}
 */
export function read_atlas_to_samples({ impostor, renderer }) {
    const renderer_ = renderer;
    const id = impostor;

    const resolution = id.resolution;
    const frame_count = id.frame_count;
    const frame_size = resolution / frame_count;

    // 1. Read all three MRT slices to CPU as RGBA8.
    const base_bytes     = read_atlas_texture_rgba_u8(renderer_, id.rt.texture[0]); // albedo + alpha
    const geometry_bytes = read_atlas_texture_rgba_u8(renderer_, id.rt.texture[1]); // normal + depth
    const orm_bytes      = read_atlas_texture_rgba_u8(renderer_, id.rt.texture[2]); // ORM

    // 2. Pick the UV encoder that matches the capture type. We only need
    //    uv_to_unit here — the frame index gives us the (u, v) directly.
    let encoder;
    if (id.capture_type === ImpostorCaptureType.FullSphere) {
        encoder = new OctahedralUvEncoder();
    } else if (id.capture_type === ImpostorCaptureType.Hemisphere) {
        encoder = new HemiOctahedralUvEncoder();
    } else {
        throw new Error(`Unsupported capture type '${id.capture_type}'`);
    }

    // 3. Bound the per-sample output to atlas covered pixels. Worst case
    //    every pixel is opaque, but typical opaque coverage is ~20–40%
    //    for hollow geometry. Allocate for the worst case so we never
    //    realloc inside the hot loop.
    const max_pixels = resolution * resolution;
    const samples = new Float32Array(max_pixels * ATLAS_SAMPLE_STRIDE);
    let sample_count = 0;

    // Reused scratch buffers — kept outside the inner loop so the JIT
    // doesn't allocate per pixel.
    const D_frame = [0, 0, 0];
    const pos_scratch = [0, 0, 0];
    const nrm_scratch = [0, 0, 0];

    const max_frame_index = frame_count - 1;
    const inv_255 = 1 / 255;
    const inv_frame_size = 1 / frame_size;

    const centre = id.offset;
    const radius = id.sphere_radius;

    // 4. Walk frames row-major, matching the order the baker used. For
    //    each frame, scan its tile pixel-by-pixel. We rely on
    //    renderSamplerFromTexture having read the atlas back unflipped
    //    so atlas pixel (x, y) corresponds to the same world position
    //    the bake camera wrote there.
    for (let frame_j = 0; frame_j < frame_count; frame_j++) {
        // Re-derive the bake direction for this frame from the same
        // (u, v) the baker used. Don't read it from atlas memory: it's
        // lost in the rasteriser.
        const frame_v = max_frame_index > 0 ? (frame_j / max_frame_index) : 0;
        const frame_offset_y = Math.floor(frame_size * frame_j);

        for (let frame_i = 0; frame_i < frame_count; frame_i++) {
            const frame_u = max_frame_index > 0 ? (frame_i / max_frame_index) : 0;
            const frame_offset_x = Math.floor(frame_size * frame_i);

            encoder.uv_to_unit(D_frame, [frame_u, frame_v]);

            // Per-pixel decode within this frame tile.
            const px_end = Math.ceil(frame_size);
            for (let py = 0; py < px_end; py++) {
                const global_y = frame_offset_y + py;
                if (global_y >= resolution) continue;

                // frame_v_local in [-1, +1], +0.5 to land at pixel centre.
                const frame_v_local = ((py + 0.5) * inv_frame_size) * 2 - 1;

                for (let px = 0; px < px_end; px++) {
                    const global_x = frame_offset_x + px;
                    if (global_x >= resolution) continue;

                    const pixel = (global_y * resolution + global_x) * 4;

                    // Alpha gates the sample. Background pixels (alpha=0)
                    // were never written by the bake shader — discard them.
                    const a = base_bytes[pixel + 3];
                    if (a === 0) {
                        continue;
                    }

                    const frame_u_local = ((px + 0.5) * inv_frame_size) * 2 - 1;

                    const depth01 = geometry_bytes[pixel + 3] * inv_255;
                    reconstruct_position(pos_scratch, depth01, frame_u_local, frame_v_local, D_frame, radius, centre);

                    decode_normal_to_object_space(
                        nrm_scratch,
                        geometry_bytes[pixel] * inv_255,
                        geometry_bytes[pixel + 1] * inv_255,
                        geometry_bytes[pixel + 2] * inv_255,
                        D_frame
                    );

                    const out_offset = sample_count * ATLAS_SAMPLE_STRIDE;

                    samples[out_offset]     = pos_scratch[0];
                    samples[out_offset + 1] = pos_scratch[1];
                    samples[out_offset + 2] = pos_scratch[2];

                    samples[out_offset + 3] = nrm_scratch[0];
                    samples[out_offset + 4] = nrm_scratch[1];
                    samples[out_offset + 5] = nrm_scratch[2];

                    // albedo: decode sRGB -> linear before storing. The
                    // bake atlas holds raw sRGB-encoded bytes (the
                    // GLTF colour textures are sRGB but BakeShaderStandard
                    // is a RawShaderMaterial that samples them without
                    // conversion). Naïvely averaging those bytes in
                    // sRGB space — which is what the voxeliser would
                    // do if we stored them as-is — darkens the result,
                    // compounding the small darkening the bake's own
                    // 4x supersample already introduces. Decoding here
                    // means voxelize_samples averages in linear space
                    // (correct), and the voxeliser re-encodes back to
                    // sRGB at output so the rest of the pipeline still
                    // sees the "sRGB bytes labelled linear" format the
                    // shaders expect.
                    samples[out_offset + 6] = convert_channel_sRGB_to_linear(base_bytes[pixel]     * inv_255);
                    samples[out_offset + 7] = convert_channel_sRGB_to_linear(base_bytes[pixel + 1] * inv_255);
                    samples[out_offset + 8] = convert_channel_sRGB_to_linear(base_bytes[pixel + 2] * inv_255);

                    samples[out_offset + 9]  = orm_bytes[pixel]     * inv_255; // occlusion
                    samples[out_offset + 10] = orm_bytes[pixel + 1] * inv_255; // roughness
                    samples[out_offset + 11] = orm_bytes[pixel + 2] * inv_255; // metalness

                    sample_count++;
                }
            }
        }
    }

    return { samples, sample_count };
}
