goldy 0.2.0

//! Source-to-source transform for Goldy virtual entry points.
//!
//! Shader authors write entry points annotated with `[goldy_compute]`,
//! `[goldy_vertex]`, or `[goldy_fragment]` using ergonomic typed parameters
//! (`Scattered<T>`, `ThreadId`, etc.).  This module transforms those into
//! standard `[shader("...")]` entry points that the Slang compiler accepts,
//! with a fixed 16-word signature mapping to the 128-byte `PushLayout`.
//!
//! # Push constant layout
//!
//! ```text
//! Words  0– 7: _bw0.._bw7  (16 × u16 bindless indices packed 2-per-word, bytes 0–31)
//! Words  8–15: _uw0.._uw7  (8 × u32 user scalar params full-precision, bytes 32–63)
//! Bytes 64–127: reserved
//! ```
//!
//! # Before / After (compute example)
//!
//! ```slang
//! // Input written by the shader author:
//! [goldy_compute]
//! [numthreads(64, 1, 1)]
//! void cs_main(TimeUniforms cfg, Scattered<uint> data, ThreadId id) {
//!     data[id.x] = data[id.x] + cfg.base;
//! }
//!
//! // Generated by this module (prepended to the source):
//! [shader("compute")]
//! [numthreads(64, 1, 1)]
//! void cs_main(uniform uint _bw0, ..., uniform uint _bw7,
//!              uniform uint _uw0, ..., uniform uint _uw7,
//!              uint3 _sv0 : SV_DispatchThreadID) {
//!     TimeUniforms cfg = goldy_broadcast<TimeUniforms>(_bw0 & 0xFFFFu);
//!     Scattered<uint> data = goldy_scattered<uint>((_bw0 >> 16u) & 0xFFFFu);
//!     ThreadId id = ThreadId(_sv0);
//!     _goldy_user_cs_main(cfg, data, id);
//! }
//! ```

use std::borrow::Cow;
use std::ops::Range;

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Extract the inner element type name for each bindless slot from the
/// `[goldy_*]` entry points in `source`.
///
/// For `Scattered<T>` / `BufRO<T>` → `Some("T")`, for `ByteAddress` →
/// `None`, for broadcasts → `Some(struct_name)`, for textures/samplers/SVs
/// → `None`.  Fragment entries take precedence, matching [`extract_push_constant_categories`].
pub fn extract_binding_element_type_names(source: &str) -> Vec<Option<String>> {
    let entries = find_all_entries(source);
    if entries.is_empty() {
        return Vec::new();
    }

    fn type_name_for_param(param: &Param) -> Option<String> {
        match &param.kind {
            ParamKind::Resource => {
                let ty = param.ty.trim();
                fn inner(ty: &str, prefix: &str) -> Option<String> {
                    if ty.starts_with(prefix) && ty.ends_with('>') {
                        Some(ty[prefix.len()..ty.len() - 1].to_string())
                    } else {
                        None
                    }
                }
                inner(ty, "Scattered<")
                    .or_else(|| inner(ty, "BufRO<"))
                    .or_else(|| inner(ty, "Interpolated<"))
                    .or_else(|| inner(ty, "DirectSpatial<"))
            }
            ParamKind::Broadcast => Some(param.ty.clone()),
            _ => None,
        }
    }

    fn extract_from_params(params: &[ParamItem]) -> Vec<Option<String>> {
        let mut names: Vec<Option<String>> = Vec::new();
        for item in params {
            match item {
                ParamItem::Single(p) => {
                    if matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast) {
                        names.push(type_name_for_param(p));
                    }
                }
                ParamItem::Conditional {
                    then_params,
                    else_params,
                    ..
                } => {
                    let then_count = then_params
                        .iter()
                        .filter(|p| matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast))
                        .count();
                    let else_count = else_params
                        .iter()
                        .filter(|p| matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast))
                        .count();
                    let max_slots = then_count.max(else_count);
                    for i in 0..max_slots {
                        let then_name = then_params
                            .iter()
                            .filter(|p| matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast))
                            .nth(i)
                            .and_then(type_name_for_param);
                        let else_name = else_params
                            .iter()
                            .filter(|p| matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast))
                            .nth(i)
                            .and_then(type_name_for_param);
                        if then_name == else_name {
                            names.push(then_name);
                        } else {
                            names.push(None);
                        }
                    }
                }
            }
        }
        names
    }

    let mut fragment_names: Option<Vec<Option<String>>> = None;
    let mut fallback_names: Option<Vec<Option<String>>> = None;

    for entry in &entries {
        let names = extract_from_params(&entry.params);
        if !names.is_empty() {
            if entry.stage == Stage::Fragment {
                fragment_names = Some(names);
            } else if fallback_names.is_none() || entry.stage == Stage::Compute {
                fallback_names = Some(names);
            }
        }
    }

    fragment_names.or(fallback_names).unwrap_or_default()
}

/// Extract per-bindless-slot [`ResourceCategory`](crate::types::ResourceCategory) expectations from the
/// `[goldy_*]` entry points in `source`.
///
/// Returns a vector indexed by bindless slot. Each element is `Some(category)`
/// when the slot's type is unambiguous, or `None` when the source has no
/// `[goldy_*]` annotations.  Multiple entry points (vertex + fragment) are
/// merged: fragment takes precedence for shared slots.
pub fn extract_push_constant_categories(source: &str) -> Vec<Option<crate::types::ResourceCategory>> {
    use crate::types::ResourceCategory;

    let entries = find_all_entries(source);
    if entries.is_empty() {
        return Vec::new();
    }

    fn category_for_param(param: &Param) -> Option<ResourceCategory> {
        match &param.kind {
            ParamKind::Resource => {
                let ty = param.ty.trim();
                if ty.starts_with("Scattered<") || ty.starts_with("BufRO<") || ty == "ByteAddress" {
                    Some(ResourceCategory::Scattered)
                } else if ty.starts_with("Interpolated<") {
                    Some(ResourceCategory::Texture)
                } else if ty.starts_with("DirectSpatial<") {
                    Some(ResourceCategory::StorageImage)
                } else if ty == "Filter" {
                    Some(ResourceCategory::Sampler)
                } else {
                    None
                }
            }
            ParamKind::Broadcast => Some(ResourceCategory::Broadcast),
            _ => None,
        }
    }

    fn extract_from_params(params: &[ParamItem]) -> Vec<Option<ResourceCategory>> {
        let mut cats: Vec<Option<ResourceCategory>> = Vec::new();
        for item in params {
            match item {
                ParamItem::Single(p) => {
                    if matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast) {
                        cats.push(category_for_param(p));
                    }
                }
                ParamItem::Conditional {
                    then_params,
                    else_params,
                    ..
                } => {
                    let then_count = then_params
                        .iter()
                        .filter(|p| matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast))
                        .count();
                    let else_count = else_params
                        .iter()
                        .filter(|p| matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast))
                        .count();
                    let max_slots = then_count.max(else_count);
                    for i in 0..max_slots {
                        let then_cat = then_params
                            .iter()
                            .filter(|p| matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast))
                            .nth(i)
                            .and_then(category_for_param);
                        let else_cat = else_params
                            .iter()
                            .filter(|p| matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast))
                            .nth(i)
                            .and_then(category_for_param);
                        if then_cat == else_cat {
                            cats.push(then_cat);
                        } else {
                            cats.push(None);
                        }
                    }
                }
            }
        }
        cats
    }

    // Fragment shader expectations take precedence (resource binding typically lives there).
    let mut fragment_cats: Option<Vec<Option<ResourceCategory>>> = None;
    let mut fallback_cats: Option<Vec<Option<ResourceCategory>>> = None;

    for entry in &entries {
        let cats = extract_from_params(&entry.params);
        if !cats.is_empty() {
            if entry.stage == Stage::Fragment {
                fragment_cats = Some(cats);
            } else if fallback_cats.is_none() || entry.stage == Stage::Compute {
                fallback_cats = Some(cats);
            }
        }
    }

    fragment_cats.or(fallback_cats).unwrap_or_default()
}

/// Extract per-slot DX12 bindless view expectations (`UAV` vs `SRV`) from `[goldy_*]` entry points.
///
/// Merging rules match [`extract_push_constant_categories`].
#[cfg(all(feature = "dx12", target_os = "windows"))]
pub(crate) fn extract_push_constant_slot_kinds(source: &str) -> Vec<Option<crate::types::BindlessSlotKind>> {
    use crate::types::BindlessSlotKind;

    let entries = find_all_entries(source);
    if entries.is_empty() {
        return Vec::new();
    }

    fn slot_kind_for_param(param: &Param) -> Option<BindlessSlotKind> {
        match &param.kind {
            ParamKind::Resource => {
                let ty = param.ty.trim();
                if ty.starts_with("Scattered<") || ty == "ByteAddress" {
                    Some(BindlessSlotKind::StorageUav)
                } else if ty.starts_with("BufRO<") {
                    Some(BindlessSlotKind::ReadOnlySrv)
                } else {
                    None
                }
            }
            ParamKind::Broadcast => Some(BindlessSlotKind::UniformCbv),
            _ => None,
        }
    }

    fn extract_from_params(params: &[ParamItem]) -> Vec<Option<BindlessSlotKind>> {
        let mut kinds: Vec<Option<BindlessSlotKind>> = Vec::new();
        for item in params {
            match item {
                ParamItem::Single(p) => {
                    if matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast) {
                        kinds.push(slot_kind_for_param(p));
                    }
                }
                ParamItem::Conditional {
                    then_params,
                    else_params,
                    ..
                } => {
                    let then_count = then_params
                        .iter()
                        .filter(|p| matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast))
                        .count();
                    let else_count = else_params
                        .iter()
                        .filter(|p| matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast))
                        .count();
                    let max_slots = then_count.max(else_count);
                    for i in 0..max_slots {
                        let then_kind = then_params
                            .iter()
                            .filter(|p| matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast))
                            .nth(i)
                            .and_then(slot_kind_for_param);
                        let else_kind = else_params
                            .iter()
                            .filter(|p| matches!(p.kind, ParamKind::Resource | ParamKind::Broadcast))
                            .nth(i)
                            .and_then(slot_kind_for_param);
                        if then_kind == else_kind {
                            kinds.push(then_kind);
                        } else {
                            kinds.push(None);
                        }
                    }
                }
            }
        }
        kinds
    }

    let mut fragment_kinds: Option<Vec<Option<BindlessSlotKind>>> = None;
    let mut fallback_kinds: Option<Vec<Option<BindlessSlotKind>>> = None;

    for entry in &entries {
        let kinds = extract_from_params(&entry.params);
        if !kinds.is_empty() {
            if entry.stage == Stage::Fragment {
                fragment_kinds = Some(kinds);
            } else if fallback_kinds.is_none() || entry.stage == Stage::Compute {
                fallback_kinds = Some(kinds);
            }
        }
    }

    fragment_kinds.or(fallback_kinds).unwrap_or_default()
}

