onedrop_codegen/

compiler.rs

//! Shader compiler with naga validation

use crate::error::{CodegenError, Result};
use crate::prelude::SHADER_UNIFORMS_WGSL;
use onedrop_hlsl::{MAX_USER_TEXTURE_SLOTS, TextureBindingPlan};
use std::collections::HashMap;
use std::fmt::Write as _;
use std::sync::{Arc, Mutex};

/// Re-export so the renderer can spell the same constant in its bind group
/// builder. Single source of truth: the wrapper declares `MAX_USER_TEXTURE_SLOTS`
/// bindings; the comp pipeline allocates that many entries.
pub const USER_TEXTURE_SLOTS: usize = MAX_USER_TEXTURE_SLOTS;

/// First @binding index used by the user-texture array. The
/// 14 bindings below this are claimed by built-ins (uniforms, main texture,
/// blur, noise, four sampler variants) — see [`USER_COMP_BINDINGS`].
pub const USER_TEXTURE_FIRST_BINDING: u32 = 15;

/// Compiled shader with validated module
#[derive(Clone)]
pub struct CompiledShader {
    /// Original WGSL source
    pub source: String,

    /// Validated naga module
    pub module: Arc<naga::Module>,

    /// Module info for validation
    pub info: Arc<naga::valid::ModuleInfo>,
}

/// Shader compiler with caching
pub struct ShaderCompiler {
    cache: Arc<Mutex<HashMap<String, CompiledShader>>>,
    validator: naga::valid::Validator,
}

impl ShaderCompiler {
    pub fn new() -> Self {
        Self {
            cache: Arc::new(Mutex::new(HashMap::new())),
            validator: naga::valid::Validator::new(
                naga::valid::ValidationFlags::all(),
                naga::valid::Capabilities::all(),
            ),
        }
    }

    /// Compile and validate a WGSL shader
    pub fn compile(&mut self, source: &str) -> Result<CompiledShader> {
        // Check cache first
        {
            let cache = self.cache.lock().unwrap();
            if let Some(compiled) = cache.get(source) {
                log::debug!("Shader cache hit");
                return Ok(compiled.clone());
            }
        }

        log::debug!("Compiling shader ({} bytes)", source.len());

        // Parse WGSL
        let module = naga::front::wgsl::parse_str(source)
            .map_err(|e| CodegenError::Compilation(format!("WGSL parse error: {:?}", e)))?;

        // Validate
        let info = self
            .validator
            .validate(&module)
            .map_err(|e| CodegenError::Compilation(format!("Validation error: {:?}", e)))?;

        let compiled = CompiledShader {
            source: source.to_string(),
            module: Arc::new(module),
            info: Arc::new(info),
        };

        // Cache it
        {
            let mut cache = self.cache.lock().unwrap();
            cache.insert(source.to_string(), compiled.clone());
        }

        log::debug!("Shader compiled and cached successfully");

        Ok(compiled)
    }

    /// Get cache statistics
    pub fn cache_stats(&self) -> CacheStats {
        let cache = self.cache.lock().unwrap();
        CacheStats {
            size: cache.len(),
            total_source_bytes: cache.values().map(|s| s.source.len()).sum(),
        }
    }

    /// Clear the cache
    pub fn clear_cache(&mut self) {
        let mut cache = self.cache.lock().unwrap();
        cache.clear();
        log::debug!("Shader cache cleared");
    }

    /// Translate a MilkDrop 2 comp shader from HLSL, wrap it into a complete
    /// WGSL fragment module (prelude + texture bindings + standard MD2
    /// fragment inputs + entry points), and validate via naga.
    ///
    /// Returns the wrapped WGSL on success — the renderer's `CompPipeline`
    /// can feed it to `wgpu::Device::create_shader_module` directly. On
    /// failure, the error string captures whichever stage tripped (regex
    /// translation, WGSL parse, validation).
    ///
    /// Compile rates depend on the AST-based translator in [`onedrop_hlsl`];
    /// presets the translator can't fully translate fall back to a gamma-only
    /// default shader at the caller level so the render path stays alive.
    pub fn compile_user_comp_shader(&mut self, hlsl: &str) -> Result<CompiledShader> {
        let plan = TextureBindingPlan::empty();
        self.compile_user_comp_shader_with_plan(hlsl, &plan)
    }

    /// Variant of [`Self::compile_user_comp_shader`] that threads a
    /// [`TextureBindingPlan`] through the translator and wrapper. The
    /// renderer scans the HLSL, resolves each `sampler sampler_X;` against
    /// its texture pool, and hands the result over here so the emitted
    /// WGSL routes `tex2D(sampler_clouds, …)` onto the right user-texture
    /// binding and exposes the matching `texsize_clouds` constant inside
    /// `fs_main`.
    pub fn compile_user_comp_shader_with_plan(
        &mut self,
        hlsl: &str,
        plan: &TextureBindingPlan,
    ) -> Result<CompiledShader> {
        let translated = onedrop_hlsl::translate_shader_with_plan(hlsl, plan)
            .map_err(|e| CodegenError::Compilation(format!("HLSL→WGSL translate: {e}")))?;
        let wrapped = wrap_user_comp_shader_with_plan(&translated, plan);
        self.compile(&wrapped)
    }

