-- work in pewgross
-- do not share !!

---@class UnitRGBA: Vector4 [r, g, b, a]; all values are 0-1.
local UnitRGBA = {};

---@param vector Vector4 a Vector4 matching the format of UnitRGBA ([r, g, b, a]; all values are 0-1)
---@return UnitRGBA? output nil the vector with the new metatable if valid, otherwise nil
---@return string? error a string describing the error if invalid, otherwise nil
---@nodiscard
---@see UnitRGBA.fromUnitStrict to throw an error instead of returning nil
function UnitRGBA.fromUnit(vector)
    for i, v in ipairs({vector:unpack()}) do
        if type(v) ~= "number" then
            return nil, "Vector component " .. i .. " is not a number";
        end
        if v < 0 or v > 1 then
            return nil, "Vector component " .. i .. " is not in the range 0-1; got " .. v;
        end
    end
    return setmetatable({ vector = vector }, UnitRGBA), nil;
end
---@param vector Vector4 a Vector4 matching the format of UnitRGBA ([r, g, b, a]; all values are 0-1)
---@return UnitRGBA vector the vector with the new metatable if valid. otherwise, an error is thrown
---@nodiscard
---@see UnitRGBA.fromUnit to return nil instead of throwing an error
function UnitRGBA.fromUnitStrict(vector)
    local result, why = UnitRGBA.fromUnit(vector);
    if result == nil then
        error("Invalid vector: " .. why);
    end
    return result;
end
---@param vector Vector4 a Vector4 matching the format of UnitRGBA ([r, g, b, a]; all values are 0-255)
---@return UnitRGBA? output nil the vector with the new metatable if valid, otherwise nil
---@return string? error a string describing the error if invalid, otherwise nil
---@nodiscard
---@see UnitRGBA.fromU8Strict to throw an error instead of returning nil
function UnitRGBA.fromU8(vector)
    for i, v in ipairs(vector) do
        if type(v) ~= "number" then
            return nil, "Vector component " .. i .. " is not a number";
        end
        if v < 0 or v > 255 then
            return nil, "Vector component " .. i .. " is not in the range 0-255";
        end
    end
    return setmetatable({ vector = vector / 255 }, UnitRGBA), nil;
end
---@param vector Vector4 a Vector4 matching the format of UnitRGBA ([r, g, b, a]; all values are 0-255)
---@return UnitRGBA vector the vector with the new metatable if valid. otherwise, an error is thrown
---@nodiscard
---@see UnitRGBA.fromU8 to return nil instead of throwing an error
function UnitRGBA.fromU8Strict(vector)
    local result, why = UnitRGBA.fromU8(vector);
    if result == nil then
        error("Invalid vector: " .. why);
    end
    return result;
end

---@param vector Vector4 a Vector4 matching the format of UnitRGBA ([r, g, b, a]; all values are 0-1)
---@return UnitRGBA vector the vector with the new metatable
---@nodiscard
function UnitRGBA.fromUnitClamping(vector)
    vector[1] = math.clamp(vector[1], 0, 1);
    vector[2] = math.clamp(vector[2], 0, 1);
    vector[3] = math.clamp(vector[3], 0, 1);
    vector[4] = math.clamp(vector[4], 0, 1);
    return setmetatable({ vector = vector }, UnitRGBA);
end

---@param hex string a color in hexadecimal format (e.g. "#FF00FF")
---@return UnitRGBA? output the color with the new metatable if valid, otherwise nil
---@return string? error a string describing the error if invalid, otherwise nil
---@nodiscard
---@see UnitRGBA.fromHexStrict to throw an error instead of returning nil
function UnitRGBA.fromHex(hex)
    if type(hex) ~= "string" then
        return nil, "not a string";
    end
    if hex:sub(1, 1) == "#" then
        hex = hex:sub(2);
    end
    if hex:len() ~= 6 and hex:len() ~= 8 then
        return nil, "not 6 or 8 characters long";
    end
    local vector = vec(0, 0, 0, 1);
    for i = 1, 3 do
        local value = tonumber(hex:sub(i * 2 - 1, i * 2), 16);
        if value == nil then
            return nil, "contains non-hex characters";
        end
        vector[i] = value / 255;
    end
    if hex:len() == 8 then
        local value = tonumber(hex:sub(7, 8), 16);
        if value == nil then
            return nil, "contains non-hex characters";
        end
        vector[4] = value / 255;
    else
        vector[4] = 1;
    end
    return setmetatable({ vector = vector }, UnitRGBA), nil;
end
---@param hex string a color in hexadecimal format (e.g. "#FF00FF")
---@return UnitRGBA color the color with the new metatable if valid. otherwise, an error is thrown
---@nodiscard
---@see UnitRGBA.fromHex to return nil instead of throwing an error
function UnitRGBA.fromHexStrict(hex)
    local result, why = UnitRGBA.fromHex(hex);
    if result == nil then
        error("Invalid hex color: " .. why);
    end
    return result;