/// Transform a Slang source string that may contain `[goldy_*]` virtual entry
/// points into a standard Slang source with generated `[shader(...)]` wrappers.
///
/// If no `[goldy_*]` attributes are present, the source is returned unchanged.
pub fn transform_virtual_main(source: &str) -> String {
    let entries = find_all_entries(source);
    if entries.is_empty() {
        return source.to_string();
    }

    // Generate wrapper code (in declaration order — first entry first).
    let mut wrapper_block = String::from("// [generated by goldy virtual_main — do not edit]\n");
    for entry in &entries {
        wrapper_block.push_str(&emit_wrapper(entry));
        wrapper_block.push_str("\n\n");
    }

    // Modify the original source (remove [goldy_*] attrs, rename fns).
    // Process in *reverse* position order so earlier byte offsets stay valid.
    let mut modified = source.to_string();
    for entry in entries.iter().rev() {
        apply_entry_transforms(&mut modified, entry);
    }

    // Reset Slang's line counter to 1 so that diagnostics for the user's code
    // are reported at the original source lines, not offset by the generated prefix.
    wrapper_block + "#line 1\n" + &modified
}

/// WebGPU compute lowering for `[goldy_compute]` entry points.
///
/// Unlike the native lowering, resources become fixed `@group(0)` bindings
/// (expressed as HLSL registers for Slang) and are passed directly to the
/// renamed user function. This intentionally supports only the portable,
/// bounded resource model; scalar dispatch parameters still require a future
/// uniform-ring ABI.
pub fn transform_virtual_main_webgpu_compute(source: &str) -> Result<String, String> {
    let entries = find_all_entries(source);
    if entries.is_empty() {
        return Ok(source.to_string());
    }

    fn resource_decl(param: &Param, binding: u32, name: &str) -> Result<String, String> {
        let ty = param.ty.trim();
        let (decl_ty, register) = if ty.starts_with("Scattered<") || ty == "ByteAddress" {
            (ty.to_string(), 'u')
        } else if ty.starts_with("BufRO<") || ty.starts_with("Interpolated<") {
            (ty.to_string(), 't')
        } else if ty.starts_with("DirectSpatial<") {
            (ty.to_string(), 'u')
        } else if ty == "Filter" {
            (ty.to_string(), 's')
        } else if matches!(param.kind, ParamKind::Broadcast) {
            (format!("ConstantBuffer<{ty}>"), 'b')
        } else {
            return Err(format!("unsupported WebGPU resource parameter type `{ty}`"));
        };
        Ok(format!("{decl_ty} {name} : register({register}{binding}, space0);"))
    }

    let mut generated = String::from("// [generated by goldy virtual_main WebGPU lowering — do not edit]\n");
    for entry in &entries {
        if entry.stage != Stage::Compute {
            return Err("WebGPU prototype only lowers [goldy_compute] entry points".to_string());
        }

        let mut binding = 0u32;
        let mut sv_index = 0u32;
        let mut signature = Vec::new();
        let mut body = String::new();
        let mut call_args = Vec::new();

        for item in &entry.params {
            let ParamItem::Single(param) = item else {
                return Err("conditional parameters are not supported by the WebGPU prototype".to_string());
            };
            match &param.kind {
                ParamKind::Resource | ParamKind::Broadcast => {
                    let global = format!("_goldy_wgpu_binding_{binding}");
                    generated.push_str(&resource_decl(param, binding, &global)?);
                    generated.push('\n');
                    if matches!(param.kind, ParamKind::Broadcast) {
                        call_args.push(format!("*{global}"));
                    } else {
                        call_args.push(global);
                    }
                    binding += 1;
                }
                ParamKind::SystemValue(sv) => {
                    let arg = format!("_sv{sv_index}");
                    signature.push(format!("{} {} : {}", sv.primitive(), arg, sv.semantic()));
                    body.push_str(&format!("    {} {} = {}({});\n", param.ty, param.name, param.ty, arg));
                    call_args.push(param.name.clone());
                    sv_index += 1;
                }
                ParamKind::Scalar => {
                    return Err(format!(
                        "WebGPU prototype does not yet lower scalar dispatch parameter `{}`; bind a broadcast buffer",
                        param.name
                    ));
                }
                ParamKind::PassThrough => {
                    return Err(format!(
                        "WebGPU compute entry has unsupported pass-through parameter `{}`",
                        param.name
                    ));
                }
            }
        }

        generated.push_str(entry.stage.shader_attr());
        generated.push('\n');
        if let Some((x, y, z)) = entry.numthreads {
            generated.push_str(&format!("[numthreads({x}, {y}, {z})]\n"));
        }
        generated.push_str(&format!(
            "{} {}({}) {{\n{}",
            entry.return_type,
            entry.fn_name,
            signature.join(", "),
            body
        ));
        let user_fn = format!("_goldy_user_{}", entry.fn_name);
        if entry.return_type == "void" {
            generated.push_str(&format!("    {user_fn}({});\n", call_args.join(", ")));
        } else {
            generated.push_str(&format!("    return {user_fn}({});\n", call_args.join(", ")));
        }
        generated.push_str("}\n\n");
    }

    let mut modified = source.to_string();
    for entry in entries.iter().rev() {
        apply_entry_transforms(&mut modified, entry);
    }
    Ok(generated + "#line 1\n" + &modified)
}

/// CUDA compute lowering for `[goldy_compute]` entry points.
///
/// Unlike the native push-constant / bindless-heap lowering, each buffer resource
/// becomes a `uniform StructuredBuffer` / `RWStructuredBuffer` entry-point
/// parameter. Slang's CUDA target lowers those to `{T* data; size_t count}`
/// descriptors that the backend fills from `CUdeviceptr` at launch time.
///
/// Broadcast structs are passed as a one-element read-only structured buffer and
/// loaded at index 0. Textures, samplers, scalars, and graphics stages are rejected.
pub fn transform_virtual_main_cuda_compute(source: &str) -> Result<String, String> {
    let entries = find_all_entries(source);
    if entries.is_empty() {
        return Ok(source.to_string());
    }

    fn resource_param_ty(param: &Param) -> Result<String, String> {
        let ty = param.ty.trim();
        if ty.starts_with("Scattered<") || ty == "ByteAddress" || ty.starts_with("BufRO<") {
            Ok(ty.to_string())
        } else if matches!(param.kind, ParamKind::Broadcast) {
            // One-element read-only buffer; wrapper loads `[0]`.
            Ok(format!("StructuredBuffer<{ty}>"))
        } else if ty.starts_with("Interpolated<") || ty.starts_with("DirectSpatial<") || ty == "Filter" {
            Err(format!(
                "CUDA prototype does not support texture/sampler parameter type `{ty}`"
            ))
        } else {
            Err(format!("unsupported CUDA resource parameter type `{ty}`"))
        }
    }

    let mut generated = String::from("// [generated by goldy virtual_main CUDA lowering — do not edit]\n");
    for entry in &entries {
        if entry.stage != Stage::Compute {
            return Err("CUDA prototype only lowers [goldy_compute] entry points".to_string());
        }

        let mut signature = Vec::new();
        let mut body = String::new();
        let mut call_args = Vec::new();
        let mut sv_index = 0u32;
        let mut binding = 0u32;

        for item in &entry.params {
            let ParamItem::Single(param) = item else {
                return Err("conditional parameters are not supported by the CUDA prototype".to_string());
            };
            match &param.kind {
                ParamKind::Resource | ParamKind::Broadcast => {
                    let param_ty = resource_param_ty(param)?;
                    let global = format!("_goldy_cuda_binding_{binding}");
                    signature.push(format!("uniform {param_ty} {global}"));
                    if matches!(param.kind, ParamKind::Broadcast) {
                        body.push_str(&format!("    {} {} = {}[0];\n", param.ty, param.name, global));
                        call_args.push(param.name.clone());
                    } else {
                        call_args.push(global);
                    }
                    binding += 1;
                }
                ParamKind::SystemValue(sv) => {
                    let arg = format!("_sv{sv_index}");
                    signature.push(format!("{} {} : {}", sv.primitive(), arg, sv.semantic()));
                    body.push_str(&format!("    {} {} = {}({});\n", param.ty, param.name, param.ty, arg));
                    call_args.push(param.name.clone());
                    sv_index += 1;
                }
                ParamKind::Scalar => {
                    return Err(format!(
                        "CUDA prototype does not yet lower scalar dispatch parameter `{}`; bind a broadcast buffer",
                        param.name
                    ));
                }
                ParamKind::PassThrough => {
                    return Err(format!(
                        "CUDA compute entry has unsupported pass-through parameter `{}`",
                        param.name
                    ));
                }
            }
        }

        generated.push_str(entry.stage.shader_attr());
        generated.push('\n');
        if let Some((x, y, z)) = entry.numthreads {
            generated.push_str(&format!("[numthreads({x}, {y}, {z})]\n"));
        }
        generated.push_str(&format!(
            "{} {}({}) {{\n{}",
            entry.return_type,
            entry.fn_name,
            signature.join(", "),
            body
        ));
        let user_fn = format!("_goldy_user_{}", entry.fn_name);
        if entry.return_type == "void" {
            generated.push_str(&format!("    {user_fn}({});\n", call_args.join(", ")));
        } else {
            generated.push_str(&format!("    return {user_fn}({});\n", call_args.join(", ")));
        }
        generated.push_str("}\n\n");
    }

    let mut modified = source.to_string();
    for entry in entries.iter().rev() {
        apply_entry_transforms(&mut modified, entry);
    }
    Ok(generated + "#line 1\n" + &modified)
}

/// Slang translation-unit source after the optional virtual-main rewrite.
///
/// This matches the compiler gate: [`transform_virtual_main`] runs only when `source`
/// contains `[goldy_compute]`, `[goldy_vertex]`, or `[goldy_fragment]` as substrings
/// (the same [`str::contains`] checks previously used before calling Slang).
///
/// Callers that feed Slang directly (disk cache keys, `add_translation_unit_source_string`) should
/// use this so the hashed text matches the compiled text.
pub fn effective_slang_source_for_compile(source: &str) -> Cow<'_, str> {
    if source.contains("[goldy_compute]") || source.contains("[goldy_vertex]") || source.contains("[goldy_fragment]") {
        Cow::Owned(transform_virtual_main(source))
    } else {
        Cow::Borrowed(source)
    }
}

// ---------------------------------------------------------------------------
// Stage
// ---------------------------------------------------------------------------

/// The shader stage for a virtual entry point.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Stage {
    Compute,
    Vertex,
    Fragment,
}

impl Stage {
    fn shader_attr(self) -> &'static str {
        match self {
            Stage::Compute => r#"[shader("compute")]"#,
            Stage::Vertex => r#"[shader("vertex")]"#,
            Stage::Fragment => r#"[shader("fragment")]"#,
        }
    }
}

// ---------------------------------------------------------------------------
// Parameter model
// ---------------------------------------------------------------------------

/// Classification of a single parameter in a `[goldy_*]` entry point.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ParamKind {
    /// A Goldy resource handle type → becomes `uniform uint` resource slot resolved via goldy_*.
    Resource,
    /// A user-defined struct → becomes `uniform uint` resource slot resolved via goldy_broadcast.
    ///
    /// In compute shaders every unrecognised struct is a broadcast (constant buffer).
    /// In vertex/fragment shaders every unrecognised struct *except the last* is a broadcast;
    /// the last one is the stage input and stays as [`ParamKind::PassThrough`].
    Broadcast,
    /// A Goldy SV wrapper type → becomes an `SV_*`-annotated parameter.
    SystemValue(SvKind),
    /// A plain scalar type (uint, float, int, …) → becomes `uniform T` entry-point parameter.
    Scalar,
    /// Stage input varying / vertex-attribute struct → passed through unchanged.
    PassThrough,
}