    /// Compile a user-authored MD2 warp shader (the per-pixel feedback
    /// stage that runs before the composite pass) through the same
    /// HLSL→WGSL translator the comp pipeline uses, then validate via
    /// naga.
    ///
    /// The wrapper at [`wrap_user_warp_shader_with_plan`] mirrors the
    /// comp wrapper's bind-group layout — same `ShaderUniforms`, same
    /// texture binding numbers, same `GetPixel`/`GetBlur*` helpers — so
    /// the renderer can reuse the comp pipeline's texture-plan
    /// machinery. The only structural differences:
    ///
    /// 1. The vertex shader consumes the warp-mesh `WarpVertex`
    ///    (`pos_clip` + `uv_warp`) instead of generating a fullscreen
    ///    triangle. `uv_warp` is the per-vertex MD2-formula-warped UV
    ///    the CPU writes each frame; the rasteriser interpolates it
    ///    across each triangle.
    /// 2. `ret` is seeded from `sampler_main_texture` at the warped
    ///    UV. The renderer is expected to bind `prev_texture` to the
    ///    `sampler_main_texture` slot so MD2's "previous frame at the
    ///    warped UV" semantics fall out for free.
    pub fn compile_user_warp_shader_with_plan(
        &mut self,
        hlsl: &str,
        plan: &TextureBindingPlan,
    ) -> Result<CompiledShader> {
        let translated = onedrop_hlsl::translate_shader_with_plan(hlsl, plan)
            .map_err(|e| CodegenError::Compilation(format!("HLSL→WGSL translate: {e}")))?;
        let wrapped = wrap_user_warp_shader_with_plan(&translated, plan);
        self.compile(&wrapped)
    }

    /// Empty-plan convenience for [`Self::compile_user_warp_shader_with_plan`].
    pub fn compile_user_warp_shader(&mut self, hlsl: &str) -> Result<CompiledShader> {
        let plan = TextureBindingPlan::empty();
        self.compile_user_warp_shader_with_plan(hlsl, &plan)
    }
}

/// Wrap a translated WGSL fragment body into a complete shader module
/// (prelude + texture bindings + entry points).
///
/// The user body is pasted into `fs_main` after standard MD2 inputs (`uv`,
/// `uv_orig`, `rad`, `ang`, `ret`) have been declared as locals.
/// The fragment returns `vec4<f32>(ret * uniforms.gamma_adj, 1.0)`, so a
/// well-behaved MD2 shader that assigns to `ret` produces a correct result.
///
/// Made `pub` so tests and the CLI can inspect the wrapped output before
/// shipping it to naga.
pub fn wrap_user_comp_shader(translated_body: &str) -> String {
    wrap_user_comp_shader_with_plan(translated_body, &TextureBindingPlan::empty())
}

/// Wrap a translated WGSL fragment body for the **warp pass** into a
/// complete shader module. Mirror of [`wrap_user_comp_shader_with_plan`]
/// — same uniforms, same texture bindings, same MD2 private-state
/// declarations — but with a vertex shader that consumes the warp-mesh
/// `WarpVertex` (`pos_clip` + `uv_warp`) and a fragment that seeds `uv`
/// from `uv_warp` and `ret` from the previous-frame texture at that UV.
///
/// The renderer is expected to bind `prev_texture` to the
/// `sampler_main_texture` slot when running the warp pipeline so the
/// MD2 convention "sample the previous frame at the warped UV" lands
/// naturally. Helper functions (`GetPixel` / `GetBlur*`) therefore
/// also read from `prev_texture`'s view.
pub fn wrap_user_warp_shader_with_plan(translated_body: &str, plan: &TextureBindingPlan) -> String {
    let (lifted_fns, body) = match translated_body.split_once(onedrop_hlsl::LIFTED_FN_SENTINEL) {
        Some((lifted, body)) => (lifted, body.trim_start_matches('\n')),
        None => ("", translated_body),
    };

    let mut s = String::with_capacity(
        SHADER_UNIFORMS_WGSL.len() + USER_COMP_HELPERS.len() + translated_body.len() + 1536,
    );
    s.push_str(SHADER_UNIFORMS_WGSL);
    s.push_str(USER_COMP_BINDINGS);
    s.push_str(MD2_MATH_CONSTANTS);
    s.push_str(MD2_NOISE_TEXSIZE_CONSTANTS);
    s.push_str(MD2_PRIVATE_STATE_DECLS);
    s.push_str(USER_COMP_HELPERS);
    if !lifted_fns.trim().is_empty() {
        s.push_str("\n// ---- lifted user-defined functions ----\n");
        s.push_str(lifted_fns);
        s.push('\n');
    }
    s.push_str(USER_WARP_VERTEX);
    s.push_str(USER_WARP_FRAGMENT_PREFIX);
    append_user_texsize_constants(&mut s, plan);
    s.push_str(body);
    s.push_str(USER_WARP_FRAGMENT_SUFFIX);
    s
}

