-- webp.lua — Pure Lua VP8L (lossless WebP) decoder for Love2D (LuaJIT)
-- Supports: VP8L (lossless) only. VP8 (lossy) is not feasible in pure Lua.
-- Usage: local WebP = require("webp"); local img = WebP.load("sprite.webp")

local WebP = {}

local bit = require("bit")
local band   = bit.band
local bor    = bit.bor
local lshift = bit.lshift
local rshift = bit.rshift
local tobit  = bit.tobit

-- ─── Bitstream reader ────────────────────────────────────────────────────────
-- VP8L is LSB-first. We maintain a 32-bit window since LuaJIT bit ops are 32-bit.

local function newBitReader(data)
    local r = {
        data   = data,
        pos    = 1,      -- next byte to load (1-based)
        window = 0,      -- current bit window (32-bit)
        bits   = 0,      -- valid bits in window
    }

    function r:fill()
        while self.bits <= 24 and self.pos <= #self.data do
            local byte = self.data:byte(self.pos)
            self.window = bor(self.window, lshift(byte, self.bits))
            self.bits   = self.bits + 8
            self.pos    = self.pos + 1
        end
    end

    function r:read(n)
        if self.bits < n then self:fill() end
        local v = band(self.window, lshift(1, n) - 1)
        self.window = rshift(self.window, n)
        self.bits   = self.bits - n
        return v
    end

    function r:readBool()
        return self:read(1) == 1
    end

    r:fill()
    return r
end

-- ─── Huffman trees ───────────────────────────────────────────────────────────

local function buildHuffmanTable(codeLengths)
    local n = #codeLengths
    local counts = {}
    for i = 0, 15 do counts[i] = 0 end
    for _, cl in ipairs(codeLengths) do
        if cl > 0 then counts[cl] = counts[cl] + 1 end
    end

    local nextCode = {}
    local code = 0
    counts[0] = 0
    for bits = 1, 15 do
        code = lshift(code + counts[bits - 1], 1)
        nextCode[bits] = code
    end

    local htable = {}
    for i = 1, n do
        local len = codeLengths[i]
        if len > 0 then
            htable[nextCode[len]] = { sym = i - 1, len = len }
            nextCode[len] = nextCode[len] + 1
        end
    end
    return htable
end

local function decodeHuffman(br, htable)
    local code = 0
    for len = 1, 15 do
        code = bor(lshift(code, 1), br:read(1))
        local entry = htable[code]
        if entry and entry.len == len then
            return entry.sym
        end
    end
    error("Invalid Huffman code")
end

-- ─── Code length decoding ────────────────────────────────────────────────────

local CODE_LENGTH_ORDER = {17,18,0,1,2,3,4,5,16,6,7,8,9,10,11,12,13,14,15}

