#include "common/Exception.h"
#include "tensor/CpuKernels.h"
#include "tensor/KernelGenWgslInternal.h"
#include "tensor/Quant.h"

#include <algorithm>
#include <cmath>
#include <sstream>

namespace eve::tensor {
namespace wgsl_detail {


std::string header(int localX, int localY) {
    return "const workgroupX = " + std::to_string(localX) + "u;\nconst workgroupY = " + std::to_string(localY) + "u;\n";
}
std::string bufferDecl(int binding, const char *name) {
    return "@group(0) @binding(" + std::to_string(binding) + ") var<storage, read_write> " + name + ": array<f32>;\n";
}
std::string bufferDeclUint(int binding, const char *name) {
    return "@group(0) @binding(" + std::to_string(binding) + ") var<storage, read_write> " + name + ": array<u32>;\n";
}

/**
 * WGSL dequantizer `float bval(uint elemIdx)` for a packed weight buffer
 * `bufName` (byte-addressed inside a uint SSBO) plus optional per-group
 * scales `bs`.
 */
std::string emitQuantizedBVal(DType dt, int group, const char *bufName = "b") {
    std::ostringstream os;
    const int          g = group > 0 ? group : 1;
    switch (dt) {
        case DType::Int8:
            os << "fn bval(idx: u32) -> f32 {\n"
               << "  var word: u32 = " << bufName << "[idx >> 2u];\n"
               << "  var byte: u32 = (word >> ((idx & 3u) * 8u)) & 0xFFu;\n"
               << "  var v: i32 = i32(byte);\n"
               << "  if (v >= 128) { v -= 256; }\n"
               << "  return f32(v) * bs[idx / " << g << "u];\n"
               << "}\n";
            break;
        case DType::Int4:
            os << "fn bval(idx: u32) -> f32 {\n"
               << "  var word: u32 = " << bufName << "[idx >> 3u];\n"
               << "  var byte: u32 = (word >> (((idx >> 1u) & 3u) * 8u)) & 0xFFu;\n"
               << "  var nib: u32 = select(byte >> 4u, byte & 0xFu, (idx & 1u) == 0u);\n"
               << "  var v: i32 = i32(nib);\n"
               << "  if (v >= 8) { v -= 16; }\n"
               << "  return f32(v) * bs[idx / " << g << "u];\n"
               << "}\n";
            break;
        case DType::Fp16:
            os << "fn bval(idx: u32) -> f32 {\n"
               << "  var word: u32 = " << bufName << "[idx >> 1u];\n"
               << "  var hb: u32 = (word >> ((idx & 1u) * 16u)) & 0xFFFFu;\n"
               << "  return unpack2x16float(hb).x;\n"
               << "}\n";
            break;
        case DType::Fp8E4M3:
            os << "fn bval(idx: u32) -> f32 {\n"
               << "  var word: u32 = " << bufName << "[idx >> 2u];\n"
               << "  var byte: u32 = (word >> ((idx & 3u) * 8u)) & 0xFFu;\n"
               << "  var s: i32 = select(1, -1, (byte & 0x80u) != 0u);\n"
               << "  var e: i32 = i32((byte >> 3u) & 0xFu);\n"
               << "  var m: i32 = i32(byte & 0x7u);\n"
               << "  var v: f32 = select(exp2(f32(e - 7)) * (1.0 + f32(m) / 8.0),\n"
               << "                      exp2(-6.0) * f32(m) / 8.0, e == 0);\n"
               << "  return f32(s) * v * bs[idx / " << g << "u];\n"
               << "}\n";
            break;
        case DType::Fp4E2M1:
            os << "fn bval(idx: u32) -> f32 {\n"
               << "  var word: u32 = " << bufName << "[idx >> 3u];\n"
               << "  var byte: u32 = (word >> (((idx >> 1u) & 3u) * 8u)) & 0xFFu;\n"
               << "  var nib: u32 = select(byte >> 4u, byte & 0xFu, (idx & 1u) == 0u);\n"
               << "  var s: i32 = select(1, -1, (nib & 8u) != 0u);\n"
               << "  var e: i32 = i32((nib >> 1u) & 3u);\n"
               << "  var m: i32 = i32(nib & 1u);\n"
               << "  var v: f32 = select(exp2(f32(e - 1)) * (1.0 + 0.5 * f32(m)),\n"
               << "                      0.5 * f32(m), e == 0);\n"
               << "  return f32(s) * v * bs[idx / " << g << "u];\n"
               << "}\n";
            break;
        default: break;
    }
    return os.str();
}

std::string pushConstant() { return ""; }

int groupsFor(int count) { return (count + kLocalSize - 1) / kLocalSize; }

std::string scalarStr(float v) {
    if (v == int(v) && std::fabs(v) < 1e9f) return std::to_string(int(v)) + ".0";
    std::ostringstream os;
    os << v << "f";
    return os.str();
}

/**
 * Emit the elementwise chain expression. `rootVar` is the value entering the
 * first chain node (e.g. the matmul accumulator). Input fetches reference
 * buffers by `inputBind` + "[" + indexExpr + "]". Returns the final variable.
 */
struct ChainContext {
    const Graph                    &graph;
    const FusedGroup               &group;
    const std::vector<std::string> &indexExprs;  // per group input
    std::string                     rootVar;
    std::ostringstream             &os;
    int                             temp = 0;
    std::string                     biasIndexExpr;  // e.g. "jj" (matmul) / "f" (conv), empty = none