/// Same as [`wrap_user_comp_shader`], but emits per-preset `texsize_<NAME>`
/// constants from the supplied [`TextureBindingPlan`]. The user-texture
/// bindings themselves are always declared (so the bind-group layout the
/// comp pipeline owns stays stable); only the WGSL identifiers the user
/// body can reference (`texsize_clouds`, …) change per preset.
pub fn wrap_user_comp_shader_with_plan(translated_body: &str, plan: &TextureBindingPlan) -> String {
    // `translate_shader` may emit module-scope user functions in a prefix
    // region separated by `LIFTED_FN_SENTINEL`. Split here so the lifted
    // text lands BEFORE `fs_main` (where it belongs) and the body lands
    // INSIDE `fs_main`. When no sentinel is present, the whole translated
    // text is treated as body and the module-scope region stays empty.
    let (lifted_fns, body) = match translated_body.split_once(onedrop_hlsl::LIFTED_FN_SENTINEL) {
        Some((lifted, body)) => (lifted, body.trim_start_matches('\n')),
        None => ("", translated_body),
    };

    let mut s = String::with_capacity(
        SHADER_UNIFORMS_WGSL.len() + USER_COMP_HELPERS.len() + translated_body.len() + 1536,
    );
    s.push_str(SHADER_UNIFORMS_WGSL);
    s.push_str(USER_COMP_BINDINGS);
    s.push_str(MD2_MATH_CONSTANTS);
    s.push_str(MD2_NOISE_TEXSIZE_CONSTANTS);
    s.push_str(MD2_PRIVATE_STATE_DECLS);
    s.push_str(USER_COMP_HELPERS);
    if !lifted_fns.trim().is_empty() {
        s.push_str("\n// ---- lifted user-defined functions ----\n");
        s.push_str(lifted_fns);
        s.push('\n');
    }
    s.push_str(USER_COMP_VERTEX);
    s.push_str(USER_COMP_FRAGMENT_PREFIX);
    append_user_texsize_constants(&mut s, plan);
    s.push_str(body);
    s.push_str(USER_COMP_FRAGMENT_SUFFIX);
    s
}

/// Emit a `let texsize_<NAME>: vec4<f32> = vec4<f32>(w, h, 1/w, 1/h);` line
/// for each filled slot in the plan. Preset code reads these directly:
/// `tex2D(sampler_clouds, uv + texsize_clouds.zw * dt)` is a common pattern
/// for "step one texel along the texture dimension".
fn append_user_texsize_constants(out: &mut String, plan: &TextureBindingPlan) {
    if plan.slot_count() == 0 {
        return;
    }
    out.push_str("    // ---- user texture pack: per-preset texsize constants ----\n");
    for slot in plan.slots() {
        let Some(name) = &slot.pool_name else {
            continue;
        };
        // Pool names are already lowercase-canonical. Skip any that contain
        // characters that aren't valid WGSL identifiers — the codegen
        // wrapper can't reference them otherwise. Real preset texture
        // names in the wild are all ASCII alnum + underscore.
        if !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
            continue;
        }
        let [w, h, rw, rh] = slot.texsize;
        let _ = writeln!(
            out,
            "    let texsize_{name}: vec4<f32> = vec4<f32>({w:.1}, {h:.1}, {rw}, {rh});"
        );
    }
}

const USER_COMP_BINDINGS: &str = r#"
@group(0) @binding(1)
var sampler_main_texture: texture_2d<f32>;

@group(0) @binding(2)
var sampler_main: sampler;

// Blur pyramid. Bindings 3/4/5 are progressively softer Gaussian-
// blurred copies of the warp output, populated each frame by
// `BlurPipeline`. They share `sampler_main` as their sampler — both
// the comp pass and these blur textures want linear filtering.
@group(0) @binding(3)
var sampler_blur1_texture: texture_2d<f32>;

@group(0) @binding(4)
var sampler_blur2_texture: texture_2d<f32>;

@group(0) @binding(5)
var sampler_blur3_texture: texture_2d<f32>;

// Procedural noise pack. Five fixed-resolution textures generated at
// engine init from a deterministic seed. The translator routes
// `tex2D(sampler_noise_lq|mq|hq, …)` and `tex3D(sampler_noisevol_lq|hq, …)`
// onto these bindings; user shaders that reference unknown noise sampler
// names still fall back to `sampler_main`.
@group(0) @binding(6)
var sampler_noise_lq_texture: texture_2d<f32>;

@group(0) @binding(7)
var sampler_noise_mq_texture: texture_2d<f32>;

@group(0) @binding(8)
var sampler_noise_hq_texture: texture_2d<f32>;

@group(0) @binding(9)
var sampler_noisevol_lq_texture: texture_3d<f32>;

@group(0) @binding(10)
var sampler_noisevol_hq_texture: texture_3d<f32>;

// MD2 sampler variants: filter (filtered=f / point=p) × address (wrap=w /
// clamp=c). The translator rewrites `tex2D(sampler_fw_main, uv)` →
// `textureSample(sampler_main_texture, sampler_fw, uv)`, and similarly
// for `_fc_` / `_pw_` / `_pc_`. Authored presets pick the variant per
// sample site; the runtime always has all four available.
@group(0) @binding(11)
var sampler_fw: sampler;

@group(0) @binding(12)
var sampler_fc: sampler;

@group(0) @binding(13)
var sampler_pw: sampler;

@group(0) @binding(14)
var sampler_pc: sampler;

// ---- User texture pack ----
// Disk-loaded textures referenced by `sampler sampler_<NAME>;` declarations
// in the preset's HLSL. The translator routes `tex2D(sampler_<NAME>, uv)`
// onto one of these slots based on the per-preset
// `onedrop_hlsl::TextureBindingPlan`. Unfilled slots get a 1×1 white
// fallback at bind time so the WGSL parses and the comp pass draws cleanly
// even when the texture pool is empty.
@group(0) @binding(15)
var sampler_user_0_texture: texture_2d<f32>;

@group(0) @binding(16)
var sampler_user_1_texture: texture_2d<f32>;

@group(0) @binding(17)
var sampler_user_2_texture: texture_2d<f32>;

@group(0) @binding(18)
var sampler_user_3_texture: texture_2d<f32>;

@group(0) @binding(19)
var sampler_user_4_texture: texture_2d<f32>;