/// The kind of system-value a `SystemValue` param maps to.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SvKind {
    DispatchThreadId, // SV_DispatchThreadID  uint3
    GroupThreadId,    // SV_GroupThreadID     uint3
    GroupId,          // SV_GroupID           uint3
    VertexId,         // SV_VertexID          uint
    InstanceId,       // SV_InstanceID        uint
    IsFrontFace,      // SV_IsFrontFace       bool
}

impl SvKind {
    /// HLSL semantic name.
    fn semantic(self) -> &'static str {
        match self {
            SvKind::DispatchThreadId => "SV_DispatchThreadID",
            SvKind::GroupThreadId => "SV_GroupThreadID",
            SvKind::GroupId => "SV_GroupID",
            SvKind::VertexId => "SV_VertexID",
            SvKind::InstanceId => "SV_InstanceID",
            SvKind::IsFrontFace => "SV_IsFrontFace",
        }
    }

    /// The primitive type for the generated SV parameter.
    fn primitive(self) -> &'static str {
        match self {
            SvKind::DispatchThreadId | SvKind::GroupThreadId | SvKind::GroupId => "uint3",
            SvKind::VertexId | SvKind::InstanceId => "uint",
            SvKind::IsFrontFace => "bool",
        }
    }
}

/// A single parameter from a `[goldy_*]` entry-point signature.
#[derive(Debug, Clone)]
pub struct Param {
    /// Parameter name as written in the user source (e.g., `"data"`, `"id"`).
    pub name: String,
    /// Full type string as written (e.g., `"Scattered<uint>"`, `"ThreadId"`).
    pub ty: String,
    /// Classified kind.
    pub kind: ParamKind,
}

/// A parameter item in a `[goldy_*]` entry-point signature.
///
/// Parameters may be wrapped in preprocessor conditionals (`#ifdef`/`#else`/`#endif`)
/// directly inside the function's parameter list (as in `fine.slang`'s MSAA variants).
/// This enum represents both the simple and conditional cases.
#[derive(Debug, Clone)]
pub enum ParamItem {
    Single(Param),
    /// A `#ifdef condition … #else … #endif` block directly inside the parameter list.
    Conditional {
        condition: String,
        then_params: Vec<Param>,
        else_params: Vec<Param>,
    },
}

// ---------------------------------------------------------------------------
// Entry definition
// ---------------------------------------------------------------------------

/// A fully parsed `[goldy_*]` entry point.
#[derive(Debug, Clone)]
pub struct EntryDef {
    pub stage: Stage,
    pub fn_name: String,
    pub return_type: String,
    /// Return semantic (e.g., `"SV_Target"`) — fragment only.
    pub return_semantic: Option<String>,
    pub params: Vec<ParamItem>,
    /// `[numthreads(X, Y, Z)]` if present (compute only).
    pub numthreads: Option<(u32, u32, u32)>,

    // Byte ranges in the *original* source for the apply_transforms step.
    /// The `[goldy_*]` attribute span (including brackets).
    pub goldy_attr_range: Range<usize>,
    /// The `[numthreads(...)]` span, if found adjacent to the entry.
    pub numthreads_attr_range: Option<Range<usize>>,
    /// The span of just the function name (e.g., `"cs_main"`).
    pub fn_name_range: Range<usize>,
    /// The `: SV_Target` span in the user function declaration (fragment only).
    pub return_semantic_range: Option<Range<usize>>,
}

// ---------------------------------------------------------------------------
// Finding entries
// ---------------------------------------------------------------------------

const GOLDY_STAGES: &[(&str, Stage)] = &[
    ("goldy_compute", Stage::Compute),
    ("goldy_vertex", Stage::Vertex),
    ("goldy_fragment", Stage::Fragment),
];

/// Find all `[goldy_*]` entry points in `source`, sorted by position.
fn find_all_entries(source: &str) -> Vec<EntryDef> {
    let mut entries: Vec<EntryDef> = Vec::new();
    let bytes = source.as_bytes();
    let mut search_from = 0;

    loop {
        // Find the earliest [goldy_*] from `search_from`.
        let mut found: Option<(usize, usize, Stage)> = None;
        for &(token, stage) in GOLDY_STAGES {
            let needle = format!("[{}]", token);
            if let Some(pos) = find_substr(source, search_from, &needle) {
                if found.is_none() || pos < found.unwrap().0 {
                    found = Some((pos, pos + needle.len(), stage));
                }
            }
        }

        let (attr_start, attr_end, stage) = match found {
            Some(f) => f,
            None => break,
        };

        // Skip if it's inside a line comment.
        if is_in_line_comment(source, attr_start) {
            search_from = attr_end;
            continue;
        }

        // Parse this entry.
        match parse_entry(source, attr_start, attr_end, stage, bytes) {
            Some(entry) => {
                search_from = entry.fn_name_range.end;
                entries.push(entry);
            }
            None => {
                search_from = attr_end;
            }
        }
    }

    entries
}

/// Parse a single entry, given the position of its `[goldy_*]` attribute.
fn parse_entry(source: &str, attr_start: usize, attr_end: usize, stage: Stage, _bytes: &[u8]) -> Option<EntryDef> {
    // -----------------------------------------------------------------------
    // 1. Scan backward from attr_start to find the beginning of the attribute
    //    group (so we can find [numthreads] that appears BEFORE [goldy_*]).
    // -----------------------------------------------------------------------
    let group_start = find_attr_group_start(source, attr_start);

    // -----------------------------------------------------------------------
    // 2. Scan forward from attr_end to collect remaining attributes and find
    //    the function declaration.
    // -----------------------------------------------------------------------
    let mut pos = attr_end;
    let mut numthreads_attr_range: Option<Range<usize>> = None;
    let mut numthreads: Option<(u32, u32, u32)> = None;

    // Scan forward through any attributes that follow [goldy_*].
    loop {
        pos = skip_whitespace_and_comments(source, pos);
        if source[pos..].starts_with('[') {
            let (attr_content, bracket_end) = scan_bracket_block(source, pos)?;
            let trimmed = attr_content.trim();

            // Is it [numthreads(...)]?
            if trimmed.starts_with("numthreads") {
                if let Some(nt) = parse_numthreads(trimmed) {
                    numthreads = Some(nt);
                    numthreads_attr_range = Some(pos..bracket_end);
                }
            }
            // Skip [shader("...")] — it will be replaced by our generated one.
            pos = bracket_end;
        } else {
            break;
        }
    }

    // Also check the backward group for [numthreads] if not found forward.
    if numthreads.is_none() {
        if let Some((nt, nt_range)) = find_numthreads_in_range(source, group_start, attr_start) {
            numthreads = Some(nt);
            numthreads_attr_range = Some(nt_range);
        }
    }

    // -----------------------------------------------------------------------
    // 3. Parse the function signature starting at `pos`.
    // -----------------------------------------------------------------------
    pos = skip_whitespace_and_comments(source, pos);

    // Return type: everything up to the function name.
    let (return_type, pos) = scan_identifier(source, pos)?;
    let pos = skip_whitespace_and_comments(source, pos);

    // Function name.
    let fn_name_start = pos;
    let (fn_name, pos) = scan_identifier(source, pos)?;
    let fn_name_end = pos;

    let pos = skip_whitespace_and_comments(source, pos);
    if !source[pos..].starts_with('(') {
        return None; // Not a function declaration
    }

    // Parameter list (balanced parens, handles nested angle brackets).
    let params_open = pos;
    let params_close = find_matching_close(source, params_open, '(', ')')?;
    let params_str = &source[params_open + 1..params_close];
    let mut params = parse_params_items(params_str);
    reclassify_passthrough(&mut params, stage);

    let pos = params_close + 1;
    let pos = skip_whitespace_and_comments(source, pos);

    // Optional return semantic: `: SV_Target` etc.
    let (return_semantic, return_semantic_range) = if source[pos..].starts_with(':') {
        let colon_pos = pos;
        let pos = skip_whitespace_and_comments(source, colon_pos + 1);
        let (semantic, end_pos) = scan_identifier(source, pos).unwrap_or_else(|| (String::new(), pos));
        if !semantic.is_empty() {
            let sem_range = colon_pos..end_pos;
            (Some(semantic), Some(sem_range))
        } else {
            (None, None)
        }
    } else {
        (None, None)
    };

    Some(EntryDef {
        stage,
        fn_name,
        return_type,
        return_semantic,
        params,
        numthreads,
        goldy_attr_range: attr_start..attr_end,
        numthreads_attr_range,
        fn_name_range: fn_name_start..fn_name_end,
        return_semantic_range,
    })
}

// ---------------------------------------------------------------------------
// Parameter parsing
// ---------------------------------------------------------------------------

/// Strip `//`-to-end-of-line comments from each line, preserving newlines so that
/// multi-line parameter lists stay structurally intact for `#ifdef` handling.
fn strip_line_comments(s: &str) -> String {
    s.lines()
        .map(|line| {
            if let Some(idx) = line.find("//") {
                &line[..idx]
            } else {
                line
            }
        })
        .collect::<Vec<_>>()
        .join("\n")
}

/// Parse a comma-separated parameter list, handling nested `<...>` brackets.
///
/// Inline `//` comments are stripped first — without this, trailing comments on
/// parameter lines become part of the next parameter's type string, producing
/// invalid generated code that causes Slang ICEs.
fn parse_params(params_str: &str) -> Vec<Param> {
    let cleaned = strip_line_comments(params_str);
    let mut params = Vec::new();
    let mut current = String::new();
    let mut depth: i32 = 0;

    for ch in cleaned.chars() {
        match ch {
            '<' | '(' | '[' => {
                depth += 1;
                current.push(ch);
            }
            '>' | ')' | ']' => {
                depth -= 1;
                current.push(ch);
            }
            ',' if depth == 0 => {
                let trimmed = current.trim();
                if !trimmed.is_empty() {
                    if let Some(p) = parse_single_param(trimmed) {
                        params.push(p);
                    }
                }
                current.clear();
            }
            _ => current.push(ch),
        }
    }

    let trimmed = current.trim();
    if !trimmed.is_empty() {
        if let Some(p) = parse_single_param(trimmed) {
            params.push(p);
        }
    }

    params
}

/// Parse a parameter list that may contain `#ifdef`/`#else`/`#endif` directives.
///
/// For simple lists with no preprocessor conditionals, delegates to [`parse_params`].
/// For lists that contain `#ifdef`, segments the text at preprocessor boundaries and
/// returns [`ParamItem::Conditional`] items wrapping the alternative parameter sets.
fn parse_params_items(params_str: &str) -> Vec<ParamItem> {
    if !params_str.contains("#ifdef") {
        return parse_params(params_str).into_iter().map(ParamItem::Single).collect();
    }

    let mut items: Vec<ParamItem> = Vec::new();
    let mut outer_text = String::new();

    let mut in_ifdef = false;
    let mut condition = String::new();
    let mut then_text = String::new();
    let mut else_text = String::new();
    let mut in_else = false;

    for line in params_str.lines() {
        let trimmed = line.trim();

        // Skip blank lines and `//` comment lines: they may contain commas or other
        // punctuation that would confuse the comma-splitting in parse_params.
        if trimmed.is_empty() || trimmed.starts_with("//") {
            continue;
        }

        if let Some(rest) = trimmed.strip_prefix("#ifdef ") {
            if !in_ifdef {
                // Flush accumulated outer text as regular params.
                for p in parse_params(&outer_text) {
                    items.push(ParamItem::Single(p));
                }
                outer_text.clear();
                condition = rest.trim().to_string();
                then_text.clear();
                else_text.clear();
                in_else = false;
                in_ifdef = true;
            } else {
                // Nested #ifdef: not supported; treat line as literal text.
                if in_else {
                    else_text.push_str(line);
                    else_text.push('\n');
                } else {
                    then_text.push_str(line);
                    then_text.push('\n');
                }
            }
        } else if trimmed == "#else" && in_ifdef {
            in_else = true;
        } else if trimmed == "#endif" && in_ifdef {
            let then_params = parse_params(&then_text);
            let else_params = parse_params(&else_text);
            items.push(ParamItem::Conditional {
                condition: std::mem::take(&mut condition),
                then_params,
                else_params,
            });
            in_ifdef = false;
            in_else = false;
            then_text.clear();
            else_text.clear();
        } else if in_ifdef {
            if in_else {
                else_text.push_str(line);
                else_text.push('\n');
            } else {
                then_text.push_str(line);
                then_text.push('\n');
            }
        } else {
            outer_text.push_str(line);
            outer_text.push('\n');
        }
    }

    // Flush any remaining outer text.
    for p in parse_params(&outer_text) {
        items.push(ParamItem::Single(p));
    }

    items
}