    std::string operand(int nodeId) {
        if (nodeId == group.biasNode && !biasIndexExpr.empty()) return "bias[" + biasIndexExpr + "]";
        const auto it = std::find(group.inputs.begin(), group.inputs.end(), nodeId);
        if (it != group.inputs.end()) {
            const size_t idx = static_cast<size_t>(it - group.inputs.begin());
            return "a" + std::to_string(idx) + "[" + indexExprs[idx] + "]";
        }
        return vars[static_cast<size_t>(nodeId)];
    }

    std::vector<std::string> vars;

    std::string emit(int nodeId) {
        const GraphNode  &nd = graph.node(nodeId);
        const std::string x =
            (nd.in0 == group.nodes.front() && rootVar != "") ? rootVar : (nd.in0 >= 0 ? operand(nd.in0) : rootVar);
        std::string expr;
        switch (nd.type) {
            case OpType::Add: expr = "(" + x + " + " + operand(nd.in1) + ")"; break;
            case OpType::Sub: expr = "(" + x + " - " + operand(nd.in1) + ")"; break;
            case OpType::Multiply: expr = "(" + x + " * " + operand(nd.in1) + ")"; break;
            case OpType::Divide: expr = "(" + x + " / " + operand(nd.in1) + ")"; break;
            case OpType::AddScalar: expr = "(" + x + " + " + scalarStr(nd.s0) + ")"; break;
            case OpType::SubScalar: expr = "(" + x + " - " + scalarStr(nd.s0) + ")"; break;
            case OpType::MulScalar: expr = "(" + x + " * " + scalarStr(nd.s0) + ")"; break;
            case OpType::DivScalar: expr = "(" + x + " / " + scalarStr(nd.s0) + ")"; break;
            case OpType::PowScalar: expr = "pow(" + x + ", " + scalarStr(nd.s0) + ")"; break;
            case OpType::Neg: expr = "(-" + x + ")"; break;
            case OpType::Abs: expr = "abs(" + x + ")"; break;
            case OpType::Sqrt: expr = "sqrt(" + x + ")"; break;
            case OpType::Exp: expr = "exp(" + x + ")"; break;
            case OpType::Log: expr = "log(" + x + ")"; break;
            case OpType::Sin: expr = "sin(" + x + ")"; break;
            case OpType::Cos: expr = "cos(" + x + ")"; break;
            case OpType::Tanh: expr = "tanh(" + x + ")"; break;
            case OpType::Relu: expr = "max(" + x + ", 0.0)"; break;
            case OpType::Sigmoid: expr = "(1.0 / (1.0 + exp(-" + x + ")))"; break;
            case OpType::Gelu:
                expr = "(0.5 * " + x + " * (1.0 + tanh(0.7978845608028654 * (" + x + " + 0.044715 * " + x + " * " + x +
                       " * " + x + "))))";
                break;
            case OpType::Silu: expr = "(" + x + " / (1.0 + exp(-" + x + ")))"; break;
            case OpType::Clamp: {
                float lo = nd.s0, hi = nd.s1;
                if (lo > hi) std::swap(lo, hi);
                expr = "clamp(" + x + ", " + scalarStr(lo) + ", " + scalarStr(hi) + ")";
                break;
            }
            case OpType::MaximumScalar: expr = "max(" + x + ", " + scalarStr(nd.s0) + ")"; break;
            case OpType::MinimumScalar: expr = "min(" + x + ", " + scalarStr(nd.s0) + ")"; break;
            case OpType::Where:
                expr = "select(" + operand(nd.in2) + ", " + operand(nd.in1) + ", " + operand(nd.in0) + " > 0.5)";
                break;
            default: return "";
        }
        const std::string var = "t" + std::to_string(temp++);
        os << "  var " << var << ": f32 = " << expr << ";\n";
        if (vars.size() <= static_cast<size_t>(nodeId)) vars.resize(static_cast<size_t>(nodeId) + 1);
        vars[static_cast<size_t>(nodeId)] = var;
        return var;
    }
};

/**
 * Per-input flat-index expression for an elementwise group. When an input has
 * the exact output shape, the identity index is used; otherwise a broadcast
 * index is decomposed from the output flat index with baked strides.
 */
void emitInputIndexExprs(std::ostringstream &os, const Graph &g, const FusedGroup &grp,
                         std::vector<std::string> &indexExprs) {
    const GraphNode &on   = g.node(grp.outputNode);
    const int        rank = on.rank;
    std::vector<int> S(static_cast<size_t>(rank), 1);
    if (rank > 0) {
        S[static_cast<size_t>(rank - 1)] = 1;
        for (int k = rank - 2; k >= 0; --k) S[static_cast<size_t>(k)] = S[static_cast<size_t>(k + 1)] * on.dims[k + 1];
    }
    indexExprs.resize(grp.inputs.size());
    for (size_t k = 0; k < grp.inputs.size(); ++k) {
        const GraphNode &inN       = g.node(grp.inputs[k]);
        bool             identical = inN.rank == rank;
        if (identical) {
            for (int d = 0; d < rank; ++d)
                if (inN.dims[d] != on.dims[d]) {
                    identical = false;
                    break;
                }
        }
        if (identical) {
            indexExprs[k] = "i_";
            continue;
        }
        const int        pad = rank - inN.rank;
        std::vector<int> stride(static_cast<size_t>(rank), 0);
        for (int d = 0; d < rank; ++d) {
            const int dim = d < pad ? 1 : inN.dims[d - pad];
            if (dim == 1) continue;
            int s = 1;
            for (int t = d + 1; t < rank; ++t) {
                const int dt = t < pad ? 1 : inN.dims[t - pad];
                if (dt != 1) s *= dt;
            }
            stride[static_cast<size_t>(d)] = s;
        }
        const std::string var = "idx" + std::to_string(k);
        os << "  var " << var << ": u32 = 0u;\n";
        for (int d = 0; d < rank; ++d) {
            if (stride[static_cast<size_t>(d)] == 0) continue;
            os << "  " << var << " += ((i_ / " << S[static_cast<size_t>(d)] << "u) % " << on.dims[d] << "u) * "
               << stride[static_cast<size_t>(d)] << "u;\n";
        }
        indexExprs[k] = var;
    }
}

void genElementwise(const Graph &g, const FusedGroup &grp, KernelSpec &out) {
    if (grp.inputs.size() > size_t(kMaxKernelBindings - 1))
        throw eve::Exception("Tensor WGSL: unsupported kernel variant or binding count");
    const GraphNode   &on = g.node(grp.outputNode);
    std::ostringstream os;
    os << header(kLocalSize);
    for (size_t k = 0; k < grp.inputs.size(); ++k) os << bufferDecl(int(k), ("a" + std::to_string(k)).c_str());
    os << bufferDecl(int(grp.inputs.size()), "o");
    os << pushConstant();
    os << "@compute @workgroup_size(workgroupX, workgroupY)\nfn main(@builtin(global_invocation_id) globalId: "
          "vec3<u32>, @builtin(local_invocation_id) localId: vec3<u32>, @builtin(workgroup_id) groupId: vec3<u32>) {\n";
    os << "  var i_: u32 = globalId.x;\n";
    os << "  if (i_ >= " << on.size << "u) { return; }\n";
    std::vector<std::string> indexExprs;
    emitInputIndexExprs(os, g, grp, indexExprs);
    ChainContext ctx{g, grp, indexExprs, "", os, 0, ""};
    for (int u : grp.nodes) ctx.emit(u);
    os << "  o[i_] = " << ctx.vars[static_cast<size_t>(grp.outputNode)] << ";\n";
    os << "}\n";
    out.pass1.clear();
    out.pass2      = os.str();
    out.groupsX2   = groupsFor(on.size);
    out.inputCount = int(grp.inputs.size());
    return;
}

/** Emit matmul kernel; `tiled` selects the shared-memory 16x16 variant. */
void genMatMul(const Graph &g, const FusedGroup &grp, bool tiled, KernelSpec &out) {
    const GraphNode &mm      = g.node(grp.nodes.front());
    const GraphNode &A       = g.node(mm.in0);
    const GraphNode &B       = g.node(mm.in1);
    const bool       batched = mm.rank == 3;
    const int        batch   = batched ? A.dims[0] : 1;
    const int        m       = A.dims[batched ? 1 : 0];
    const int        k       = A.dims[batched ? 2 : 1];
    const int        n       = B.dims[batched ? 2 : 1];
    const bool       hasBias = grp.biasNode >= 0;
    const bool       bQuant  = q::isQuantDType(static_cast<DType>(B.dtype)) && !B.constBytes.empty();
    if (bQuant && tiled)
        throw eve::Exception("Tensor WGSL: unsupported kernel variant or binding count");  // quantized weights use the
                                                                                           // naive variant only
    const int scalesBinding = hasBias ? 3 : 2;
    const int bindingOut    = bQuant ? (hasBias ? 4 : 3) : (hasBias ? 3 : 2);

    std::ostringstream os;
    if (!tiled) {
        os << header(kLocalSize);
    } else {
        os << header(16, 16);
    }
    os << bufferDecl(0, "a");
    if (bQuant)
        os << bufferDeclUint(1, "b");
    else
        os << bufferDecl(1, "b");
    if (hasBias) os << bufferDecl(2, "bias");
    if (bQuant && B.dtype != static_cast<int>(DType::Fp16)) os << bufferDecl(scalesBinding, "bs");
    os << bufferDecl(bindingOut, "o");
    os << pushConstant();
    if (bQuant) os << emitQuantizedBVal(static_cast<DType>(B.dtype), B.qGroup);

    std::vector<std::string> indexExprs(grp.inputs.size(), "i_");
    ChainContext             ctx{g, grp, indexExprs, "r", os, 0, ""};
    ctx.biasIndexExpr = "jj";

    if (!tiled) {
        os << "@compute @workgroup_size(workgroupX, workgroupY)\nfn main(@builtin(global_invocation_id) globalId: "
              "vec3<u32>, @builtin(local_invocation_id) localId: vec3<u32>, @builtin(workgroup_id) groupId: vec3<u32>) "
              "{\n";
        os << "  var i_: u32 = globalId.x;\n";
        os << "  if (i_ >= " << batch * m * n << "u) { return; }\n";
        if (batched) {
            os << "  var bb: u32 = i_ / " << m * n << "u;\n";
            os << "  var rem: u32 = i_ % " << m * n << "u;\n";
            os << "  var ii: u32 = rem / " << n << "u;\n";
            os << "  var jj: u32 = rem % " << n << "u;\n";
        } else {
            os << "  var ii: u32 = i_ / " << n << "u;\n";
            os << "  var jj: u32 = i_ % " << n << "u;\n";
        }
        os << "  var r: f32 = 0.0;\n";
        os << "  for (var t: u32 = 0u; t < " << k << "u; t++) {\n";
        if (batched) {
            os << "    r += a[bb * " << m * k << "u + ii * " << k << "u + t] * "
               << (bQuant ? "bval(bb * " + std::to_string(k * n) + "u + t * " + std::to_string(n) + "u + jj)"
                          : "b[bb * " + std::to_string(k * n) + "u + t * " + std::to_string(n) + "u + jj]")
               << ";\n";
        } else {
            os << "    r += a[ii * " << k << "u + t] * "
               << (bQuant ? "bval(t * " + std::to_string(n) + "u + jj)" : "b[t * " + std::to_string(n) + "u + jj]")
               << ";\n";
        }
        os << "  }\n";
        for (int u : grp.epilogue) ctx.emit(u);
        const std::string finalExpr =
            grp.epilogue.empty() ? std::string("r") : ctx.vars[static_cast<size_t>(grp.outputNode)];
        os << "  o[i_] = " << finalExpr << ";\n";
        os << "}\n";
        out.groupsX2 = groupsFor(batch * m * n);
    } else {
        ctx.biasIndexExpr = "col";
        const int gx      = (n + 15) / 16;
        const int gy      = (m + 15) / 16;
        os << "var<workgroup> As: array<array<f32, 17>, 16>;\n";
        os << "var<workgroup> Bs: array<array<f32, 17>, 16>;\n";
        os << "@compute @workgroup_size(workgroupX, workgroupY)\nfn main(@builtin(global_invocation_id) globalId: "
              "vec3<u32>, @builtin(local_invocation_id) localId: vec3<u32>, @builtin(workgroup_id) groupId: vec3<u32>) "
              "{\n";
        os << "  var tx: u32 = localId.x;\n";
        os << "  var ty: u32 = localId.y;\n";
        os << "  var row: u32 = globalId.y;\n";
        os << "  var col: u32 = globalId.x;\n";
        os << "  var acc: f32 = 0.0;\n";
        os << "  for (var tile: u32 = 0u; tile < " << (k + 15) / 16 << "u; tile++) {\n";
        os << "    var aCol: u32 = tile * 16u + tx;\n";
        os << "    var bRow: u32 = tile * 16u + ty;\n";
        os << "    As[ty][tx] = 0.0; Bs[ty][tx] = 0.0;\n";
        os << "    if (aCol < " << k << "u && row < " << m << "u) { As[ty][tx] = a[row * " << k << "u + aCol]; }\n";
        os << "    if (bRow < " << k << "u && col < " << n << "u) { Bs[ty][tx] = b[bRow * " << n << "u + col]; }\n";
        os << "    workgroupBarrier();\n";
        os << "    for (var t: u32 = 0u; t < 16u; t++) { acc += As[ty][t] * Bs[t][tx]; }\n";
        os << "    workgroupBarrier();\n";
        os << "  }\n";
        os << "  if (row >= " << m << "u || col >= " << n << "u) { return; }\n";
        os << "  var i_: u32 = row * " << n << "u + col;\n";
        os << "  var r: f32 = acc;\n";
        for (int u : grp.epilogue) ctx.emit(u);
        const std::string finalExpr =
            grp.epilogue.empty() ? std::string("r") : ctx.vars[static_cast<size_t>(grp.outputNode)];
        os << "  o[i_] = " << finalExpr << ";\n}\n";
        out.groupsX2 = gx;
        out.groupsY2 = gy;
    }
    out.pass1.clear();
    out.pass2         = os.str();
    out.inputCount    = hasBias ? 3 : 2;
    out.qDtype        = bQuant ? B.dtype : 0;
    out.qGroup        = B.qGroup;
    out.scalesBinding = bQuant && B.dtype != static_cast<int>(DType::Fp16) ? scalesBinding : -1;
    out.outputBinding = bQuant ? bindingOut : -1;
    if (tiled && batched)
        throw eve::Exception(
            "Tensor WGSL: unsupported kernel variant or binding count");  // batched matmul uses the naive variant
    return;
}

void genConv(const Graph &g, const FusedGroup &grp, KernelSpec &out) {
    const GraphNode &cn     = g.node(grp.nodes.front());
    const GraphNode &X      = g.node(cn.in0);
    const GraphNode &Wt     = g.node(cn.in1);
    const bool       is1d   = cn.type == OpType::Conv1d;
    const int        stride = cn.i0, pad = cn.i1;
    const bool       hasBias     = grp.biasNode >= 0;
    const bool       hasConvBias = cn.in2 >= 0;
    const int        bindingOut  = 2 + int(hasConvBias) + int(hasBias);

    std::ostringstream os;
    os << header(kLocalSize);
    os << bufferDecl(0, "x");
    os << bufferDecl(1, "w");
    if (hasConvBias) os << bufferDecl(2, "convBias");
    if (hasBias) os << bufferDecl(hasConvBias ? 3 : 2, "bias");
    os << bufferDecl(bindingOut, "o");
    os << pushConstant();
    std::vector<std::string> indexExprs(grp.inputs.size(), "i_");
    ChainContext             ctx{g, grp, indexExprs, "r", os, 0, ""};
    ctx.biasIndexExpr = is1d ? "ol" : "ow";

    os << "@compute @workgroup_size(workgroupX, workgroupY)\nfn main(@builtin(global_invocation_id) globalId: "
          "vec3<u32>, @builtin(local_invocation_id) localId: vec3<u32>, @builtin(workgroup_id) groupId: vec3<u32>) {\n";
    os << "  var i_: u32 = globalId.x;\n";
    os << "  if (i_ >= " << g.node(grp.outputNode).size << "u) { return; }\n";
    if (is1d) {
        const int C = X.dims[1], L = X.dims[2];
        const int F = Wt.dims[0], K = Wt.dims[2];
        os << "  var n_: u32 = i_ / (" << F << "u * " << cn.dims[2] << "u);\n";
        os << "  var rem: u32 = i_ % (" << F << "u * " << cn.dims[2] << "u);\n";
        os << "  var f: u32 = rem / " << cn.dims[2] << "u;\n";
        os << "  var ol: u32 = rem % " << cn.dims[2] << "u;\n";
        os << "  var r: f32 = " << (hasConvBias ? "convBias[f]" : "0.0") << ";\n";
        os << "  for (var c: u32 = 0u; c < " << C << "u; c++) {\n";
        os << "    for (var kk: u32 = 0u; kk < " << K << "u; kk++) {\n";
        os << "      var il: i32 = i32(ol) * " << stride << " + i32(kk) - " << pad << ";\n";
        os << "      if (il < 0 || il >= " << L << ") { continue; }\n";
        os << "      r += x[(n_ * " << C << "u + c) * " << L << "u + u32(il)] * w[(f * " << C << "u + c) * " << K
           << "u + kk];\n";
        os << "    }\n";
        os << "  }\n";
    } else {
        const int C = X.dims[1], H = X.dims[2], W = X.dims[3];
        const int F = Wt.dims[0], KH = Wt.dims[2], KW = Wt.dims[3];
        os << "  var n_: u32 = i_ / (" << F << "u * " << cn.dims[2] << "u * " << cn.dims[3] << "u);\n";
        os << "  var rem: u32 = i_ % (" << F << "u * " << cn.dims[2] << "u * " << cn.dims[3] << "u);\n";
        os << "  var f: u32 = rem / (" << cn.dims[2] << "u * " << cn.dims[3] << "u);\n";
        os << "  var rem2: u32 = rem % (" << cn.dims[2] << "u * " << cn.dims[3] << "u);\n";
        os << "  var oh: u32 = rem2 / " << cn.dims[3] << "u;\n";
        os << "  var ow: u32 = rem2 % " << cn.dims[3] << "u;\n";
        os << "  var r: f32 = " << (hasConvBias ? "convBias[f]" : "0.0") << ";\n";
        os << "  for (var c: u32 = 0u; c < " << C << "u; c++) {\n";
        os << "    for (var kh: u32 = 0u; kh < " << KH << "u; kh++) {\n";
        os << "      var ih: i32 = i32(oh) * " << stride << " + i32(kh) - " << pad << ";\n";
        os << "      if (ih < 0 || ih >= " << H << ") { continue; }\n";
        os << "      for (var kw: u32 = 0u; kw < " << KW << "u; kw++) {\n";
        os << "        var iw: i32 = i32(ow) * " << stride << " + i32(kw) - " << pad << ";\n";
        os << "        if (iw < 0 || iw >= " << W << ") { continue; }\n";
        os << "        r += x[((n_ * " << C << "u + c) * " << H << "u + u32(ih)) * " << W << "u + u32(iw)] * w[((f * "
           << C << "u + c) * " << KH << "u + kh) * " << KW << "u + kw];\n";
        os << "      }\n";
        os << "    }\n";
        os << "  }\n";
    }
    for (int u : grp.epilogue) ctx.emit(u);
    const std::string finalExpr =
        grp.epilogue.empty() ? std::string("r") : ctx.vars[static_cast<size_t>(grp.outputNode)];
    os << "  o[i_] = " << finalExpr << ";\n";
    os << "}\n";
    out.pass1.clear();
    out.pass2      = os.str();
    out.groupsX2   = groupsFor(g.node(grp.outputNode).size);
    out.inputCount = 2 + int(hasConvBias) + int(hasBias);
    return;
}

void genPool(const Graph &g, const FusedGroup &grp, KernelSpec &out) {
    const GraphNode   &pn    = g.node(grp.outputNode);
    const GraphNode   &X     = g.node(pn.in0);
    const int          ksize = pn.i0, stride = pn.i1, pad = pn.i2;
    const int          C = X.dims[1], H = X.dims[2], W = X.dims[3];
    const int          OH = pn.dims[2], OW = pn.dims[3];
    const bool         maxPool = pn.type == OpType::MaxPool2d;
    std::ostringstream os;
    os << header(kLocalSize);
    os << bufferDecl(0, "in_");
    os << bufferDecl(1, "o");
    os << pushConstant();
    os << "@compute @workgroup_size(workgroupX, workgroupY)\nfn main(@builtin(global_invocation_id) globalId: "
          "vec3<u32>, @builtin(local_invocation_id) localId: vec3<u32>, @builtin(workgroup_id) groupId: vec3<u32>) {\n";
    os << "  var i_: u32 = globalId.x;\n";
    os << "  if (i_ >= " << pn.size << "u) { return; }\n";
    os << "  var n_: u32 = i_ / (" << C << "u * " << OH << "u * " << OW << "u);\n";
    os << "  var rem: u32 = i_ % (" << C << "u * " << OH << "u * " << OW << "u);\n";
    os << "  var c: u32 = rem / (" << OH << "u * " << OW << "u);\n";
    os << "  var rem2: u32 = rem % (" << OH << "u * " << OW << "u);\n";
    os << "  var oh: u32 = rem2 / " << OW << "u;\n";
    os << "  var ow: u32 = rem2 % " << OW << "u;\n";
    os << "  var acc: f32 = " << (maxPool ? "-3.402823e38" : "0.0") << ";\n";
    os << "  var valid: i32 = 0;\n";
    os << "  for (var kh: i32 = 0; kh < " << ksize << "; kh++) {\n";
    os << "    var ih: i32 = i32(oh) * " << stride << " + kh - " << pad << ";\n";
    os << "    if (ih < 0 || ih >= " << H << ") { continue; }\n";
    os << "    for (var kw: i32 = 0; kw < " << ksize << "; kw++) {\n";
    os << "      var iw: i32 = i32(ow) * " << stride << " + kw - " << pad << ";\n";
    os << "      if (iw < 0 || iw >= " << W << ") { continue; }\n";
    os << "      var v: f32 = in_[((n_ * " << C << "u + c) * " << H << "u + u32(ih)) * " << W << "u + u32(iw)];\n";
    if (maxPool) {
        os << "      acc = max(acc, v);\n";
    } else {
        os << "      acc += v; valid++;\n";
    }
    os << "    }\n";
    os << "  }\n";
    if (!maxPool) os << "  if (valid > 0) { acc /= f32(valid); }\n";
    os << "  o[i_] = acc;\n";
    os << "}\n";
    out.pass1.clear();
    out.pass2      = os.str();
    out.groupsX2   = groupsFor(pn.size);
    out.inputCount = 1;
    return;
}

}  // namespace wgsl_detail

void generateKernelWgslImpl(const Graph &graph, const FusedGroup &group, KernelSpec &out) {
    using namespace wgsl_detail;
    out = KernelSpec{};
    switch (group.kind) {
        case GroupKind::Elementwise: return genElementwise(graph, group, out);
        case GroupKind::MatMul:
            // naive variant first; the runtime autotunes between naive and tiled
            return genMatMul(graph, group, false, out);
        case GroupKind::Conv1d:
        case GroupKind::Conv2d: return genConv(graph, group, out);
        case GroupKind::MaxPool2d:
        case GroupKind::AvgPool2d: return genPool(graph, group, out);
        case GroupKind::Softmax: return genSoftmax(graph, group, out);
        case GroupKind::LayerNorm: return genNorm(graph, group, false, out);
        case GroupKind::RMSNorm: return genNorm(graph, group, true, out);
        case GroupKind::Reduce:
        case GroupKind::ArgMax: return genReduceOrArgmax(graph, group, group.kind == GroupKind::ArgMax, out);
        case GroupKind::Embedding: return genEmbedding(graph, group, out);
        case GroupKind::Concat: return genConcat(graph, group, out);
        case GroupKind::Slice: return genSlice(graph, group, out);
        case GroupKind::Permute: return genPermute(graph, group, out);
        case GroupKind::Resize2d: return genResize2d(graph, group, out);
        case GroupKind::Sdpa: return genSdpa(graph, group, out);
        case GroupKind::Alias: return;  // no kernel; pure buffer alias
    }
    throw eve::Exception("Tensor WGSL: unsupported kernel variant or binding count");
}

Result<KernelSpec> generateWgslKernel(const Graph &graph, const FusedGroup &group, KernelVariant variant) {
    try {
        KernelSpec out;
        if (variant == KernelVariant::TiledMatMul) {
            if (group.kind != GroupKind::MatMul)
                return Result<KernelSpec>::failure(Diagnostic::error(
                    DiagnosticCode::Unsupported, "Tiled WGSL requires a matrix product", "tensor.kernel"));
            wgsl_detail::genMatMul(graph, group, true, out);
        } else {
            generateKernelWgslImpl(graph, group, out);
        }
        wgsl_detail::specializeInputBindings(graph, group, out);
        return Result<KernelSpec>::success(std::move(out));
    } catch (const std::exception &error) {
        return Result<KernelSpec>::failure(
            Diagnostic::error(DiagnosticCode::Unsupported, error.what(), "tensor.kernel"));
    }
}

}  // namespace eve::tensor