@group(0) @binding(20)
var sampler_user_5_texture: texture_2d<f32>;

@group(0) @binding(21)
var sampler_user_6_texture: texture_2d<f32>;

@group(0) @binding(22)
var sampler_user_7_texture: texture_2d<f32>;

// Previous-frame display for the echo blend. Bound by
// the comp pipeline at binding 23 (immediately after the user-texture
// slots). User shaders can sample this directly if they want to roll
// their own echo; the default pipeline applies the standard MD2 blend
// in `shaders/comp.wgsl`.
@group(0) @binding(23)
var sampler_prev_main_texture: texture_2d<f32>;
"#;

/// MD2 math constants (mirrors projectM's `milkdrop-shaders.h`).
///
/// User shaders reference these by short bare names. `M_PI_2` (11 hits) and
/// `M_INV_PI_2` (8 hits) are the most common in the 200-preset survey;
/// `M_PI` and `M_INV_PI` are added for completeness — they cost nothing.
const MD2_MATH_CONSTANTS: &str = r#"
const M_PI: f32 = 3.14159265358979;
const M_PI_2: f32 = 1.57079632679489;
const M_INV_PI: f32 = 0.31830988618379;
const M_INV_PI_2: f32 = 0.15915494309189;
"#;

/// Procedural noise pack `texsize_*` vec4 constants. Module-scope (was
/// previously emitted inside `fs_main`) so lifted user functions that
/// reference `texsize_noise_lq.zw` for a per-texel offset compile cleanly.
/// The values are fixed at engine init (see `onedrop-renderer::noise`), so
/// promoting them from `let` to `const` is a free win.
const MD2_NOISE_TEXSIZE_CONSTANTS: &str = r#"
const texsize_noise_lq: vec4<f32>      = vec4<f32>(256.0, 256.0, 0.00390625, 0.00390625);
const texsize_noise_lq_lite: vec4<f32> = vec4<f32>( 32.0,  32.0, 0.03125,    0.03125);
const texsize_noise_mq: vec4<f32>      = vec4<f32>( 64.0,  64.0, 0.015625,   0.015625);
const texsize_noise_hq: vec4<f32>      = vec4<f32>( 32.0,  32.0, 0.03125,    0.03125);
const texsize_noisevol_lq: vec4<f32>   = vec4<f32>( 32.0,  32.0, 0.03125,    0.03125);
const texsize_noisevol_hq: vec4<f32>   = vec4<f32>(  8.0,   8.0, 0.125,      0.125);
"#;

/// MD2 per-invocation private state hoisted to module scope.
///
/// Every binding a user shader can reference by its bare MD2 name
/// (`q1`..`q32`, `texsize`, `time`, `bass`, `slow_roam_cos`, `uv`, `rad`,
/// `ang`, `ret`, …) is declared here as `var<private>`. That makes them
/// visible to lifted user-defined functions (which live at module scope
/// alongside these declarations) — the previous design pinned them as
/// `let` bindings inside `fs_main`, so a lifted `float3 lavcol(float t)
/// { return pow(.5, 2*t*slow_roam_cos); }` would fail naga parse with
/// "no definition in scope for identifier: `slow_roam_cos`".
///
/// WGSL `var<private>` is per-invocation, so each fragment thread has its
/// own copy — assigning these at the top of `fs_main` from `uniforms.*` /
/// per-pixel values is exactly what the original `let` site did, just one
/// scope higher.
const MD2_PRIVATE_STATE_DECLS: &str = r#"
// ---- MD2 wrapper private state (per-invocation, module scope) ----
// fs_main + every lifted user function reads/writes these. Values are
// assigned at the top of fs_main; reading them before assignment is
// undefined behavior, but every code path in our wrapper writes first.
var<private> uv: vec2<f32>;
var<private> uv_orig: vec2<f32>;
var<private> rad: f32;
var<private> ang: f32;
var<private> ret: vec3<f32>;
// `_md2_color` is a wrapper scratch slot — kept under a `_md2_`-prefixed
// name to avoid colliding with the very common user-declared `color`.
var<private> _md2_color: vec3<f32>;

var<private> texsize: vec4<f32>;
var<private> aspect: vec4<f32>;
var<private> time: f32;
var<private> fps: f32;
var<private> frame: f32;
var<private> progress: f32;
var<private> bass: f32;
var<private> mid: f32;
var<private> treb: f32;
var<private> vol: f32;
var<private> bass_att: f32;
var<private> mid_att: f32;
var<private> treb_att: f32;
var<private> vol_att: f32;
var<private> rand_preset: vec4<f32>;
var<private> rand_frame: vec4<f32>;
var<private> slow_roam_cos: vec4<f32>;
var<private> slow_roam_sin: vec4<f32>;
var<private> roam_cos: vec4<f32>;
var<private> roam_sin: vec4<f32>;
var<private> blur1_min: f32;
var<private> blur1_max: f32;
var<private> blur2_min: f32;
var<private> blur2_max: f32;
var<private> blur3_min: f32;
var<private> blur3_max: f32;
var<private> hue_shader: vec3<f32>;
// MD1 legacy alias for `texsize` (some old presets read `g_fTexSize.zw`).
var<private> g_fTexSize: vec4<f32>;