/// Parse a single `type name` parameter.
fn parse_single_param(s: &str) -> Option<Param> {
    let s = s.trim();
    if s.is_empty() {
        return None;
    }

    // Skip `uniform` keyword if present (old-style params in shaders that mix styles).
    let s = if let Some(rest) = s.strip_prefix("uniform") {
        rest.trim_start()
    } else {
        s
    };

    // Find the split point between type and name: the last "whitespace gap"
    // that is not inside angle brackets or parentheses.
    let mut last_space: Option<usize> = None;
    let mut depth: i32 = 0;
    let chars: Vec<(usize, char)> = s.char_indices().collect();
    for (byte_pos, ch) in &chars {
        match ch {
            '<' | '(' | '[' => depth += 1,
            '>' | ')' | ']' => depth -= 1,
            ' ' | '\t' | '\n' | '\r' if depth == 0 => {
                last_space = Some(*byte_pos);
            }
            _ => {}
        }
    }

    let (ty_str, name) = if let Some(sp) = last_space {
        let ty = s[..sp].trim();
        let name = s[sp + 1..].trim();
        // If name contains spaces it likely has a : SV_ annotation (old-style param).
        // Strip the annotation: `id : SV_DispatchThreadID` → name = "id".
        let name = name.split(':').next().unwrap_or(name).trim();
        (ty.to_string(), name.to_string())
    } else {
        // No space — probably just a type with no name (shouldn't happen normally).
        (s.to_string(), String::new())
    };

    let kind = classify_type(&ty_str);
    Some(Param { name, ty: ty_str, kind })
}

/// Build the RHS initialiser expression for a resource parameter.
///
/// Resource type aliases (`Scattered<T>`, `BufRO<T>`, etc.) are plain typealiases
/// for the underlying buffer/texture handle types. The generated wrapper calls the
/// corresponding `goldy_*` free function to resolve the resource slot index to
/// a live handle. Broadcast (constant-buffer) params are handled by their own
/// `ParamKind::Broadcast` arm — only handle types appear here.
fn resource_init_expr(ty: &str, slot_var: &str) -> String {
    let ty = ty.trim();

    // Extract inner type T from Wrapper<T> for generic free-function calls.
    fn inner(ty: &str, prefix: &str) -> Option<String> {
        if ty.starts_with(prefix) && ty.ends_with('>') {
            Some(ty[prefix.len()..ty.len() - 1].to_string())
        } else {
            None
        }
    }

    if let Some(t) = inner(ty, "Scattered<") {
        return format!("goldy_scattered<{}>({})", t, slot_var);
    }
    if let Some(t) = inner(ty, "BufRO<") {
        return format!("goldy_buf_ro<{}>({})", t, slot_var);
    }
    if let Some(t) = inner(ty, "Interpolated<") {
        return format!("goldy_interpolated<{}>({})", t, slot_var);
    }
    if let Some(t) = inner(ty, "DirectSpatial<") {
        return format!("goldy_direct_spatial<{}>({})", t, slot_var);
    }
    if ty == "ByteAddress" {
        return format!("goldy_byte_address({})", slot_var);
    }
    if ty == "Filter" {
        return format!("goldy_filter({})", slot_var);
    }
    // Unrecognised resource type: fall back to calling __init(slot).
    format!("{}({})", ty, slot_var)
}

/// Post-pass: reclassify `PassThrough` parameters as `Broadcast` based on stage.
///
/// After initial classification, every unrecognised struct lands in `PassThrough`.
/// This pass promotes those to `Broadcast` (constant-buffer fetch) except for the
/// stage-input parameter in vertex/fragment shaders, which genuinely passes through.
///
/// Rule:
/// - Compute: ALL `PassThrough` → `Broadcast` (no stage inputs exist in compute).
/// - Vertex/Fragment: all `PassThrough` except the last `Single(PassThrough)` → `Broadcast`.
///   The last one is the vertex attribute struct or fragment varying.
fn reclassify_passthrough(params: &mut [ParamItem], stage: Stage) {
    // In vertex/fragment, the last Single(PassThrough) is the stage input — keep it,
    // but ONLY if it's truly the final parameter.  When system-value params (VertexId,
    // InstanceId, etc.) follow the last PassThrough, there is no stage-input struct
    // and all PassThroughs are broadcasts.
    let preserve_idx: Option<usize> = match stage {
        Stage::Vertex | Stage::Fragment => {
            let last_pt = params
                .iter()
                .rposition(|item| matches!(item, ParamItem::Single(p) if p.kind == ParamKind::PassThrough));
            if let Some(idx) = last_pt {
                let has_sv_after = params[idx + 1..]
                    .iter()
                    .any(|item| matches!(item, ParamItem::Single(p) if matches!(p.kind, ParamKind::SystemValue(_))));
                if has_sv_after {
                    None
                } else {
                    Some(idx)
                }
            } else {
                None
            }
        }
        Stage::Compute => None,
    };

    for (i, item) in params.iter_mut().enumerate() {
        match item {
            ParamItem::Single(p) if p.kind == ParamKind::PassThrough && Some(i) != preserve_idx => {
                p.kind = ParamKind::Broadcast;
            }
            ParamItem::Conditional {
                then_params,
                else_params,
                ..
            } => {
                // Conditional blocks never contain stage inputs — all PassThrough → Broadcast.
                for p in then_params.iter_mut().chain(else_params.iter_mut()) {
                    if p.kind == ParamKind::PassThrough {
                        p.kind = ParamKind::Broadcast;
                    }
                }
            }
            _ => {}
        }
    }
}

/// Classify a Slang type string into a `ParamKind`.
fn classify_type(ty: &str) -> ParamKind {
    let ty = ty.trim();

    // Resource wrappers (generic).
    if ty.starts_with("Scattered<") {
        return ParamKind::Resource;
    }
    if ty.starts_with("BufRO<") {
        return ParamKind::Resource;
    }
    if ty.starts_with("Interpolated<") {
        return ParamKind::Resource;
    }
    if ty.starts_with("DirectSpatial<") {
        return ParamKind::Resource;
    }

    // Resource wrappers (non-generic).
    if ty == "ByteAddress" {
        return ParamKind::Resource;
    }
    if ty == "Filter" {
        return ParamKind::Resource;
    }

    // System-value types.
    if ty == "ThreadId" {
        return ParamKind::SystemValue(SvKind::DispatchThreadId);
    }
    if ty == "GroupThreadId" {
        return ParamKind::SystemValue(SvKind::GroupThreadId);
    }
    if ty == "GroupId" {
        return ParamKind::SystemValue(SvKind::GroupId);
    }
    if ty == "VertexId" {
        return ParamKind::SystemValue(SvKind::VertexId);
    }
    if ty == "InstanceId" {
        return ParamKind::SystemValue(SvKind::InstanceId);
    }
    if ty == "IsFrontFace" {
        return ParamKind::SystemValue(SvKind::IsFrontFace);
    }

    // Plain scalars → uniform entry-point parameters.
    if matches!(
        ty,
        "uint"
            | "int"
            | "float"
            | "bool"
            | "half"
            | "double"
            | "uint2"
            | "uint3"
            | "uint4"
            | "int2"
            | "int3"
            | "int4"
            | "float2"
            | "float3"
            | "float4"
    ) {
        return ParamKind::Scalar;
    }

    // Everything else is a pass-through (varying input structs, etc.).
    ParamKind::PassThrough
}

// ---------------------------------------------------------------------------
// Code generation
// ---------------------------------------------------------------------------

// ---------------------------------------------------------------------------
// Wrapper builder (accumulates sig / body / call for emit_wrapper)
// ---------------------------------------------------------------------------

/// Mutable index counters shared while emitting a wrapper (`emit_wrapper`).
struct WrapperIndices<'a> {
    bindless_idx: &'a mut u32,
    user_idx: &'a mut u32,
    sv_idx: &'a mut u32,
    pt_idx: &'a mut u32,
}

/// Accumulates the three parallel output streams needed to build a generated
/// `[shader(...)]` wrapper: the function signature parameter list, the body
/// (init statements), and the argument list for the call to the user function.
///
/// Handles both plain parameters and `#ifdef`/`#else`/`#endif` conditional
/// blocks that may appear inside the parameter list.
struct WrapperBuilder {
    sig: String,
    body: String,
    call: String,
    sig_sep: bool,
    call_sep: bool,
}

impl WrapperBuilder {
    fn new() -> Self {
        Self {
            sig: String::new(),
            body: String::new(),
            call: String::new(),
            sig_sep: false,
            call_sep: false,
        }
    }

    fn push_sig(&mut self, s: &str) {
        if self.sig_sep {
            self.sig.push_str(", ");
        }
        self.sig.push_str(s);
        self.sig_sep = true;
    }

    fn push_call(&mut self, s: &str) {
        if self.call_sep {
            self.call.push_str(", ");
        }
        self.call.push_str(s);
        self.call_sep = true;
    }

    fn push_body_stmt(&mut self, s: &str) {
        if !s.is_empty() {
            self.body.push_str(s);
            self.body.push('\n');
        }
    }

    fn process_param(
        &mut self,
        param: &Param,
        bindless_idx: &mut u32,
        user_idx: &mut u32,
        sv_idx: &mut u32,
        pt_idx: &mut u32,
    ) {
        match &param.kind {
            ParamKind::Resource => {
                let k = *bindless_idx;
                *bindless_idx += 1;
                let extract = format!("goldy_frame_table_index(_rs0, {}u, _rs1, _rs2)", k);
                let init_expr = resource_init_expr(&param.ty, &extract);
                self.push_body_stmt(&format!("    {} {} = {};", param.ty, param.name, init_expr));
                self.push_call(&param.name);
            }
            ParamKind::SystemValue(sv) => {
                let gn = format!("_sv{}", *sv_idx);
                *sv_idx += 1;
                self.push_sig(&format!("{} {} : {}", sv.primitive(), gn, sv.semantic()));
                self.push_body_stmt(&format!("    {} {} = {}({});", param.ty, param.name, param.ty, gn));
                self.push_call(&param.name);
            }
            ParamKind::Broadcast => {
                let k = *bindless_idx;
                *bindless_idx += 1;
                let extract = format!("goldy_frame_table_index(_rs0, {}u, _rs1, _rs2)", k);
                self.push_body_stmt(&format!(
                    "    {} {} = goldy_broadcast<{}>({});",
                    param.ty, param.name, param.ty, extract
                ));
                self.push_call(&param.name);
            }
            ParamKind::Scalar => {
                let j = *user_idx;
                *user_idx += 1;
                // User params are full u32 words at _uw{j}.
                // Use the appropriate HLSL bit-reinterpret intrinsic for the type.
                let init_expr = match param.ty.as_str() {
                    "float" => format!("asfloat(_uw{})", j),
                    "int" => format!("asint(_uw{})", j),
                    "uint" => format!("_uw{}", j),
                    "bool" => format!("_uw{} != 0u", j),
                    _ => format!("({})_uw{}", param.ty, j),
                };
                self.push_body_stmt(&format!("    {} {} = {};", param.ty, param.name, init_expr));
                self.push_call(&param.name);
            }
            ParamKind::PassThrough => {
                let gn = format!("_pt{}", *pt_idx);
                *pt_idx += 1;
                self.push_sig(&format!("{} {}", param.ty, gn));
                self.push_call(&gn);
            }
        }
    }

