local N             = 1024           -- elements per vector
local LOCAL_SIZE    = 32             -- match add_*.comp local_size_x
local GROUPS_X      = N / LOCAL_SIZE -- work-groups in X
local SHOW_WARN     = true           -- print GLSL driver warnings

-- Capability test  (compile a 1-line compute shader)
local function hasCompute()
  if not love.graphics.newComputeShader then return false end
  return pcall(love.graphics.newComputeShader,
               "layout(local_size_x = 1) in; void computemain(){}")
end

-- Check hardware limits and bail if we'd overflow them
local function checkLimits()
  local lim = love.graphics.getSystemLimits()
  local wg   = lim.maxcomputeworkgroupsize   or {1024,1024,64}
  local grp  = lim.maxcomputeworkgroupcount  or {65535,65535,65535}
  local inv  = lim.maxcomputeinvocations     or wg[1]

  assert(LOCAL_SIZE <= wg[1] , "local_size_x too large")
  assert(GROUPS_X  <= grp[1] , "dispatch X too large")
  assert(LOCAL_SIZE <= inv   , "threads / WG too many")
end

-- Helpers to make textures / buffers
local function makeTexture(gen)
  local img = love.image.newImageData(N,1,"rgba32f")
  for x=0,N-1 do img:setPixel(x,0,gen(x),0,0,1) end
  return love.graphics.newTexture(img,
          {format="rgba32f", computewrite=true})
end

-- Helper: build a shader-storage buffer (SSBO) and upload data.
local function makeBuffer(gen)
  -- build a flat Lua array with the floats
  local tmp = {}
  for i = 1, N do tmp[i] = gen(i - 1) end     -- generator is 0-based
  -- create an empty float-format buffer
  local buf = love.graphics.newBuffer("float", N, { shaderstorage = true })

  -- upload – setArrayData is the function in the latest build
  if buf.setArrayData then
    buf:setArrayData(tmp)                     -- <-- THIS is the right call
  else
    error("Buffer:setArrayData missing – mismatched engine build?")
  end

  return buf
end

-- Pure-Lua fallback (only used if GPU lacks compute)
local function addCPU(dst,a,b)
  for i=1,N do dst[i]=a[i]+b[i] end
end

-- Globals for love.draw
local sums_tex, sums_buf = {}, {}
local usedGPU = false

function love.load()
  usedGPU = hasCompute()
  print(usedGPU and "Compute detected – using GPU"
                  or "No compute – CPU fallback")

  local genA = function(i)return i      end
  local genB = function(i)return N-i    end

  if usedGPU then
    checkLimits()

    ---------------- Texture (image2D) pass ----------------------
    local texA,texB = makeTexture(genA), makeTexture(genB)
    local texOut = love.graphics.newTexture(N,1,
                   {format="rgba32f", computewrite=true})

    local csImg = love.graphics.newComputeShader("add_image.comp")
    csImg:send("A",texA); csImg:send("B",texB); csImg:send("Out",texOut)

    local warn = csImg:getWarnings()
    if SHOW_WARN and warn~="" then print("Image warn: "..warn) end

    love.graphics.dispatchThreadgroups(csImg,GROUPS_X,1,1)

    local imgData = love.graphics.readbackTexture(texOut)
    for x=0,N-1 do sums_tex[x+1]=imgData:getPixel(x,0) end

    ---------------- Buffer (SSBO) pass --------------------------
    local bufA,bufB = makeBuffer(genA), makeBuffer(genB)
    local bufOut    = love.graphics.newBuffer("float", N,
                                              {shaderstorage=true})

    local csBuf = love.graphics.newComputeShader("add_buffer.comp")
    csBuf:send("A",bufA); csBuf:send("B",bufB); csBuf:send("Out",bufOut)

    warn = csBuf:getWarnings()
    if SHOW_WARN and warn~="" then print("Buffer warn: "..warn) end

    love.graphics.dispatchThreadgroups(csBuf,GROUPS_X,1,1)

    local raw = love.graphics.readbackBuffer(bufOut)  -- returns Data
    for i = 1, N do
      sums_buf[i] = raw:getFloat((i - 1) * 4)         -- 4 bytes per float
    end

  else
    --------------- CPU reference path --------------------------
    local A,B={},{}; for i=1,N do A[i]=genA(i-1); B[i]=genB(i-1) end
    addCPU(sums_tex,A,B); sums_buf=sums_tex            -- mirror for display
  end
end

--------------------------------------------------------------------
function love.draw()
  love.graphics.print("Compute demo – first 10 elements",10,10)
  for i=1,10 do
    local a,b=i-1,N-(i-1)
    love.graphics.print(
      string.format("%3d+%3d = %3d (tex)  %3d (buf)",
        a,b,sums_tex[i]or 0,sums_buf[i]or 0),
      10, 30+i*16)
  end
  if not usedGPU then
    love.graphics.print("**CPU fallback – compute unavailable**",10,230)
  end
end