var<private> q1: f32;
var<private> q2: f32;
var<private> q3: f32;
var<private> q4: f32;
var<private> q5: f32;
var<private> q6: f32;
var<private> q7: f32;
var<private> q8: f32;
var<private> q9: f32;
var<private> q10: f32;
var<private> q11: f32;
var<private> q12: f32;
var<private> q13: f32;
var<private> q14: f32;
var<private> q15: f32;
var<private> q16: f32;
var<private> q17: f32;
var<private> q18: f32;
var<private> q19: f32;
var<private> q20: f32;
var<private> q21: f32;
var<private> q22: f32;
var<private> q23: f32;
var<private> q24: f32;
var<private> q25: f32;
var<private> q26: f32;
var<private> q27: f32;
var<private> q28: f32;
var<private> q29: f32;
var<private> q30: f32;
var<private> q31: f32;
var<private> q32: f32;
"#;

/// MD2 helper functions every comp shader expects to find at module scope.
///
/// `GetPixel(uv)` returns the current main texture sample as `vec3<f32>` —
/// the canonical MD2 convention is to discard alpha at the comp stage.
///
/// `GetBlur1/2/3(uv)` sample the cumulative Gaussian-blur pyramid produced
/// by [`BlurPipeline`] each frame. The vast majority of in-the-wild comp
/// shaders reference at least one `GetBlur*` — without a real pyramid they
/// would degrade to `GetPixel` and lose the soft-halo look the preset
/// depends on.
const USER_COMP_HELPERS: &str = r#"
fn GetPixel(uv: vec2<f32>) -> vec3<f32> {
    return textureSample(sampler_main_texture, sampler_main, uv).rgb;
}

// MD2 alias — many in-the-wild comp shaders call GetMain(uv) instead of
// GetPixel(uv). Same semantics: sample the main render target.
fn GetMain(uv: vec2<f32>) -> vec3<f32> {
    return textureSample(sampler_main_texture, sampler_main, uv).rgb;
}

// Blur-pyramid samplers. MD2 stores the three Gaussian-blurred copies of
// the warp output in a compressed range `[blurN_min, blurN_max]` (the
// engine writes these uniforms each frame); the canonical `GetBlurN(uv)`
// helper re-expands that range to `[0, 1]` before handing it to the user
// shader. Without this remap, presets that depend on the implicit clamp
// (`color * GetBlur1(uv)`) saturate near full-white or full-black even on
// quiet content, throwing the visual balance MD2 is tuned for.
fn GetBlur1(uv: vec2<f32>) -> vec3<f32> {
    let sample = textureSample(sampler_blur1_texture, sampler_main, uv).rgb;
    return mix(vec3<f32>(blur1_min), vec3<f32>(blur1_max), sample);
}

fn GetBlur2(uv: vec2<f32>) -> vec3<f32> {
    let sample = textureSample(sampler_blur2_texture, sampler_main, uv).rgb;
    return mix(vec3<f32>(blur2_min), vec3<f32>(blur2_max), sample);
}

fn GetBlur3(uv: vec2<f32>) -> vec3<f32> {
    let sample = textureSample(sampler_blur3_texture, sampler_main, uv).rgb;
    return mix(vec3<f32>(blur3_min), vec3<f32>(blur3_max), sample);
}

/// MD2 luminance helper. 78 / 168 in-the-wild comp shaders call `lum(c)`
/// to weight the perceptual brightness of a sample. The MD2 source uses a
/// Rec. 601-ish weight `(0.32, 0.49, 0.29)` (note: not exactly normalised);
/// we mirror the weights so visual output stays close to the reference.
fn lum(c: vec3<f32>) -> f32 {
    return dot(c, vec3<f32>(0.32, 0.49, 0.29));
}
"#;

const USER_COMP_VERTEX: &str = r#"
struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) uv: vec2<f32>,
}

@vertex
fn vs_main(@builtin(vertex_index) idx: u32) -> VertexOutput {
    let x = f32(idx & 1u) * 4.0 - 1.0;
    let y = f32((idx >> 1u) & 1u) * 4.0 - 1.0;
    var out: VertexOutput;
    out.position = vec4<f32>(x, y, 0.0, 1.0);
    out.uv = vec2<f32>((x + 1.0) * 0.5, 1.0 - (y + 1.0) * 0.5);
    return out;
}
"#;