    /// Emit a `#ifdef cond … #else … #endif` conditional block.
    ///
    /// Both branches start with the same push/sv/pt indices (they are alternatives).
    /// After the block the shared indices advance to the maximum reached by either branch.
    ///
    /// The generated sig and call blocks include a trailing comma inside each branch
    /// so that the preprocessed result is always a valid comma-separated list, regardless
    /// of which branch the compiler chooses.  The caller's `sig_sep`/`call_sep` flags
    /// are reset to `false` after the block (the trailing comma takes their place).
    fn process_conditional(
        &mut self,
        cond: &str,
        then_params: &[Param],
        else_params: &[Param],
        idx: &mut WrapperIndices<'_>,
    ) {
        let start_bindless = *idx.bindless_idx;
        let start_user = *idx.user_idx;
        let start_sv = *idx.sv_idx;
        let start_pt = *idx.pt_idx;

        // Process then-branch with a fresh builder starting at the same indices.
        let mut t_bindless = start_bindless;
        let mut t_user = start_user;
        let mut t_sv = start_sv;
        let mut t_pt = start_pt;
        let mut then_b = WrapperBuilder::new();
        for p in then_params {
            then_b.process_param(p, &mut t_bindless, &mut t_user, &mut t_sv, &mut t_pt);
        }

        // Process else-branch starting at the same indices.
        let mut e_bindless = start_bindless;
        let mut e_user = start_user;
        let mut e_sv = start_sv;
        let mut e_pt = start_pt;
        let mut else_b = WrapperBuilder::new();
        for p in else_params {
            else_b.process_param(p, &mut e_bindless, &mut e_user, &mut e_sv, &mut e_pt);
        }

        // Advance shared indices to the maximum reached by either branch.
        *idx.bindless_idx = t_bindless.max(e_bindless);
        *idx.user_idx = t_user.max(e_user);
        *idx.sv_idx = t_sv.max(e_sv);
        *idx.pt_idx = t_pt.max(e_pt);

        // --- Conditional signature block ---
        // Add a leading comma before `#ifdef` if a previous param was emitted.
        if self.sig_sep {
            self.sig.push(',');
        }
        let then_sig_trail = if then_b.sig.is_empty() { "" } else { "," };
        let else_sig_trail = if else_b.sig.is_empty() { "" } else { "," };
        if else_b.sig.is_empty() {
            self.sig.push_str(&format!(
                "\n#ifdef {}\n    {}{}\n#endif\n",
                cond, then_b.sig, then_sig_trail
            ));
        } else {
            self.sig.push_str(&format!(
                "\n#ifdef {}\n    {}{}\n#else\n    {}{}\n#endif\n",
                cond, then_b.sig, then_sig_trail, else_b.sig, else_sig_trail
            ));
        }
        // The trailing comma inside the block replaces the separator for the next param.
        self.sig_sep = false;

        // --- Conditional body block ---
        if !then_b.body.is_empty() || !else_b.body.is_empty() {
            self.body.push_str(&format!("#ifdef {}\n", cond));
            self.body.push_str(&then_b.body);
            if !else_b.body.is_empty() {
                self.body.push_str("#else\n");
                self.body.push_str(&else_b.body);
            }
            self.body.push_str("#endif\n");
        }

        // --- Conditional call-args block ---
        if self.call_sep {
            self.call.push(',');
        }
        let then_call_trail = if then_b.call.is_empty() { "" } else { "," };
        let else_call_trail = if else_b.call.is_empty() { "" } else { "," };
        if else_b.call.is_empty() {
            self.call.push_str(&format!(
                "\n#ifdef {}\n    {}{}\n#endif\n",
                cond, then_b.call, then_call_trail
            ));
        } else {
            self.call.push_str(&format!(
                "\n#ifdef {}\n    {}{}\n#else\n    {}{}\n#endif\n",
                cond, then_b.call, then_call_trail, else_b.call, else_call_trail
            ));
        }
        self.call_sep = false;
    }
}

/// Emit the generated `[shader("...")]` wrapper for one entry point.
fn emit_wrapper(entry: &EntryDef) -> String {
    let mut out = String::new();

    // Stage attribute.
    out.push_str(entry.stage.shader_attr());
    out.push('\n');

    // Transfer [numthreads] for compute.
    if let Some((x, y, z)) = entry.numthreads {
        out.push_str(&format!("[numthreads({}, {}, {})]\n", x, y, z));
    }

    let user_fn_name = format!("_goldy_user_{}", entry.fn_name);

    let mut bindless_idx = 0u32;
    let mut user_idx = 0u32;
    let mut sv_idx = 0u32;
    let mut pt_idx = 0u32;
    let mut wb = WrapperBuilder::new();

    // Always emit all 8 bindless words (_bw0.._bw7), 8 user words (_uw0.._uw7),
    // frame-table dispatch base (_rs0 = PushLayout._reserved[0]), and the
    // per-context frame-table slots (_rs1/_rs2 = _reserved[1]/_reserved[2]) so
    // that push constant offsets are fixed regardless of how many params are used.
    for i in 0..8u32 {
        wb.push_sig(&format!("uniform uint _bw{}", i));
    }
    for i in 0..8u32 {
        wb.push_sig(&format!("uniform uint _uw{}", i));
    }
    wb.push_sig("uniform uint _rs0");
    wb.push_sig("uniform uint _rs1");
    wb.push_sig("uniform uint _rs2");

    for item in &entry.params {
        match item {
            ParamItem::Single(param) => {
                wb.process_param(param, &mut bindless_idx, &mut user_idx, &mut sv_idx, &mut pt_idx);
            }
            ParamItem::Conditional {
                condition,
                then_params,
                else_params,
            } => {
                let mut idx = WrapperIndices {
                    bindless_idx: &mut bindless_idx,
                    user_idx: &mut user_idx,
                    sv_idx: &mut sv_idx,
                    pt_idx: &mut pt_idx,
                };
                wb.process_conditional(condition, then_params, else_params, &mut idx);
            }
        }
    }

    // --- Function signature ---
    let ret_sem = if let Some(ref sem) = entry.return_semantic {
        format!(" : {}", sem)
    } else {
        String::new()
    };

    out.push_str(&format!(
        "{} {}({}){}",
        entry.return_type, entry.fn_name, wb.sig, ret_sem
    ));
    out.push_str(" {\n");

    // --- Body init statements ---
    out.push_str(&wb.body);

    // --- Call to user function ---
    let is_void = entry.return_type == "void";
    if is_void {
        out.push_str(&format!("    {}({});\n", user_fn_name, wb.call));
    } else {
        out.push_str(&format!("    return {}({});\n", user_fn_name, wb.call));
    }

    out.push('}');
    out
}

// ---------------------------------------------------------------------------
// Source modification
// ---------------------------------------------------------------------------

/// Modify the original source in place (applied in reverse position order).
fn apply_entry_transforms(source: &mut String, entry: &EntryDef) {
    // We must apply changes in reverse order of position to keep ranges valid.
    // Collect all (start, end, replacement) tuples and sort descending.
    let mut replacements: Vec<(usize, usize, &str)> = Vec::new();

    // 1. Remove return semantic from user fn (fragment only; goes to generated wrapper).
    if let Some(ref sr) = entry.return_semantic_range {
        replacements.push((sr.start, sr.end, ""));
    }

    // 2. Rename the function name.
    let new_fn_name = format!("_goldy_user_{}", entry.fn_name);
    replacements.push((
        entry.fn_name_range.start,
        entry.fn_name_range.end,
        // leak is fine — these strings live for the duration of compilation
        Box::leak(new_fn_name.into_boxed_str()),
    ));

    // 3. Remove [numthreads(...)] from the user function (compute only).
    //    It lives on the generated wrapper, not the user helper function.
    if let Some(ref nr) = entry.numthreads_attr_range {
        // Include any trailing whitespace/newline.
        let end = skip_whitespace_in_source(source, nr.end);
        replacements.push((nr.start, end, ""));
    }

    // 4. Remove the [goldy_*] attribute.
    {
        let end = skip_whitespace_in_source(source, entry.goldy_attr_range.end);
        replacements.push((entry.goldy_attr_range.start, end, ""));
    }

    // Sort by start position descending, apply.
    replacements.sort_by_key(|b| std::cmp::Reverse(b.0));
    for (start, end, replacement) in replacements {
        if start < source.len() && end <= source.len() && start <= end {
            source.replace_range(start..end, replacement);
        }
    }
}

/// Advance `pos` past any whitespace in `source`, stopping before non-whitespace.
fn skip_whitespace_in_source(source: &str, pos: usize) -> usize {
    let bytes = source.as_bytes();
    let mut p = pos;
    while p < bytes.len() && (bytes[p] == b' ' || bytes[p] == b'\t' || bytes[p] == b'\n' || bytes[p] == b'\r') {
        p += 1;
    }
    p
}

// ---------------------------------------------------------------------------
// Scanning helpers
// ---------------------------------------------------------------------------

/// Find `needle` in `source` starting at `from`. Returns the byte offset of
/// the first character of the match.
fn find_substr(source: &str, from: usize, needle: &str) -> Option<usize> {
    source[from..].find(needle).map(|i| from + i)
}

/// Return true if position `pos` in `source` is inside a `//` line comment.
fn is_in_line_comment(source: &str, pos: usize) -> bool {
    // Find the start of the line.
    let line_start = source[..pos].rfind('\n').map(|i| i + 1).unwrap_or(0);
    let line = &source[line_start..pos];
    // Check if `//` appears before pos on the same line, outside strings.
    line.contains("//")
}

/// Skip whitespace and `//` line comments in `source` starting at `pos`.
fn skip_whitespace_and_comments(source: &str, mut pos: usize) -> usize {
    let bytes = source.as_bytes();
    loop {
        // Skip whitespace.
        while pos < bytes.len()
            && (bytes[pos] == b' ' || bytes[pos] == b'\t' || bytes[pos] == b'\n' || bytes[pos] == b'\r')
        {
            pos += 1;
        }
        // Skip // line comments.
        if pos + 1 < bytes.len() && bytes[pos] == b'/' && bytes[pos + 1] == b'/' {
            while pos < bytes.len() && bytes[pos] != b'\n' {
                pos += 1;
            }
            continue;
        }
        // Skip /* block comments */.
        if pos + 1 < bytes.len() && bytes[pos] == b'/' && bytes[pos + 1] == b'*' {
            pos += 2;
            while pos + 1 < bytes.len() {
                if bytes[pos] == b'*' && bytes[pos + 1] == b'/' {
                    pos += 2;
                    break;
                }
                pos += 1;
            }
            continue;
        }
        break;
    }
    pos
}

