local vec3 = lovr.math.vec3
local mat4 = lovr.math.mat4
local quat = lovr.math.quat

local nParticles = 2 ^ 14
local batchSize = 128
local batches = nParticles / batchSize
local dt = 0.01

local features
local shaders = {}
local blocks = {}
local meshes = {}
local compute = {}
local hands = {}

function lovr.load()
    features = lovr.graphics.getFeatures()
    shaders.particles = lovr.graphics.newShader("shaders/particles.vert", "shaders/particles.frag")

    blocks.particles = lovr.graphics.newShaderBlock("compute", {
        -- Reading a value seems to work around whatever the rendering glitch
        -- was. So we include this flag so we only have to ask for a single
        -- value, instead of having to copy a large array. Still, annoying :(
        -- For reference, see issue comment at:
        -- https://github.com/bjornbytes/lovr/issues/211#issuecomment-580985010
        reloadFlag = "float",
        particles = {"mat4", nParticles},
    }, { readable = true })
    shaders.particles:sendBlock('particleData', blocks.particles)
    shaders.particles:send("batchSize", batchSize)

    meshes.particle = makeParticleMesh(batchSize)

    compute.particles = lovr.graphics.newComputeShader("shaders/particles.comp")
    compute.particles:sendBlock('particleData', blocks.particles)
    compute.particles:send("dt", dt)

    local xr, yr, zr = range(-1, 1), range(1, 3), range(-1, -0.5)
    local sr, vr = range(0.005, 0.01), range(-0.05, 0.05)
    local initial = initialParticleData(nParticles, xr, yr, zr, sr, vr)
    blocks.particles:send("particles", initial)

    for _, hand in ipairs(lovr.headset.getHands()) do
        local x, y, z = lovr.headset.getPose(hand)
        hands[hand] = {
            pos = lovr.math.newVec3(x, y, z),
            vel = lovr.math.newVec3(0, 0, 0),
            rot = lovr.math.newQuat(),
        }
    end
end

function lovr.draw()
    lovr.graphics.clear()
    local x, y, z = lovr.headset.getPose("head")
    shaders.particles:send("headPos", vec3(x, y, z))
    local positions = {}
    local velocities = {}
    for _, hand in ipairs(lovr.headset.getHands()) do
        local x, y, z, a, ax, ay, az = lovr.headset.getPose(hand)
        local pos = vec3(x, y, z)
        hands[hand].vel:set((pos - hands[hand].pos) / dt)
        hands[hand].pos:set(pos)
        hands[hand].rot:set(a, ax, ay, az)
        table.insert(positions, pos)
        table.insert(velocities, hands[hand].vel)
    end
    compute.particles:send("handPos", positions)
    compute.particles:send("handVel", velocities)

    lovr.graphics.compute(compute.particles, batches)
    -- Read the buffer to force SSBO synch of some sort.
    blocks.particles:read("reloadFlag")

    lovr.graphics.setColor(0.6, 0.1, 0)
    lovr.graphics.setShader(nil)
    for _, hand in pairs(hands) do
        -- local x, y, z = hand.pos:unpack()
        lovr.graphics.cube("fill", hand.pos, 0.1, hand.rot)
    end

    withDepthTest("lequal", false, function()
        lovr.graphics.setColor(1, 1, 1)
        lovr.graphics.setShader(shaders.particles)
        meshes.particle:draw(mat4(), batches)
    end)
end

function makeParticleMesh(size)
    local format = { {'vertPosition', 'float', 2}, }
    local verts = {}
    local vertMap = {}
    local insert = table.insert
    for chunk = 1, size do
        insert(verts, { 1,  1})
        insert(verts, { 1, -1})
        insert(verts, {-1, -1})
        insert(verts, {-1,  1})
        local base = 4 * (chunk - 1)
        insert(vertMap, 3 + base)
        insert(vertMap, 2 + base)
        insert(vertMap, 1 + base)
        insert(vertMap, 4 + base)
        insert(vertMap, 3 + base)
        insert(vertMap, 1 + base)
    end
    local mesh = lovr.graphics.newMesh(format, verts, "triangles", "static")
    mesh:setVertexMap(vertMap)
    return mesh
end

function range(min, max)
    local scale = max - min
    return function()
        return min + lovr.math.random() * scale
    end
end

function initialParticleData(size, xRange, yRange, zRange, sRange, vRange)
    local particles = {}
    local insert = table.insert
    for i = 1, size do
        insert(particles, {
            xRange(), yRange(), zRange(), sRange(),
            vRange(), vRange(), vRange(), 1,
            0, 0, 0, 0,
            0, 0, 0, 0,
        })
    end
    return particles
end

function withDepthTest(compareMode, write, callback)
    local oldComp, oldWrite = lovr.graphics.getDepthTest()
    lovr.graphics.setDepthTest(compareMode, write)
    callback()
    lovr.graphics.setDepthTest(oldComp, oldWrite)
end