const USER_COMP_FRAGMENT_PREFIX: &str = r#"
@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    // All MD2 bare-name bindings live at module scope (`var<private>` in
    // `MD2_PRIVATE_STATE_DECLS`) so lifted user-defined functions can read
    // them. Initialise from the per-invocation inputs and the uniform
    // struct here — WGSL `var<private>` has no static initialiser, so the
    // first thing every fragment thread does is seed its private state.
    uv = in.uv;
    uv_orig = in.uv;
    rad = length(uv - vec2<f32>(0.5)) * 1.4142135;
    ang = atan2(uv.y - 0.5, uv.x - 0.5);
    ret = textureSample(sampler_main_texture, sampler_main, uv).rgb;
    _md2_color = ret;

    texsize = uniforms.texsize;
    aspect = uniforms.aspect;
    time = uniforms.time;
    fps = uniforms.fps;
    frame = uniforms.frame;
    progress = uniforms.progress;
    bass = uniforms.bass;
    mid = uniforms.mid;
    treb = uniforms.treb;
    vol = uniforms.vol;
    bass_att = uniforms.bass_att;
    mid_att = uniforms.mid_att;
    treb_att = uniforms.treb_att;
    vol_att = uniforms.vol_att;
    rand_preset = uniforms.rand_preset;
    rand_frame = uniforms.rand_frame;
    slow_roam_cos = uniforms.slow_roam_cos;
    slow_roam_sin = uniforms.slow_roam_sin;
    roam_cos = uniforms.roam_cos;
    roam_sin = uniforms.roam_sin;
    blur1_min = uniforms.blur1_min;
    blur1_max = uniforms.blur1_max;
    blur2_min = uniforms.blur2_min;
    blur2_max = uniforms.blur2_max;
    blur3_min = uniforms.blur3_min;
    blur3_max = uniforms.blur3_max;
    // `hue_shader` is the per-frame RGB tint MD2 cycles through over time.
    // The engine writes a time-driven tri-phase oscillator into
    // `uniforms.hue_shader.xyz` each frame; here we just splice it in so
    // user shaders that multiply `ret * hue_shader` get the slow rainbow
    // rotation MD2 is tuned for.
    hue_shader = uniforms.hue_shader.xyz;
    g_fTexSize = uniforms.texsize;

    q1 = uniforms.q[0].x;
    q2 = uniforms.q[0].y;
    q3 = uniforms.q[0].z;
    q4 = uniforms.q[0].w;
    q5 = uniforms.q[1].x;
    q6 = uniforms.q[1].y;
    q7 = uniforms.q[1].z;
    q8 = uniforms.q[1].w;
    q9 = uniforms.q[2].x;
    q10 = uniforms.q[2].y;
    q11 = uniforms.q[2].z;
    q12 = uniforms.q[2].w;
    q13 = uniforms.q[3].x;
    q14 = uniforms.q[3].y;
    q15 = uniforms.q[3].z;
    q16 = uniforms.q[3].w;
    q17 = uniforms.q[4].x;
    q18 = uniforms.q[4].y;
    q19 = uniforms.q[4].z;
    q20 = uniforms.q[4].w;
    q21 = uniforms.q[5].x;
    q22 = uniforms.q[5].y;
    q23 = uniforms.q[5].z;
    q24 = uniforms.q[5].w;
    q25 = uniforms.q[6].x;
    q26 = uniforms.q[6].y;
    q27 = uniforms.q[6].z;
    q28 = uniforms.q[6].w;
    q29 = uniforms.q[7].x;
    q30 = uniforms.q[7].y;
    q31 = uniforms.q[7].z;
    q32 = uniforms.q[7].w;

    // ---- begin translated user body ----
"#;

const USER_COMP_FRAGMENT_SUFFIX: &str = r#"
    // ---- end translated user body ----

    return vec4<f32>(ret * uniforms.gamma_adj, 1.0);
}
"#;

/// Warp-pass vertex shader. Consumes the same `WarpVertex` layout the
/// renderer uses for the default warp pipeline (`pos_clip: vec2<f32>`
/// at `@location(0)` plus `uv_warp: vec2<f32>` at `@location(1)`) so
/// switching between the default and user pipelines doesn't change
/// vertex buffer plumbing on the renderer side.
///
/// Outputs `uv` (the per-pixel rasterised warp UV) and `uv_screen` (the
/// screen-space `[0, 1]` UV recovered from `pos_clip`, kept for parity
/// with `warp.wgsl`'s built-in shaders).
const USER_WARP_VERTEX: &str = r#"
struct VertexOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) uv: vec2<f32>,
    @location(1) uv_screen: vec2<f32>,
}

@vertex
fn vs_main(
    @location(0) pos_clip: vec2<f32>,
    @location(1) uv_warp: vec2<f32>,
) -> VertexOutput {
    var out: VertexOutput;
    out.position = vec4<f32>(pos_clip, 0.0, 1.0);
    out.uv = uv_warp;
    out.uv_screen = pos_clip * 0.5 + vec2<f32>(0.5);
    return out;
}
"#;