/// Scan an identifier (letters, digits, underscores) at `pos` in `source`.
/// Returns `(identifier, new_pos)`.
fn scan_identifier(source: &str, pos: usize) -> Option<(String, usize)> {
    let bytes = source.as_bytes();
    let start = pos;
    let mut end = pos;
    while end < bytes.len() && (bytes[end].is_ascii_alphanumeric() || bytes[end] == b'_') {
        end += 1;
    }
    if end == start {
        None
    } else {
        Some((source[start..end].to_string(), end))
    }
}

/// Scan the contents of a `[...]` block.  `pos` must point at `[`.
/// Returns `(inner_content, pos_after_closing_bracket)`.
fn scan_bracket_block(source: &str, pos: usize) -> Option<(String, usize)> {
    let close = find_matching_close(source, pos, '[', ']')?;
    let inner = source[pos + 1..close].to_string();
    Some((inner, close + 1))
}

/// Find the matching closing bracket for the open bracket at `open_pos`.
/// Handles nesting of the same pair; ignores other bracket types.
fn find_matching_close(source: &str, open_pos: usize, open: char, close: char) -> Option<usize> {
    let bytes = source.as_bytes();
    if bytes.get(open_pos) != Some(&(open as u8)) {
        return None;
    }
    let mut depth = 0i32;
    for (byte_idx, ch) in source[open_pos..].char_indices() {
        let abs = open_pos + byte_idx;
        if ch == open {
            depth += 1;
        } else if ch == close {
            depth -= 1;
            if depth == 0 {
                return Some(abs);
            }
        }
    }
    None
}

/// Find the start of the attribute group that contains the `[goldy_*]` at
/// `attr_start`. An attribute group is a sequence of `[...]` blocks separated
/// only by whitespace/comments with no other code in between.
fn find_attr_group_start(source: &str, attr_start: usize) -> usize {
    // Walk backward from attr_start.
    let mut pos = attr_start;
    loop {
        // Skip whitespace backward.
        let prev = skip_backward_whitespace(source, pos);
        if prev == pos {
            break;
        }
        // If the character before is `]`, it's another attribute.
        if source.as_bytes().get(prev.saturating_sub(1)) == Some(&b']') {
            // Find the matching `[`.
            let close = prev - 1;
            if let Some(open) = find_matching_open(source, close, '[', ']') {
                pos = open;
                continue;
            }
        }
        break;
    }
    pos
}

/// Walk backward from `pos` past whitespace. Returns new position.
fn skip_backward_whitespace(source: &str, pos: usize) -> usize {
    let bytes = source.as_bytes();
    let mut p = pos;
    while p > 0 {
        let c = bytes[p - 1];
        if c == b' ' || c == b'\t' || c == b'\n' || c == b'\r' {
            p -= 1;
        } else {
            break;
        }
    }
    p
}

/// Find the matching open bracket for the close bracket at `close_pos`.
fn find_matching_open(source: &str, close_pos: usize, open: char, close: char) -> Option<usize> {
    let bytes = source.as_bytes();
    if bytes.get(close_pos) != Some(&(close as u8)) {
        return None;
    }
    let mut depth = 0i32;
    // Iterate backward.
    let prefix = &source[..=close_pos];
    for (byte_idx, ch) in prefix.char_indices().rev() {
        if ch == close {
            depth += 1;
        } else if ch == open {
            depth -= 1;
            if depth == 0 {
                return Some(byte_idx);
            }
        }
    }
    None
}

/// Find a `[numthreads(...)]` attribute in `source[start..end]`.
fn find_numthreads_in_range(source: &str, start: usize, end: usize) -> Option<((u32, u32, u32), Range<usize>)> {
    let slice = &source[start..end];
    let needle = "numthreads";
    let rel_pos = slice.find(needle)?;
    // Find the opening `[` before this position.
    let bracket_pos = source[start..start + rel_pos].rfind('[').map(|i| start + i)?;
    let (inner, close_pos) = scan_bracket_block(source, bracket_pos)?;
    let nt = parse_numthreads(inner.trim())?;
    Some((nt, bracket_pos..close_pos))
}

/// Parse `numthreads(X, Y, Z)` from the inner content of `[numthreads(...)]`.
fn parse_numthreads(s: &str) -> Option<(u32, u32, u32)> {
    let [x, y, z] = super::parse_numthreads(s)?;
    Some((x, y, z))
}