end

---@return boolean opaque whether the color is *fully opaque* (as in, a == 1, not just a > 0).
---@see UnitRGBA:isTranslucent to see if the color is partially or fully transparent.
---@see UnitRGBA:isTransparent to see if the color is fully transparent.
---@nodiscard
function UnitRGBA:isOpaque()
    return self.vector[4] == 1;
end
---@return boolean transparent whether the color is *fully transparent* (as in, a == 0, not just a < 1).
---@see UnitRGBA:isTranslucent to see if the color is partially transparent.
---@see UnitRGBA:isOpaque to see if the color is fully opaque.
---@nodiscard
function UnitRGBA:isTransparent()
    return self.vector[4] == 0;
end
---@return boolean translucent whether the color is *translucent* (as in, a < 1).
---@see UnitRGBA:isTransparent to see if the color is fully transparent.
---@see UnitRGBA:isOpaque to see if the color is fully opaque.
---@nodiscard
function UnitRGBA:isTranslucent()
    return self.vector[4] < 1;
end

function UnitRGBA:toVec3()
    return vec(self.vector[1], self.vector[2], self.vector[3]);
end


---@return Vector4 vector Vector4 with format ([r, g, b, a]; all values are 0-255)
---@nodiscard
function UnitRGBA:toU8vec4()
    return vec(self.vector[1] * 255, self.vector[2] * 255, self.vector[3] * 255, self.vector[4] * 255);
end
---@return Vector3 vector Vector3 with format ([r, g, b]; all values are 0-255)
---@nodiscard
function UnitRGBA:toU8vec3()
    local a = vec(self.vector[1] * 255, self.vector[2] * 255, self.vector[3] * 255)
    return a
end