/// Fragment-shader header for user-authored warp shaders. Same shape
/// as `USER_COMP_FRAGMENT_PREFIX` (seed every MD2 `var<private>` binding
/// from `uniforms` and per-invocation inputs), but `uv`/`uv_orig` come
/// from the rasterised warp UV — not from a fullscreen-triangle screen
/// UV — and `ret` is seeded from the previous-frame texture at that
/// warped UV. The renderer binds `prev_texture` to `sampler_main_texture`
/// so `GetPixel(uv)` (defined in the shared helpers) does the right thing.
const USER_WARP_FRAGMENT_PREFIX: &str = r#"
@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    uv = in.uv;
    // `uv_orig` is the *pre-warp* screen-space UV. Warp HLSL authors
    // sometimes reference it to compare the displaced sample against
    // the pixel's own coordinate (radial fades, vignette terms).
    uv_orig = in.uv_screen;
    rad = length(in.uv_screen - vec2<f32>(0.5)) * 1.4142135;
    ang = atan2(in.uv_screen.y - 0.5, in.uv_screen.x - 0.5);
    ret = textureSample(sampler_main_texture, sampler_main, uv).rgb;
    _md2_color = ret;

    texsize = uniforms.texsize;
    aspect = uniforms.aspect;
    time = uniforms.time;
    fps = uniforms.fps;
    frame = uniforms.frame;
    progress = uniforms.progress;
    bass = uniforms.bass;
    mid = uniforms.mid;
    treb = uniforms.treb;
    vol = uniforms.vol;
    bass_att = uniforms.bass_att;
    mid_att = uniforms.mid_att;
    treb_att = uniforms.treb_att;
    vol_att = uniforms.vol_att;
    rand_preset = uniforms.rand_preset;
    rand_frame = uniforms.rand_frame;
    slow_roam_cos = uniforms.slow_roam_cos;
    slow_roam_sin = uniforms.slow_roam_sin;
    roam_cos = uniforms.roam_cos;
    roam_sin = uniforms.roam_sin;
    blur1_min = uniforms.blur1_min;
    blur1_max = uniforms.blur1_max;
    blur2_min = uniforms.blur2_min;
    blur2_max = uniforms.blur2_max;
    blur3_min = uniforms.blur3_min;
    blur3_max = uniforms.blur3_max;
    hue_shader = vec3<f32>(1.0, 1.0, 1.0);
    g_fTexSize = uniforms.texsize;

    q1 = uniforms.q[0].x;
    q2 = uniforms.q[0].y;
    q3 = uniforms.q[0].z;
    q4 = uniforms.q[0].w;
    q5 = uniforms.q[1].x;
    q6 = uniforms.q[1].y;
    q7 = uniforms.q[1].z;
    q8 = uniforms.q[1].w;
    q9 = uniforms.q[2].x;
    q10 = uniforms.q[2].y;
    q11 = uniforms.q[2].z;
    q12 = uniforms.q[2].w;
    q13 = uniforms.q[3].x;
    q14 = uniforms.q[3].y;
    q15 = uniforms.q[3].z;
    q16 = uniforms.q[3].w;
    q17 = uniforms.q[4].x;
    q18 = uniforms.q[4].y;
    q19 = uniforms.q[4].z;
    q20 = uniforms.q[4].w;
    q21 = uniforms.q[5].x;
    q22 = uniforms.q[5].y;
    q23 = uniforms.q[5].z;
    q24 = uniforms.q[5].w;
    q25 = uniforms.q[6].x;
    q26 = uniforms.q[6].y;
    q27 = uniforms.q[6].z;
    q28 = uniforms.q[6].w;
    q29 = uniforms.q[7].x;
    q30 = uniforms.q[7].y;
    q31 = uniforms.q[7].z;
    q32 = uniforms.q[7].w;

    // ---- begin translated user body ----
"#;

/// Fragment-shader trailer for the warp pass. Same convention as the
/// comp suffix (return `vec4(ret, 1.0)`), minus the `gamma_adj` factor —
/// gamma is a display-side concern applied later in the comp pass.
const USER_WARP_FRAGMENT_SUFFIX: &str = r#"
    // ---- end translated user body ----

    return vec4<f32>(ret, 1.0);
}
"#;

impl Default for ShaderCompiler {
    fn default() -> Self {
        Self::new()
    }
}

/// Cache statistics
#[derive(Debug, Clone)]
pub struct CacheStats {
    pub size: usize,
    pub total_source_bytes: usize,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_compile_simple_shader() {
        let mut compiler = ShaderCompiler::new();

        let source = r#"
@vertex
fn vs_main(@builtin(vertex_index) in_vertex_index: u32) -> @builtin(position) vec4<f32> {
    return vec4<f32>(0.0, 0.0, 0.0, 1.0);
}

@fragment
fn fs_main() -> @location(0) vec4<f32> {
    return vec4<f32>(1.0, 0.0, 0.0, 1.0);
}
"#;

        let result = compiler.compile(source);
        assert!(result.is_ok());
    }

    #[test]
    fn test_cache_hit() {
        let mut compiler = ShaderCompiler::new();

        let source = r#"
@fragment
fn fs_main() -> @location(0) vec4<f32> {
    return vec4<f32>(1.0, 0.0, 0.0, 1.0);
}
"#;

        // First compile
        let result1 = compiler.compile(source);
        assert!(result1.is_ok());

        // Second compile (should hit cache)
        let result2 = compiler.compile(source);
        assert!(result2.is_ok());

        let stats = compiler.cache_stats();
        assert_eq!(stats.size, 1);
    }

    #[test]
    fn test_invalid_shader() {
        let mut compiler = ShaderCompiler::new();

        let source = "invalid shader code";

        let result = compiler.compile(source);
        assert!(result.is_err());
    }

    #[test]
    fn empty_user_body_compiles() {
        // The empty body case validates the wrapper itself: the prelude,
        // bindings, vertex shader, and the trivial `ret = sample(main)`
        // fragment must all parse and validate without any user code.
        let mut compiler = ShaderCompiler::new();
        let result = compiler.compile_user_comp_shader("");
        assert!(
            result.is_ok(),
            "empty wrapper should compile, got: {:?}",
            result.err()
        );
    }

    #[test]
    fn trivial_user_body_compiles() {
        // A minimal MD2-style body: assign sampled color back to `ret`.
        // After translation (`float3 → vec3<f32>`, no other surprises) the
        // wrapper produces valid WGSL.
        let mut compiler = ShaderCompiler::new();
        let body = "ret = ret * 0.5;";
        let result = compiler.compile_user_comp_shader(body);
        assert!(result.is_ok(), "got: {:?}", result.err());
    }

    #[test]
    fn invalid_user_body_returns_error() {
        // Garbage in the user body must surface a Compilation error rather
        // than panic.
        let mut compiler = ShaderCompiler::new();
        let body = "this is not even close to HLSL @@@";
        let result = compiler.compile_user_comp_shader(body);
        assert!(result.is_err());
    }

    #[test]
    fn empty_warp_body_compiles() {
        // The warp wrapper is structurally the comp wrapper minus the
        // fullscreen-triangle vs_main, plus a `WarpVertex`-consuming
        // entry point. An empty body still has to parse + validate as a
        // self-contained module (vertex + fragment, same bindings as
        // comp, `ret` defaulted to the previous-frame sample).
        let mut compiler = ShaderCompiler::new();
        let result = compiler.compile_user_warp_shader("");
        assert!(
            result.is_ok(),
            "empty warp wrapper should compile, got: {:?}",
            result.err()
        );
    }