// ---------------------------------------------------------------------------
// Unit tests
// ---------------------------------------------------------------------------

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

    // -----------------------------------------------------------------------
    // classify_type tests
    // -----------------------------------------------------------------------

    #[test]
    fn classify_resource_types() {
        assert_eq!(classify_type("Scattered<uint>"), ParamKind::Resource);
        assert_eq!(classify_type("Scattered<MyStruct>"), ParamKind::Resource);
        assert_eq!(classify_type("BufRO<float4>"), ParamKind::Resource);
        assert_eq!(classify_type("Interpolated<float4>"), ParamKind::Resource);
        assert_eq!(classify_type("DirectSpatial<float4>"), ParamKind::Resource);
        assert_eq!(classify_type("ByteAddress"), ParamKind::Resource);
        assert_eq!(classify_type("Filter"), ParamKind::Resource);
    }

    #[test]
    fn classify_broadcast_types() {
        // Unrecognised struct types → PassThrough at classify_type level;
        // reclassify_passthrough upgrades them to Broadcast based on stage context.
        assert_eq!(classify_type("SceneUniforms"), ParamKind::PassThrough);
        assert_eq!(classify_type("TimeUniforms"), ParamKind::PassThrough);
        assert_eq!(classify_type("MyCustomParams"), ParamKind::PassThrough);
    }

    #[test]
    fn classify_sv_types() {
        assert_eq!(
            classify_type("ThreadId"),
            ParamKind::SystemValue(SvKind::DispatchThreadId)
        );
        assert_eq!(
            classify_type("GroupThreadId"),
            ParamKind::SystemValue(SvKind::GroupThreadId)
        );
        assert_eq!(classify_type("GroupId"), ParamKind::SystemValue(SvKind::GroupId));
        assert_eq!(classify_type("VertexId"), ParamKind::SystemValue(SvKind::VertexId));
        assert_eq!(classify_type("InstanceId"), ParamKind::SystemValue(SvKind::InstanceId));
        assert_eq!(
            classify_type("IsFrontFace"),
            ParamKind::SystemValue(SvKind::IsFrontFace)
        );
    }

    #[test]
    fn classify_scalar_types() {
        assert_eq!(classify_type("uint"), ParamKind::Scalar);
        assert_eq!(classify_type("float"), ParamKind::Scalar);
        assert_eq!(classify_type("int"), ParamKind::Scalar);
        assert_eq!(classify_type("bool"), ParamKind::Scalar);
        assert_eq!(classify_type("uint3"), ParamKind::Scalar);
        assert_eq!(classify_type("float4"), ParamKind::Scalar);
    }

    #[test]
    fn classify_passthrough_types() {
        assert_eq!(classify_type("VaryingInput"), ParamKind::PassThrough);
        assert_eq!(classify_type("VSOutput"), ParamKind::PassThrough);
        assert_eq!(classify_type("StaticVertexIn"), ParamKind::PassThrough);
    }

    // -----------------------------------------------------------------------
    // parse_params tests
    // -----------------------------------------------------------------------

    #[test]
    fn parse_single_scattered_param() {
        let params = parse_params("Scattered<uint> data");
        assert_eq!(params.len(), 1);
        assert_eq!(params[0].name, "data");
        assert_eq!(params[0].ty, "Scattered<uint>");
        assert_eq!(params[0].kind, ParamKind::Resource);
    }

    #[test]
    fn parse_thread_id_param() {
        let params = parse_params("ThreadId id");
        assert_eq!(params.len(), 1);
        assert_eq!(params[0].name, "id");
        assert_eq!(params[0].kind, ParamKind::SystemValue(SvKind::DispatchThreadId));
    }

    #[test]
    fn parse_multiple_params() {
        // parse_params alone does not run reclassify_passthrough, so MyUniforms → PassThrough.
        // The full transform pipeline upgrades it to Broadcast.
        let params = parse_params("Scattered<uint> data, MyUniforms cfg, ThreadId id, uint scale");
        assert_eq!(params.len(), 4);
        assert_eq!(params[0].kind, ParamKind::Resource);
        assert_eq!(params[1].kind, ParamKind::PassThrough); // promoted to Broadcast by reclassify
        assert_eq!(params[2].kind, ParamKind::SystemValue(SvKind::DispatchThreadId));
        assert_eq!(params[3].kind, ParamKind::Scalar);
        assert_eq!(params[3].name, "scale");
    }

    #[test]
    fn parse_params_with_nested_generics() {
        // BufRO<float4> — outer generic with inner type containing no commas
        let params = parse_params("BufRO<float4> buf, Scattered<MyStruct> data");
        assert_eq!(params.len(), 2);
        assert_eq!(params[0].kind, ParamKind::Resource);
        assert_eq!(params[1].kind, ParamKind::Resource);
    }

    #[test]
    fn parse_passthrough_param() {
        let params = parse_params("VaryingInput input");
        assert_eq!(params.len(), 1);
        assert_eq!(params[0].kind, ParamKind::PassThrough);
        assert_eq!(params[0].name, "input");
    }

    #[test]
    fn parse_empty_params() {
        let params = parse_params("");
        assert!(params.is_empty());
    }

    #[test]
    fn parse_params_strips_inline_comments() {
        let input = "BufRO<FilterUniform> uniforms_buf,\n\
                      Interpolated<float4> src_sampled, // SRV slot: hardware-sampled reads\n\
                      DirectSpatial<float4> src,        // UAV slot\n\
                      DirectSpatial<float4> dst,\n\
                      Filter linear_clamp,              // linear sampler\n\
                      ThreadId gid";
        let params = parse_params(input);
        assert_eq!(params.len(), 6, "expected 6 params, got: {params:?}");
        assert_eq!(params[0].ty, "BufRO<FilterUniform>");
        assert_eq!(params[0].name, "uniforms_buf");
        assert_eq!(params[1].ty, "Interpolated<float4>");
        assert_eq!(params[1].name, "src_sampled");
        assert_eq!(params[2].ty, "DirectSpatial<float4>");
        assert_eq!(params[2].name, "src");
        assert_eq!(params[4].ty, "Filter");
        assert_eq!(params[4].name, "linear_clamp");
    }

    // -----------------------------------------------------------------------
    // parse_numthreads tests
    // -----------------------------------------------------------------------

    #[test]
    fn numthreads_basic() {
        assert_eq!(parse_numthreads("numthreads(64, 1, 1)"), Some((64, 1, 1)));
        assert_eq!(parse_numthreads("numthreads(8, 8, 1)"), Some((8, 8, 1)));
        assert_eq!(parse_numthreads("numthreads( 32 , 2 , 1 )"), Some((32, 2, 1)));
    }

    // -----------------------------------------------------------------------
    // Full transform tests
    // -----------------------------------------------------------------------

    #[test]
    fn no_goldy_attrs_passthrough() {
        let src = r#"
[shader("compute")]
[numthreads(64, 1, 1)]
void cs_main(uniform uint data_slot, uint3 id : SV_DispatchThreadID) {
    StorageBuffer<uint> data = goldy_scattered<uint>(data_slot);
    data[id.x] = data[id.x] * 2;
}
"#;
        let result = transform_virtual_main(src);
        assert_eq!(result, src);
    }

    #[test]
    fn compute_simple_transform() {
        let src = r#"import goldy_exp;

[goldy_compute]
[numthreads(64, 1, 1)]
void cs_main(Scattered<uint> data, ThreadId id) {
    data[id.x] = data[id.x] * 2;
}
"#;
        let result = transform_virtual_main(src);
        // Wrapper must be present.
        assert!(
            result.contains("[shader(\"compute\")]"),
            "Missing [shader(\"compute\")]"
        );
        assert!(result.contains("[numthreads(64, 1, 1)]"), "Missing [numthreads]");
        // Fixed bindless/user words in signature.
        assert!(result.contains("uniform uint _bw0"), "Missing _bw0");
        assert!(result.contains("uniform uint _rs0"), "Missing _rs0");
        assert!(result.contains("uniform uint _uw0"), "Missing _uw0");
        assert!(result.contains("SV_DispatchThreadID"), "Missing SV_DispatchThreadID");
        assert!(
            result.contains("goldy_scattered<uint>(goldy_frame_table_index(_rs0, 0u, _rs1, _rs2))"),
            "Missing scattered init"
        );
        assert!(result.contains("ThreadId id = ThreadId(_sv0)"), "Missing SV init");
        // User function must be renamed.
        assert!(result.contains("_goldy_user_cs_main"), "Missing renamed user fn");
        // [goldy_compute] must be removed from user function.
        let goldy_count = result.matches("[goldy_compute]").count();
        assert_eq!(goldy_count, 0, "[goldy_compute] should be removed");
    }

    #[test]
    fn compute_scalar_resource_slot() {
        let src = r#"import goldy_exp;

[goldy_compute]
[numthreads(64, 1, 1)]
void cs_main(Scattered<uint> data, ThreadId id, uint base) {
    data[id.x + base] = 0;
}
"#;
        let result = transform_virtual_main(src);
        // Bindless region: slot 0 via frame table.
        assert!(
            result.contains("goldy_scattered<uint>(goldy_frame_table_index(_rs0, 0u, _rs1, _rs2))"),
            "Missing scattered from frame table slot 0"
        );
        // User region: _uw0 for uint base.
        assert!(result.contains("uint base = _uw0"), "Missing scalar user param");
        assert!(
            result.contains("_goldy_user_cs_main(data, id, base)"),
            "Wrong call args"
        );
    }

    #[test]
    fn vertex_transform() {
        let src = r#"import goldy_exp;

[goldy_vertex]
VSOutput vs_main(Scattered<QuadInstance> instances, VertexId vid, InstanceId iid) {
    QuadInstance inst = instances[iid.value];
    VSOutput out;
    return out;
}
"#;
        let result = transform_virtual_main(src);
        assert!(result.contains("[shader(\"vertex\")]"), "Missing vertex attr");
        assert!(result.contains("SV_VertexID"), "Missing SV_VertexID");
        assert!(result.contains("SV_InstanceID"), "Missing SV_InstanceID");
        assert!(result.contains("_goldy_user_vs_main"), "Missing renamed fn");
    }

    #[test]
    fn fragment_transform_removes_return_semantic() {
        let src = r#"import goldy_exp;

[goldy_fragment]
float4 fs_main(Broadcast<MyUniforms> uniforms, VaryingInput input) : SV_Target {
    return float4(1, 0, 0, 1);
}
"#;
        let result = transform_virtual_main(src);
        assert!(result.contains("[shader(\"fragment\")]"), "Missing fragment attr");
        assert!(result.contains(": SV_Target"), "Wrapper must keep SV_Target");
        // The user function body should no longer have : SV_Target.
        let user_fn_idx = result.find("_goldy_user_fs_main(").unwrap();
        let after_wrapper = &result[user_fn_idx..];
        // user fn declaration should not have : SV_Target.
        let decl_line = after_wrapper.lines().next().unwrap_or("");
        assert!(!decl_line.contains(": SV_Target"), "User fn should not have SV_Target");
    }

    #[test]
    fn multiple_entries_same_file() {
        let src = r#"import goldy_exp;

[goldy_compute]
[numthreads(64, 1, 1)]
void cs_main(Scattered<uint> data, ThreadId id) { data[id.x] = 0; }

[goldy_vertex]
VSOutput vs_main(Scattered<float4> verts, VertexId vid) {
    VSOutput o; return o;
}
"#;
        let result = transform_virtual_main(src);
        assert!(result.contains("[shader(\"compute\")]"), "Missing compute");
        assert!(result.contains("[shader(\"vertex\")]"), "Missing vertex");
        assert!(result.contains("_goldy_user_cs_main"), "Missing cs rename");
        assert!(result.contains("_goldy_user_vs_main"), "Missing vs rename");
        let goldy_count = result.matches("[goldy_compute]").count() + result.matches("[goldy_vertex]").count();
        assert_eq!(goldy_count, 0, "All [goldy_*] attrs should be removed");
    }

    #[test]
    fn passthrough_param_kept() {
        let src = r#"import goldy_exp;

[goldy_fragment]
float4 fs_main(Broadcast<Uniforms> u, VaryingInput input) : SV_Target {
    return float4(1,0,0,1);
}
"#;
        let result = transform_virtual_main(src);
        // VaryingInput should appear in the generated signature.
        assert!(result.contains("VaryingInput _pt0"), "Pass-through param missing");
    }

    #[test]
    fn ifdef_in_param_list() {
        // Simulates a shader like fine.slang where an #ifdef block lives inside
        // the parameter list (MSAA vs non-MSAA extra params).
        let src = r#"import goldy_exp;

[goldy_compute]
[numthreads(4, 16, 1)]
void cs_main(BufRO<uint> config,
#ifdef msaa
             BufRO<uint> mask_lut, DirectSpatial<float4> out_tex,
#else
             DirectSpatial<float4> out_tex,
#endif
             ThreadId tid) {
}
"#;
        let result = transform_virtual_main(src);

        // Wrapper must be present.
        assert!(result.contains("[shader(\"compute\")]"), "Missing compute attr");
        assert!(result.contains("[numthreads(4, 16, 1)]"), "Missing numthreads");

        // Fixed bindless/user words always present.
        assert!(result.contains("uniform uint _bw0"), "Missing _bw0");
        assert!(result.contains("uniform uint _uw0"), "Missing _uw0");

        // Generated signature must contain the conditional block.
        assert!(result.contains("#ifdef msaa"), "Missing #ifdef msaa in sig");
        assert!(result.contains("#else"), "Missing #else in sig");
        assert!(result.contains("#endif"), "Missing #endif in sig");

        // Both branches must appear in the wrapper (before preprocessing).
        assert!(result.contains("mask_lut"), "Missing mask_lut (msaa branch)");
        assert!(result.contains("out_tex"), "Missing out_tex");

        // ThreadId SV param must follow.
        assert!(result.contains("SV_DispatchThreadID"), "Missing SV_DispatchThreadID");

        // The user function must be renamed.
        assert!(result.contains("_goldy_user_cs_main"), "Missing renamed user fn");

        // The [goldy_compute] attribute must be stripped.
        assert!(!result.contains("[goldy_compute]"), "[goldy_compute] not removed");

        // #line 1 directive must appear between the generated wrapper and user source.
        assert!(result.contains("#line 1"), "Missing #line 1 directive");
    }

    #[test]
    fn line_directive_resets_between_wrapper_and_user_source() {
        // The #line 1 directive must appear *after* the generated wrapper block and
        // *before* the user's code so compiler diagnostics land on the right lines.
        let src = r#"import goldy_exp;

[goldy_compute]
[numthreads(64, 1, 1)]
void cs_main(Scattered<uint> data, ThreadId id) {
    data[id.x] = 0;
}
"#;
        let result = transform_virtual_main(src);
        let wrapper_end = result.find("_goldy_user_cs_main(").unwrap();
        // wrapper body should precede the #line directive
        let line1_pos = result.find("#line 1").unwrap();
        assert!(
            wrapper_end < line1_pos,
            "#line 1 must come after the generated wrapper function"
        );
        // User source (the function definition) must follow the #line directive
        let user_src_pos = result.rfind("void _goldy_user_cs_main(").unwrap();
        assert!(
            line1_pos < user_src_pos,
            "#line 1 must come before the user function definition"
        );
    }

    #[test]
    fn fragment_with_is_front_face() {
        let src = r#"import goldy_exp;

[goldy_fragment]
float4 fs_main(IsFrontFace front) : SV_Target {
    return front.value ? float4(1,0,0,1) : float4(0,0,1,1);
}
"#;
        let result = transform_virtual_main(src);
        assert!(result.contains("[shader(\"fragment\")]"), "Missing fragment attr");
        assert!(result.contains("SV_IsFrontFace"), "Missing SV_IsFrontFace");
        assert!(result.contains("IsFrontFace front"), "User fn must receive IsFrontFace");
        assert!(!result.contains("[goldy_fragment]"), "[goldy_fragment] not removed");
    }

    #[test]
    fn compute_bare_struct_becomes_broadcast() {
        // A plain struct type in a compute shader becomes a broadcast (constant buffer).
        let src = r#"import goldy_exp;

[goldy_compute]
[numthreads(64, 1, 1)]
void cs_main(TimeUniforms cfg, Scattered<uint> data, ThreadId id) {
    data[id.x] = data[id.x] + cfg.base;
}
"#;
        let result = transform_virtual_main(src);
        assert!(result.contains("[shader(\"compute\")]"), "Missing compute attr");
        // cfg at bindless[0] → _bw0 low half.
        assert!(
            result.contains(
                "TimeUniforms cfg = goldy_broadcast<TimeUniforms>(goldy_frame_table_index(_rs0, 0u, _rs1, _rs2))"
            ),
            "Missing broadcast init"
        );
        // data at bindless[1] → _bw0 high half.
        assert!(
            result.contains(
                "Scattered<uint> data = goldy_scattered<uint>(goldy_frame_table_index(_rs0, 1u, _rs1, _rs2))"
            ),
            "Missing resource init"
        );
    }

    #[test]
    fn fragment_bare_struct_is_broadcast_not_passthrough() {
        // In a fragment shader with a resource struct and a varying, the struct
        // (non-last) becomes broadcast; the varying (last) stays PassThrough.
        let src = r#"import goldy_exp;

[goldy_fragment]
float4 fs_main(TimeUniforms cfg, FullscreenVarying input) : SV_Target {
    return float4(cfg.time, 0.0, 0.0, 1.0);
}
"#;
        let result = transform_virtual_main(src);
        // cfg becomes broadcast (bindless slot 0).
        assert!(
            result.contains(
                "TimeUniforms cfg = goldy_broadcast<TimeUniforms>(goldy_frame_table_index(_rs0, 0u, _rs1, _rs2))"
            ),
            "Missing broadcast init"
        );
        // input is the stage varying (PassThrough) — no bindless slot for it.
        assert!(result.contains("FullscreenVarying _pt0"), "Missing passthrough param");
    }

    #[test]
    fn vertex_shader_last_struct_is_passthrough() {
        // In a vertex shader, the last unknown struct is the vertex input (PassThrough).
        let src = r#"import goldy_exp;

[goldy_vertex]
VSOutput vs_main(TimeUniforms cfg, VIn input) {
    VSOutput o;
    return o;
}
"#;
        let result = transform_virtual_main(src);
        // cfg → broadcast at bindless[0].
        assert!(
            result.contains(
                "TimeUniforms cfg = goldy_broadcast<TimeUniforms>(goldy_frame_table_index(_rs0, 0u, _rs1, _rs2))"
            ),
            "Missing broadcast init"
        );
        // VIn input → passthrough (vertex attribute struct).
        assert!(result.contains("VIn _pt0"), "Missing passthrough vertex input");
        assert!(!result.contains("goldy_broadcast<VIn>"), "VIn must not be broadcast");
    }

    #[test]
    fn filter_and_interpolated_resource_types() {
        // Filter → bindless slot; Interpolated<T> → bindless slot.
        let src = r#"import goldy_exp;

[goldy_compute]
[numthreads(8, 8, 1)]
void cs_main(Interpolated<float4> src_tex, Filter samp, Scattered<float4> dst, ThreadId id) {
    float2 uv = float2(id.x, id.y) * 0.01;
    dst[id.x + id.y * 8] = src_tex.Sample(samp, uv);
}
"#;
        let result = transform_virtual_main(src);
        // All three resource params use frame-table slots.
        assert!(result.contains("uniform uint _rs0"), "Missing _rs0");
        assert!(
            result.contains("goldy_interpolated<float4>(goldy_frame_table_index(_rs0, 0u, _rs1, _rs2))"),
            "Missing src_tex init"
        );
        assert!(
            result.contains("goldy_filter(goldy_frame_table_index(_rs0, 1u, _rs1, _rs2))"),
            "Missing samp init"
        );
        assert!(
            result.contains("goldy_scattered<float4>(goldy_frame_table_index(_rs0, 2u, _rs1, _rs2))"),
            "Missing dst init"
        );
        // SV param keeps its SV semantic.
        assert!(result.contains("SV_DispatchThreadID"), "Missing SV_DispatchThreadID");
    }

    // -----------------------------------------------------------------------
    // extract_push_constant_categories tests
    // -----------------------------------------------------------------------

    #[test]
    fn categories_compute_scattered_and_broadcast() {
        use crate::types::ResourceCategory;
        let source = r#"
[goldy_compute]
[numthreads(64, 1, 1)]
void cs_main(TimeUniforms cfg, Scattered<uint> data, ThreadId id) {
    data[id.x] = cfg.base;
}
"#;
        let cats = extract_push_constant_categories(source);
        assert_eq!(cats.len(), 2);
        assert_eq!(cats[0], Some(ResourceCategory::Broadcast));
        assert_eq!(cats[1], Some(ResourceCategory::Scattered));
    }

    #[test]
    fn categories_compute_all_resource_types() {
        use crate::types::ResourceCategory;
        let source = r#"
[goldy_compute]
[numthreads(64, 1, 1)]
void cs_main(
    Scattered<float4> buf,
    BufRO<uint> ro,
    Interpolated<float4> tex,
    DirectSpatial<float4> img,
    Filter sampler,
    ThreadId id
) {}
"#;
        let cats = extract_push_constant_categories(source);
        assert_eq!(cats.len(), 5);
        assert_eq!(cats[0], Some(ResourceCategory::Scattered));
        assert_eq!(cats[1], Some(ResourceCategory::Scattered));
        assert_eq!(cats[2], Some(ResourceCategory::Texture));
        assert_eq!(cats[3], Some(ResourceCategory::StorageImage));
        assert_eq!(cats[4], Some(ResourceCategory::Sampler));
    }

    #[cfg(all(test, feature = "dx12", target_os = "windows"))]
    #[test]
    fn slot_kinds_scattered_vs_bufro() {
        use crate::types::BindlessSlotKind;
        let source = r#"
[goldy_compute]
[numthreads(8, 8, 1)]
void cs_main(Scattered<uint> rw, BufRO<uint> ro, ThreadId id) {}
"#;
        let kinds = extract_push_constant_slot_kinds(source);
        assert_eq!(kinds.len(), 2);
        assert_eq!(kinds[0], Some(BindlessSlotKind::StorageUav));
        assert_eq!(kinds[1], Some(BindlessSlotKind::ReadOnlySrv));
    }

    #[test]
    fn categories_fragment_takes_precedence() {
        use crate::types::ResourceCategory;
        let source = r#"
[goldy_vertex]
void vs_main(Scattered<float4> verts, VertexId vid) {}

[goldy_fragment]
float4 fs_main(Interpolated<float4> tex, Filter samp) : SV_Target {
    return float4(0,0,0,1);
}
"#;
        let cats = extract_push_constant_categories(source);
        assert_eq!(cats.len(), 2);
        assert_eq!(cats[0], Some(ResourceCategory::Texture));
        assert_eq!(cats[1], Some(ResourceCategory::Sampler));
    }

    #[test]
    fn categories_empty_for_non_goldy_source() {
        let source = r#"
[shader("compute")]
[numthreads(64, 1, 1)]
void cs_main(uniform uint _bw0) {}
"#;
        let cats = extract_push_constant_categories(source);
        assert!(cats.is_empty());
    }

    // -----------------------------------------------------------------------
    // extract_binding_element_type_names tests
    // -----------------------------------------------------------------------

    #[test]
    fn binding_type_names_compute_basic() {
        let source = r#"
[goldy_compute]
[numthreads(64, 1, 1)]
void cs_main(TimeUniforms cfg, Scattered<uint> data, ThreadId id) {
    data[id.x] = cfg.base;
}
"#;
        let names = extract_binding_element_type_names(source);
        assert_eq!(names.len(), 2);
        assert_eq!(names[0], Some("TimeUniforms".into()));
        assert_eq!(names[1], Some("uint".into()));
    }

    #[test]
    fn binding_type_names_all_resource_types() {
        let source = r#"
[goldy_compute]
[numthreads(64, 1, 1)]
void cs_main(
    Scattered<float4> buf,
    BufRO<uint> ro,
    Interpolated<float4> tex,
    DirectSpatial<float4> img,
    Filter sampler,
    ThreadId id
) {}
"#;
        let names = extract_binding_element_type_names(source);
        assert_eq!(names.len(), 5);
        assert_eq!(names[0], Some("float4".into()));
        assert_eq!(names[1], Some("uint".into()));
        assert_eq!(names[2], Some("float4".into()));
        assert_eq!(names[3], Some("float4".into()));
        assert_eq!(names[4], None); // Filter has no inner type
    }

    #[test]
    fn categories_bufro_and_scattered_struct() {
        use crate::types::ResourceCategory;
        let source = r#"
            import goldy_exp;
            struct Pair { uint a; uint b; };
            [goldy_compute]
            [numthreads(64, 1, 1)]
            void cs_main(BufRO<Pair> input, Scattered<Pair> output, ThreadId id) {
                output[id.x] = input[id.x];
            }
        "#;
        let cats = extract_push_constant_categories(source);
        assert_eq!(cats.len(), 2);
        assert_eq!(cats[0], Some(ResourceCategory::Scattered));
        assert_eq!(cats[1], Some(ResourceCategory::Scattered));
    }

    #[test]
    fn binding_type_names_bufro_and_scattered_struct() {
        let source = r#"
            import goldy_exp;

            struct Pair { uint a; uint b; };

            [goldy_compute]
            [numthreads(64, 1, 1)]
            void cs_main(BufRO<Pair> input, Scattered<Pair> output, ThreadId id) {
                output[id.x] = input[id.x];
            }
        "#;
        let names = extract_binding_element_type_names(source);
        assert_eq!(names.len(), 2);
        assert_eq!(names[0], Some("Pair".into()));
        assert_eq!(names[1], Some("Pair".into()));
    }

    #[test]
    fn binding_type_names_empty_for_non_goldy() {
        let source = r#"
[shader("compute")]
[numthreads(64, 1, 1)]
void cs_main(uniform uint _bw0) {}
"#;
        let names = extract_binding_element_type_names(source);
        assert!(names.is_empty());
    }

    // -----------------------------------------------------------------------
    // CUDA compute lowering tests
    // -----------------------------------------------------------------------

    #[test]
    fn cuda_compute_passthrough_without_goldy_attrs() {
        let src = r#"
[shader("compute")]
[numthreads(1, 1, 1)]
void cs_main(uniform RWStructuredBuffer<uint> values, uint3 id : SV_DispatchThreadID) {
    values[id.x] = values[id.x] * 2;
}
"#;
        let result = transform_virtual_main_cuda_compute(src).unwrap();
        assert_eq!(result, src);
    }

    #[test]
    fn cuda_compute_scattered_and_thread_id() {
        let src = r#"import goldy_exp;

[goldy_compute]
[numthreads(64, 1, 1)]
void cs_main(Scattered<uint> values, ThreadId id) {
    values[id.x] = values[id.x] * 2;
}
"#;
        let result = transform_virtual_main_cuda_compute(src).unwrap();
        assert!(result.contains("[shader(\"compute\")]"), "{result}");
        assert!(result.contains("[numthreads(64, 1, 1)]"), "{result}");
        assert!(
            result.contains("uniform Scattered<uint> _goldy_cuda_binding_0"),
            "{result}"
        );
        assert!(result.contains("SV_DispatchThreadID"), "{result}");
        assert!(
            result.contains("_goldy_user_cs_main(_goldy_cuda_binding_0, id)"),
            "{result}"
        );
        assert!(!result.contains("[goldy_compute]"), "{result}");
        assert!(!result.contains("_bw0"), "must not use native push ABI: {result}");
    }

    #[test]
    fn cuda_compute_two_buffers_preserve_order() {
        let src = r#"import goldy_exp;

[goldy_compute]
[numthreads(1, 1, 1)]
void cs_main(BufRO<uint> input, Scattered<uint> output, ThreadId id) {
    output[id.x] = input[id.x] * 2;
}
"#;
        let result = transform_virtual_main_cuda_compute(src).unwrap();
        assert!(result.contains("uniform BufRO<uint> _goldy_cuda_binding_0"), "{result}");
        assert!(
            result.contains("uniform Scattered<uint> _goldy_cuda_binding_1"),
            "{result}"
        );
        assert!(
            result.contains("_goldy_user_cs_main(_goldy_cuda_binding_0, _goldy_cuda_binding_1, id)"),
            "{result}"
        );
    }

    #[test]
    fn cuda_compute_broadcast_loads_element_zero() {
        let src = r#"import goldy_exp;

struct Params { uint mul; };

[goldy_compute]
[numthreads(1, 1, 1)]
void cs_main(Params cfg, Scattered<uint> values, ThreadId id) {
    values[id.x] = values[id.x] * cfg.mul;
}
"#;
        let result = transform_virtual_main_cuda_compute(src).unwrap();
        assert!(
            result.contains("uniform StructuredBuffer<Params> _goldy_cuda_binding_0"),
            "{result}"
        );
        assert!(result.contains("Params cfg = _goldy_cuda_binding_0[0];"), "{result}");
        assert!(
            result.contains("_goldy_user_cs_main(cfg, _goldy_cuda_binding_1, id)"),
            "{result}"
        );
    }

    #[test]
    fn cuda_compute_rejects_scalar_and_texture() {
        let scalar = r#"
[goldy_compute]
[numthreads(1, 1, 1)]
void cs_main(Scattered<uint> values, ThreadId id, uint base) {}
"#;
        let err = transform_virtual_main_cuda_compute(scalar).unwrap_err();
        assert!(err.contains("scalar"), "{err}");

        let texture = r#"
[goldy_compute]
[numthreads(1, 1, 1)]
void cs_main(Interpolated<float4> tex, Filter samp, ThreadId id) {}
"#;
        let err = transform_virtual_main_cuda_compute(texture).unwrap_err();
        assert!(err.contains("texture/sampler"), "{err}");
    }
}