---@param hash? boolean? whether to include the hash symbol (#) in the output. defaults to true
---@param alpha? boolean? whether to include the alpha channel in the output. defaults to true
---@return string hex the color in RGBA hexadecimal format (e.g. "#FF00FF")
---@nodiscard
function UnitRGBA:toHex(hash, alpha)
    local hex = (hash == nil or hash) and "#" or "";
    local iter = (alpha == nil or alpha) and 4 or 3;
    for i = 1, iter do
        local value = math.round(self.vector[i] * 255);
        hex = hex .. string.format("%02X", value);
    end
    return hex;
end

function UnitRGBA:__index(key)
    return rawget(UnitRGBA, key) or self.vector[key];
end
function UnitRGBA:__tostring()
    return "UnitRGBA(" .. table.concat(self.vector, ", ") .. ")";
end

function UnitRGBA:__sub(other)
    return UnitRGBA.fromUnitClamping(vec(
        self.vector[1] - other.vector[1],
        self.vector[2] - other.vector[2],
        self.vector[3] - other.vector[3],
        self.vector[4] - other.vector[4]
    ));
end

--- Transparent white.
UnitRGBA.TRANSPARENT_WHITE = UnitRGBA.fromUnitStrict(vec(1, 1, 1, 0));
--- Transparent black.
UnitRGBA.TRANSPARENT_BLACK = UnitRGBA.fromUnitStrict(vec(0, 0, 0, 0));
--- Opaque white.
UnitRGBA.WHITE = UnitRGBA.fromUnitStrict(vec(1, 1, 1, 1));
--- Opaque black.
UnitRGBA.BLACK = UnitRGBA.fromUnitStrict(vec(0, 0, 0, 1));

---@class Keypoint - A keypoint (interpolation point definition) in a gradient.
---@field public color UnitRGBA color of the keypoint
---@field public time number a number between 0 and 1 (inclusive on both ends)
local Keypoint = {};

---Returns a new Keypoint object.
---@param color UnitRGBA color of the keypoint
---@param time number a number between 0 and 1 (inclusive on both ends)
---@return Keypoint
---@nodiscard
function Keypoint.new(color, time)
    return setmetatable({
        color = color,
        time = time
    }, Keypoint);
end

---@class Gradient A gradient definition.
---@field public keypoints Keypoint[] a list of keypoint objects
---@field public loop boolean whether the gradient should loop when given out-of-bounds time values. if false, an error is instead thrown.
local Gradient = {};

---Returns a new Gradient object.
---@param keypoints Keypoint[] a list of keypoint objects
---@param loop boolean? whether the gradient should loop when given out-of-bounds time values. if false, an error is instead thrown. defaults to true.
---@return Gradient
---@nodiscard
function Gradient.new(keypoints, loop)
    return setmetatable({
        keypoints = keypoints,
        loop = loop == nil or loop
    }, Gradient);
end

local function easeInOutQuad(t)
    return t < 0.5 and 2 * t * t or 1 - (-2 * t + 2)^2 / 2
end

local oklab = require("scripts.libs.oklab")
---Returns the color of the gradient at a given time.
---@param time number a number between 0 and 1 (inclusive on both ends) (can extend beyond this range if Gradient.loop is true)
---@return UnitRGBA? color the color of the gradient at the given time, or nil if the time is out of bounds and Gradient.loop is false
function Gradient:at(time)
    if time < 0 or time > 1 then
        if not self.loop then
            return nil;
        end
        time = time % 1;
    end

        if time == 0 then return self.keypoints[1].color
    elseif time == 1 then return self.keypoints[#self.keypoints].color
    end

    for i = 1, #self.keypoints - 1 do
        local this = self.keypoints[i]
        local next = self.keypoints[i + 1]
        if time >= this.time and time < next.time then
            local t = easeInOutQuad((time - this.time) / (next.time - this.time));

            local alpha = math.lerp(this.color.a, next.color.a, t)
            local La, Aa, Ba = oklab.srgbToLinear(this.color.xyz):unpack()
            local Lb, Ab, Bb = oklab.srgbToLinear(next.color.xyz):unpack()
            local mixed = vec(
                math.lerp(La, Lb, easeInOutQuad(t)),
                math.lerp(Aa, Ab, t),
                math.lerp(Ba, Bb, t)
            )

            mixed = oklab.linearToSrgb(mixed);
            mixed = vec(mixed.x, mixed.y, mixed.z, alpha)
            return UnitRGBA.fromUnitClamping(mixed);
        end
    end

    if #self.keypoints == 1 then
        return self.keypoints[1].color;
    end

    error("Gradient.at: time " .. time .. " is out of bounds");
end

---@param colors string[] a list of colors in hexadecimal format (e.g. "#FF00FF")
---@param loop boolean? whether the gradient should loop when given out-of-bounds time values. if false, an error is instead thrown. defaults to true.
---@return Gradient gradient the gradient with the new metatable if valid. otherwise, an error is thrown
function Gradient.distributedHex(colors, loop)
    local keypoints = {};
    for i, hex in ipairs(colors) do
        local time = (i - 1) / (#colors - 1);
        local color = UnitRGBA.fromHexStrict(hex);
        table.insert(keypoints, Keypoint.new(color, time));
    end
    return Gradient.new(keypoints, loop);
end

function Gradient:__index(key)
    return rawget(Gradient, key) or self[key];
end


---@class SimpleGradientBuilder
---@field public colors UnitRGBA[] a list of colors
local SimpleGradientBuilder = {};
---@return SimpleGradientBuilder builder a simple gradient builder
function SimpleGradientBuilder.new()
    return setmetatable({ colors = {} }, { __index = SimpleGradientBuilder });
end
---Reflect the colors of the gradient; this makes it nicer when looping.
---@param reflect_last boolean? whether to also reflect the final color, such that it would appear twice in a row at the middle. if false, it'll only be present once
---@return SimpleGradientBuilder self builder for chaining
function SimpleGradientBuilder:reflect(reflect_last)
    ---@type UnitRGBA[]
    local reflected = {};
    for i = 1, #self.colors - 1 do
        table.insert(reflected, self.colors[i]);
    end
    table.insert(reflected, self.colors[#self.colors])
    if reflect_last then
        table.insert(reflected, self.colors[#self.colors])
    end
    for i = #self.colors - 1, 1, -1 do
        table.insert(reflected, self.colors[i])
    end
    self.colors = reflected
    return self;
end
---"Thicken" the gradient by duplicating every keypoint. This means that the actual colors themselves will have as much presence as the transitions.
---@param amount? integer how many times to duplicate each keypoint. defaults to 1.
---@return SimpleGradientBuilder self builder for chaining
function SimpleGradientBuilder:thicken(amount)
    amount = amount or 1;
    local thickened = {};
    for i = 1, #self.colors do
        for _ = 0, amount do
            table.insert(thickened, self.colors[i])
        end
    end
    self.colors = thickened;
    return self;
end
---@param loop boolean? whether the gradient should loop when given out-of-bounds time values. if false, an error is instead thrown. defaults to true.
---@return Gradient gradient the constructed gradient
function SimpleGradientBuilder:build(loop)
    loop = loop == nil or loop;

    ---@type Keypoint[]
    local keypoints = {};
    for i, color in ipairs(self.colors) do
        table.insert(keypoints, Keypoint.new(color, (i - 1) / (#self.colors - 1)));
    end

    return setmetatable({
        keypoints = keypoints,
        loop = loop
    }, Gradient);
end
---@param colors string[] | string hexadecimal color(s) to add
function SimpleGradientBuilder:add(colors)
    if type(colors) == "string" then
        colors = {colors}
    end
    for _, color in ipairs(colors) do
        table.insert(self.colors, UnitRGBA.fromHexStrict(color))
    end
    return self;
end

return {
    UnitRGBA = UnitRGBA,
    Keypoint = Keypoint,
    Gradient = Gradient,
    SimpleGradientBuilder = SimpleGradientBuilder,
}