    #[test]
    fn trivial_warp_body_compiles() {
        // The canonical MD2 warp pattern from corpus preset 427.milk:
        //     ret = tex2D(sampler_main, uv).xyz * 0.85;
        //     ret -= 0.022;
        // Exercises every layer (translator, type-aware passes, wrapper)
        // and proves the warp wrap accepts the same HLSL idioms the
        // comp wrap already handles.
        let mut compiler = ShaderCompiler::new();
        let body = "ret = tex2D(sampler_main, uv).xyz * 0.85;\nret -= 0.022;";
        let result = compiler.compile_user_warp_shader(body);
        assert!(result.is_ok(), "got: {:?}", result.err());
    }

    #[test]
    fn warp_wrapper_consumes_warpvertex_layout() {
        // The renderer ships `pos_clip` at @location(0) and `uv_warp` at
        // @location(1); the wrapper must consume the same layout so the
        // default and user warp pipelines can share a vertex buffer.
        let wrapped =
            wrap_user_warp_shader_with_plan("", &onedrop_hlsl::TextureBindingPlan::empty());
        assert!(
            wrapped.contains("@location(0) pos_clip: vec2<f32>"),
            "warp vertex shader must read pos_clip at @location(0)"
        );
        assert!(
            wrapped.contains("@location(1) uv_warp: vec2<f32>"),
            "warp vertex shader must read uv_warp at @location(1)"
        );
        // Differs from the comp wrapper: no `gamma_adj` factor on the
        // final return (gamma is applied later in the comp pass).
        assert!(
            !wrapped.contains("ret * uniforms.gamma_adj"),
            "warp output must not apply gamma_adj"
        );
        assert!(
            wrapped.contains("return vec4<f32>(ret, 1.0)"),
            "warp output must return ret unchanged"
        );
    }

    #[test]
    fn wrapper_includes_prelude_and_bindings() {
        let wrapped = wrap_user_comp_shader("// body");
        assert!(wrapped.contains("struct ShaderUniforms"), "no prelude");
        assert!(
            wrapped.contains("sampler_main_texture"),
            "no main texture binding"
        );
        assert!(wrapped.contains("@vertex"), "no vertex stage");
        assert!(wrapped.contains("@fragment"), "no fragment stage");
        assert!(wrapped.contains("// body"), "user body not pasted");
        // The wrapper always declares the user-texture slots,
        // even with an empty plan — the bind-group layout the renderer
        // owns is fixed-size, so the WGSL must reference all of them or
        // wgpu rejects the pipeline.
        for slot in 0..USER_TEXTURE_SLOTS {
            let needle = format!("sampler_user_{slot}_texture");
            assert!(
                wrapped.contains(&needle),
                "user texture binding {slot} missing from wrapper"
            );
        }
    }

    #[test]
    fn wrapper_with_plan_emits_texsize_constants() {
        let mut plan = TextureBindingPlan::empty();
        plan.add_slot(
            Some("clouds".to_string()),
            [256.0, 128.0, 1.0 / 256.0, 1.0 / 128.0],
            &[("sampler_clouds".to_string(), "sampler_fw")],
        )
        .unwrap();
        plan.add_slot(
            Some("worms".to_string()),
            [64.0, 64.0, 1.0 / 64.0, 1.0 / 64.0],
            &[("sampler_worms".to_string(), "sampler_fw")],
        )
        .unwrap();

        let wrapped = wrap_user_comp_shader_with_plan("// body", &plan);
        assert!(
            wrapped.contains("let texsize_clouds: vec4<f32> = vec4<f32>(256.0, 128.0,"),
            "got: {wrapped}"
        );
        assert!(
            wrapped.contains("let texsize_worms: vec4<f32> = vec4<f32>(64.0, 64.0,"),
            "got: {wrapped}"
        );
    }

    #[test]
    fn user_shader_referencing_plan_texture_compiles() {
        // End-to-end: a real MD2-style body that samples `sampler_clouds`
        // must translate, validate, and use the plan-routed binding.
        let mut compiler = ShaderCompiler::new();
        let mut plan = TextureBindingPlan::empty();
        plan.add_slot(
            Some("clouds".to_string()),
            [256.0, 256.0, 1.0 / 256.0, 1.0 / 256.0],
            &[("sampler_clouds".to_string(), "sampler_fw")],
        )
        .unwrap();
        let hlsl = "ret = tex2D(sampler_clouds, uv).xyz;";
        let compiled = compiler
            .compile_user_comp_shader_with_plan(hlsl, &plan)
            .expect("plan-driven user shader must validate");
        // The compiled source carries the user-binding routing, not the
        // sampler_main fallback.
        assert!(compiled.source.contains("sampler_user_0_texture"));
        assert!(!compiled.source.contains("/*was: sampler_clouds*/"));
    }

    #[test]
    fn test_clear_cache() {
        let mut compiler = ShaderCompiler::new();

        let source = r#"
@fragment
fn fs_main() -> @location(0) vec4<f32> {
    return vec4<f32>(1.0, 0.0, 0.0, 1.0);
}
"#;

        compiler.compile(source).unwrap();
        assert_eq!(compiler.cache_stats().size, 1);

        compiler.clear_cache();
        assert_eq!(compiler.cache_stats().size, 0);
    }
}