local function readCodeLengths(br, n)
    local numCLCodes = br:read(4) + 4
    local clLengths = {}
    for i = 1, 19 do clLengths[i] = 0 end
    for i = 1, numCLCodes do
        clLengths[CODE_LENGTH_ORDER[i] + 1] = br:read(3)
    end
    local clTable = buildHuffmanTable(clLengths)

    local lengths = {}
    local prev = 8
    while #lengths < n do
        local sym = decodeHuffman(br, clTable)
        if sym <= 15 then
            lengths[#lengths + 1] = sym
            if sym ~= 0 then prev = sym end
        elseif sym == 16 then
            local rep = br:read(2) + 3
            for _ = 1, rep do lengths[#lengths + 1] = prev end
        elseif sym == 17 then
            local rep = br:read(3) + 3
            for _ = 1, rep do lengths[#lengths + 1] = 0 end
        elseif sym == 18 then
            local rep = br:read(7) + 11
            for _ = 1, rep do lengths[#lengths + 1] = 0 end
        end
    end
    return lengths
end

-- ─── Prefix code reading ─────────────────────────────────────────────────────

local function readPrefixCode(br, alphabetSize)
    local simpleCode = br:read(1)
    if simpleCode == 1 then
        local numSyms = br:read(1) + 1
        local firstSym = br:read(1)
        local sym1 = br:read(firstSym == 1 and 8 or 1)
        local lengths = {}
        for i = 1, alphabetSize do lengths[i] = 0 end
        lengths[sym1 + 1] = 1
        if numSyms == 2 then
            local sym2 = br:read(8)
            lengths[sym2 + 1] = 1
        end
        return buildHuffmanTable(lengths)
    else
        local lengths = readCodeLengths(br, alphabetSize)
        return buildHuffmanTable(lengths)
    end
end

-- ─── Color cache ─────────────────────────────────────────────────────────────

local function newColorCache(bits)
    local size = lshift(1, bits)
    local cache = {}
    for i = 0, size - 1 do cache[i] = 0 end
    return {
        size = size,
        data = cache,
        insert = function(self, color)
            -- hash: (0x1e35a7bd * color) >> (32 - bits), masked to cache size
            local hash = tobit(0x1e35a7bd * color)
            local idx  = band(rshift(hash, 32 - bits), size - 1)
            self.data[idx] = color
        end,
        lookup = function(self, idx)
            return self.data[idx]
        end,
    }
end

-- ─── Transform type constants ────────────────────────────────────────────────

local TRANSFORM_PREDICTOR      = 0
local TRANSFORM_COLOR          = 1
local TRANSFORM_SUBTRACT_GREEN = 2
local TRANSFORM_COLOR_INDEXING = 3

-- ─── Prefix length/distance tables ──────────────────────────────────────────

local PREFIX_EXTRA_BITS = {
    0,0,0,0, 1,1,2,2, 3,3,4,4, 5,5,6,6, 7,7,8,8, 9,9,10,10, 11,11,12,12, 13,13
}
local PREFIX_OFFSET = {
    0,1,2,3, 4,6,8,12, 16,24,32,48, 64,96,128,192, 256,384,512,768,
    1024,1536,2048,3072, 4096,6144,8192,12288, 16384,24576
}

local function prefixToValue(br, code)
    if code < 4 then return code end
    local extra  = PREFIX_EXTRA_BITS[code + 1] or 0
    local offset = PREFIX_OFFSET[code + 1] or 0
    return offset + br:read(extra)
end

-- ─── VP8L distance offset table (120 entries) ────────────────────────────────

local DIST_MAP = {
    {0,1},{1,0},{1,1},{-1,1},{0,2},{2,0},{1,2},{-1,2},
    {2,1},{-2,1},{2,2},{-2,2},{0,3},{3,0},{1,3},{-1,3},
    {3,1},{-3,1},{2,3},{-2,3},{3,2},{-3,2},{0,4},{4,0},
    {1,4},{-1,4},{4,1},{-4,1},{3,3},{-3,3},{2,4},{-2,4},
    {4,2},{-4,2},{0,5},{3,4},{-3,4},{4,3},{-4,3},{5,0},
    {1,5},{-1,5},{5,1},{-5,1},{2,5},{-2,5},{5,2},{-5,2},
    {4,4},{-4,4},{3,5},{-3,5},{5,3},{-5,3},{0,6},{6,0},
    {1,6},{-1,6},{6,1},{-6,1},{2,6},{-2,6},{6,2},{-6,2},
    {4,5},{-4,5},{5,4},{-5,4},{3,6},{-3,6},{6,3},{-6,3},
    {0,7},{7,0},{1,7},{-1,7},{5,5},{-5,5},{7,1},{-7,1},
    {4,6},{-4,6},{6,4},{-6,4},{2,7},{-2,7},{7,2},{-7,2},
    {3,7},{-3,7},{7,3},{-7,3},{5,6},{-5,6},{6,5},{-6,5},
    {8,0},{4,7},{-4,7},{7,4},{-7,4},{8,1},{8,2},{6,6},
    {-6,6},{8,3},{5,7},{-5,7},{7,5},{-7,5},{8,4},{6,7},
    {-6,7},{7,6},{-7,6},{8,5},{8,6},{7,7},{-7,7},{8,7},
}

-- ─── Main VP8L decode (recursive for transform sub-images) ───────────────────

local function decodeVP8L(br, width, height)

    -- color cache
    local hasCCache  = br:readBool()
    local ccache     = nil
    local ccacheBits = 0
    if hasCCache then
        ccacheBits = br:read(4)
        ccache     = newColorCache(ccacheBits)
    end

    -- collect transforms (applied in reverse after decode)
    local transforms = {}
    while br:readBool() do
        local ttype = br:read(2)
        local t = { ttype = ttype }
        if ttype == TRANSFORM_PREDICTOR or ttype == TRANSFORM_COLOR then
            t.sizeBits = br:read(3) + 2
            local tw = math.floor((width  + lshift(1, t.sizeBits) - 1) / lshift(1, t.sizeBits))
            local th = math.floor((height + lshift(1, t.sizeBits) - 1) / lshift(1, t.sizeBits))
            t.data = decodeVP8L(br, tw, th)
        elseif ttype == TRANSFORM_COLOR_INDEXING then
            t.numColors = br:read(8) + 1
            t.colors    = decodeVP8L(br, t.numColors, 1)
        end
        -- SUBTRACT_GREEN has no extra data
        transforms[#transforms + 1] = t
    end

    -- entropy/meta-Huffman image
    local groupBits  = 0
    local groupImage = nil
    local numGroups  = 1
    if br:readBool() then
        groupBits = br:read(3) + 2
        local gw  = math.floor((width  + lshift(1, groupBits) - 1) / lshift(1, groupBits))
        local gh  = math.floor((height + lshift(1, groupBits) - 1) / lshift(1, groupBits))
        groupImage = decodeVP8L(br, gw, gh)
        local maxG = 0
        for _, v in ipairs(groupImage) do
            local g = band(rshift(v, 8), 0xffff)
            if g > maxG then maxG = g end
        end
        numGroups = maxG + 1
    end

    -- alphabet sizes
    local ccSize    = hasCCache and lshift(1, ccacheBits) or 0
    local alphabetG = 256 + 24 + ccSize

    -- read Huffman tables for each group (G, R, B, A + distance)
    local huffGroups = {}
    for g = 1, numGroups do
        huffGroups[g] = {
            G    = readPrefixCode(br, alphabetG),
            R    = readPrefixCode(br, 256),
            B    = readPrefixCode(br, 256),
            A    = readPrefixCode(br, 256),
            dist = readPrefixCode(br, 40),
        }
    end

    -- decode pixels
    local pixels    = {}
    local numPixels = width * height
    local px, py    = 0, 0

    local function getGroup()
        if not groupImage then return huffGroups[1] end
        local gx  = rshift(px, groupBits)
        local gy  = rshift(py, groupBits)
        local gw  = math.floor((width + lshift(1, groupBits) - 1) / lshift(1, groupBits))
        local idx = gy * gw + gx + 1
        local g   = band(rshift(groupImage[idx] or 0, 8), 0xffff)
        return huffGroups[g + 1] or huffGroups[1]
    end

    while #pixels < numPixels do
        local hg   = getGroup()
        local code = decodeHuffman(br, hg.G)

        if code < 256 then
            -- literal ARGB (green first, then R, B, A)
            local g = code
            local r = decodeHuffman(br, hg.R)
            local b = decodeHuffman(br, hg.B)
            local a = decodeHuffman(br, hg.A)
            local color = bor(bor(bor(lshift(a, 24), lshift(r, 16)), lshift(g, 8)), b)
            pixels[#pixels + 1] = color
            if ccache then ccache:insert(color) end

        elseif code < 256 + 24 then
            -- LZ77 back-reference
            local lenCode  = code - 256
            local length   = prefixToValue(br, lenCode) + 1
            local distCode = decodeHuffman(br, hg.dist)
            local dist     = prefixToValue(br, distCode) + 1

            -- remap small distances through VP8L spatial table
            if dist <= 120 then
                local d    = DIST_MAP[dist]
                local srcX = px - d[1]
                local srcY = py - d[2]
                dist = py * width + px - (srcY * width + srcX)
            end

            local src = #pixels - dist + 1
            for i = 0, length - 1 do
                local c = pixels[src + i] or 0
                pixels[#pixels + 1] = c
                if ccache then ccache:insert(c) end
            end

        else
            -- color cache reference
            local cacheIdx = code - 256 - 24
            pixels[#pixels + 1] = ccache:lookup(cacheIdx)
        end

        px = px + 1
        if px >= width then px = 0; py = py + 1 end
    end

    -- ─── Apply transforms in reverse order ───────────────────────────────────

    for i = #transforms, 1, -1 do
        local t = transforms[i]

        -- SUBTRACT_GREEN: R += G, B += G (mod 256)
        if t.ttype == TRANSFORM_SUBTRACT_GREEN then
            for idx = 1, #pixels do
                local c = pixels[idx]
                local a = band(rshift(c, 24), 0xff)
                local r = band(rshift(c, 16), 0xff)
                local g = band(rshift(c,  8), 0xff)
                local b = band(c, 0xff)
                r = band(r + g, 0xff)
                b = band(b + g, 0xff)
                pixels[idx] = bor(bor(bor(lshift(a, 24), lshift(r, 16)), lshift(g, 8)), b)
            end

        -- COLOR_INDEXING: replace each pixel's green channel with palette entry
        elseif t.ttype == TRANSFORM_COLOR_INDEXING then
            local bpp = 8
            if     t.numColors <= 2  then bpp = 1
            elseif t.numColors <= 4  then bpp = 2
            elseif t.numColors <= 16 then bpp = 4 end

            local newPixels = {}
            if bpp == 8 then
                for _, c in ipairs(pixels) do
                    local idx = band(rshift(c, 8), 0xff)
                    newPixels[#newPixels + 1] = t.colors[idx + 1] or 0
                end
            else
                local pxPerByte = 8 / bpp
                local mask = lshift(1, bpp) - 1
                for _, c in ipairs(pixels) do
                    local packed = band(rshift(c, 8), 0xff)
                    for p = 0, pxPerByte - 1 do
                        if #newPixels < width * height then
                            local idx = band(rshift(packed, p * bpp), mask)
                            newPixels[#newPixels + 1] = t.colors[idx + 1] or 0
                        end
                    end
                end
            end
            pixels = newPixels

        -- PREDICTOR: undo per-block spatial prediction
        elseif t.ttype == TRANSFORM_PREDICTOR then
            local sb = t.sizeBits
            local tw = math.floor((width + lshift(1, sb) - 1) / lshift(1, sb))

            local function getpx(x, y)
                if x < 0 then x = 0 end
                if y < 0 then return tobit(0xff000000) end
                return pixels[y * width + x + 1] or 0
            end

            local function addARGB(a, b)
                local aa  = band(rshift(a, 24), 0xff)
                local ar  = band(rshift(a, 16), 0xff)
                local ag  = band(rshift(a,  8), 0xff)
                local ab  = band(a, 0xff)
                local ba  = band(rshift(b, 24), 0xff)
                local br_ = band(rshift(b, 16), 0xff)
                local bg  = band(rshift(b,  8), 0xff)
                local bb  = band(b, 0xff)
                return bor(bor(bor(
                    lshift(band(aa + ba,  0xff), 24),
                    lshift(band(ar + br_, 0xff), 16)),
                    lshift(band(ag + bg,  0xff),  8)),
                    band(ab + bb, 0xff))
            end

            local function avgARGB(a, b)
                local aa  = band(rshift(a, 24), 0xff)
                local ar  = band(rshift(a, 16), 0xff)
                local ag  = band(rshift(a,  8), 0xff)
                local ab  = band(a, 0xff)
                local ba  = band(rshift(b, 24), 0xff)
                local br_ = band(rshift(b, 16), 0xff)
                local bg  = band(rshift(b,  8), 0xff)
                local bb  = band(b, 0xff)
                return bor(bor(bor(
                    lshift(rshift(aa + ba,  1), 24),
                    lshift(rshift(ar + br_, 1), 16)),
                    lshift(rshift(ag + bg,  1),  8)),
                    rshift(ab + bb, 1))
            end

            local function clampByte(v)
                return math.max(0, math.min(255, v))
            end

            local function clampAddARGB(a, b)
                local aa  = band(rshift(a, 24), 0xff)
                local ar  = band(rshift(a, 16), 0xff)
                local ag  = band(rshift(a,  8), 0xff)
                local ab  = band(a, 0xff)
                local ba  = band(rshift(b, 24), 0xff)
                local br_ = band(rshift(b, 16), 0xff)
                local bg  = band(rshift(b,  8), 0xff)
                local bb  = band(b, 0xff)
                return bor(bor(bor(
                    lshift(clampByte(aa + ba),  24),
                    lshift(clampByte(ar + br_), 16)),
                    lshift(clampByte(ag + bg),   8)),
                    clampByte(ab + bb))
            end

            local function subARGB(a, b)
                local aa  = band(rshift(a, 24), 0xff)
                local ar  = band(rshift(a, 16), 0xff)
                local ag  = band(rshift(a,  8), 0xff)
                local ab  = band(a, 0xff)
                local ba  = band(rshift(b, 24), 0xff)
                local br_ = band(rshift(b, 16), 0xff)
                local bg  = band(rshift(b,  8), 0xff)
                local bb  = band(b, 0xff)
                return bor(bor(bor(
                    lshift(band(aa - ba,  0xff), 24),
                    lshift(band(ar - br_, 0xff), 16)),
                    lshift(band(ag - bg,  0xff),  8)),
                    band(ab - bb, 0xff))
            end

            local function halfSubARGB(a, b)
                local aa  = band(rshift(a, 24), 0xff)
                local ar  = band(rshift(a, 16), 0xff)
                local ag  = band(rshift(a,  8), 0xff)
                local ab  = band(a, 0xff)
                local ba  = band(rshift(b, 24), 0xff)
                local br_ = band(rshift(b, 16), 0xff)
                local bg  = band(rshift(b,  8), 0xff)
                local bb  = band(b, 0xff)
                return bor(bor(bor(
                    lshift(rshift(aa - ba,  1), 24),
                    lshift(rshift(ar - br_, 1), 16)),
                    lshift(rshift(ag - bg,  1),  8)),
                    rshift(ab - bb, 1))
            end

            local function selectARGB(l, tp, tl)
                local function absdiff(x, y, s)
                    return math.abs(band(rshift(x, s), 0xff) - band(rshift(y, s), 0xff))
                end
                local function dist(x, y)
                    return absdiff(x,y,24) + absdiff(x,y,16) + absdiff(x,y,8) + absdiff(x,y,0)
                end
                return dist(l, tl) <= dist(tp, tl) and l or tp
            end

            for iy = 0, height - 1 do
                for ix = 0, width - 1 do
                    if not (ix == 0 and iy == 0) then
                        local pidx = iy * width + ix + 1
                        local L  = getpx(ix - 1, iy)
                        local T  = getpx(ix,     iy - 1)
                        local TL = getpx(ix - 1, iy - 1)
                        local TR = getpx(ix + 1, iy - 1)
                        local tx   = rshift(ix, sb)
                        local ty   = rshift(iy, sb)
                        local mode = band(t.data[ty * tw + tx + 1] or 0, 0xff)

                        local pred
                        if    mode == 0  then pred = tobit(0xff000000)
                        elseif mode == 1  then pred = L
                        elseif mode == 2  then pred = T
                        elseif mode == 3  then pred = TR
                        elseif mode == 4  then pred = TL
                        elseif mode == 5  then pred = avgARGB(avgARGB(L, TR), T)
                        elseif mode == 6  then pred = avgARGB(L, TL)
                        elseif mode == 7  then pred = avgARGB(L, T)
                        elseif mode == 8  then pred = avgARGB(TL, T)
                        elseif mode == 9  then pred = avgARGB(T, TR)
                        elseif mode == 10 then pred = avgARGB(avgARGB(L, TL), avgARGB(T, TR))
                        elseif mode == 11 then pred = selectARGB(L, T, TL)
                        elseif mode == 12 then pred = clampAddARGB(L, subARGB(T, TL))
                        elseif mode == 13 then
                            local avg = avgARGB(L, T)
                            pred = clampAddARGB(avg, halfSubARGB(avg, TL))
                        else pred = L end

                        pixels[pidx] = addARGB(pixels[pidx], pred)
                    end
                end
            end

        -- COLOR: undo green/red channel correlations
        elseif t.ttype == TRANSFORM_COLOR then
            local sb = t.sizeBits
            local tw = math.floor((width + lshift(1, sb) - 1) / lshift(1, sb))
            for iy = 0, height - 1 do
                for ix = 0, width - 1 do
                    local pidx = iy * width + ix + 1
                    local c    = pixels[pidx] or 0
                    local tx   = rshift(ix, sb)
                    local ty   = rshift(iy, sb)
                    local m    = t.data[ty * tw + tx + 1] or 0

                    -- unpack signed bytes from transform pixel (stored as ARGB)
                    local g2r = band(rshift(m, 16), 0xff)
                    local r2b = band(rshift(m,  8), 0xff)
                    local g2b = band(m, 0xff)
                    if g2r >= 128 then g2r = g2r - 256 end
                    if r2b >= 128 then r2b = r2b - 256 end
                    if g2b >= 128 then g2b = g2b - 256 end

                    local a = band(rshift(c, 24), 0xff)
                    local r = band(rshift(c, 16), 0xff)
                    local g = band(rshift(c,  8), 0xff)
                    local b = band(c, 0xff)

                    r = band(r + math.floor(g2r * g / 256), 0xff)
                    b = band(b + math.floor(g2b * g / 256) + math.floor(r2b * r / 256), 0xff)

                    pixels[pidx] = bor(bor(bor(lshift(a, 24), lshift(r, 16)), lshift(g, 8)), b)
                end
            end
        end
    end

    return pixels
end

-- ─── Public API ──────────────────────────────────────────────────────────────

function WebP.decode(data)
    assert(data:sub(1, 4) == "RIFF", "Not a RIFF file")
    assert(data:sub(9, 12) == "WEBP", "Not a WEBP file")

    local fourCC = data:sub(13, 16)
    assert(fourCC == "VP8L", "Only lossless VP8L WebP is supported (got: " .. fourCC .. ")")

    -- byte 21 = VP8L signature (0x2f), bitstream starts at byte 22
    assert(data:byte(21) == 0x2f, "Invalid VP8L signature byte")
    local br = newBitReader(data:sub(22))

    local width  = br:read(14) + 1
    local height = br:read(14) + 1
    br:readBool()       -- alpha hint flag, unused
    local version = br:read(3)
    assert(version == 0, "Unsupported VP8L version: " .. version)

    local pixels = decodeVP8L(br, width, height)

    -- VP8L stores ARGB; Love2D mapPixel expects normalised RGBA floats
    local imageData = love.image.newImageData(width, height, "rgba8")
    imageData:mapPixel(function(x, y)
        local c = pixels[y * width + x + 1] or 0
        local a = band(rshift(c, 24), 0xff)
        local r = band(rshift(c, 16), 0xff)
        local g = band(rshift(c,  8), 0xff)
        local b = band(c, 0xff)
        return r / 255, g / 255, b / 255, a / 255
    end)

    return love.graphics.newImage(imageData)
end

function WebP.load(path)
    local data = love.filesystem.read(path)
    assert(data, "Could not read file: " .. path)
    return WebP.decode(data)
end

return WebP