/*
    KWin - the KDE window manager
    This file is part of the KDE project.

    SPDX-FileCopyrightText: 2024 Xaver Hugl <xaver.hugl@gmail.com>

    SPDX-License-Identifier: GPL-2.0-or-later
*/
#include "drm_colorop.h"
#include "drm_blob.h"
#include "drm_commit.h"
#include "drm_gpu.h"
#include "drm_object.h"
#include "utils/envvar.h"

#include <ranges>

namespace KWin
{

DrmAbstractColorOp::DrmAbstractColorOp(DrmAbstractColorOp *next, Features features, const QString &name)
    : m_next(next)
    , m_features(features)
    , m_name(name)
{
}

DrmAbstractColorOp::~DrmAbstractColorOp()
{
}

DrmAbstractColorOp *DrmAbstractColorOp::next() const
{
    return m_next;
}

static const auto s_disableAmdgpuWorkaround = environmentVariableBoolValue("KWIN_DRM_DISABLE_AMD_GAMMA_WORKAROUND");

bool DrmAbstractColorOp::canBypass() const
{
    return m_features & Feature::Bypass;
}

bool DrmAbstractColorOp::supportsMultipleOps() const
{
    return m_features & Feature::MultipleOps;
}

QString DrmAbstractColorOp::name() const
{
    return m_name;
}

struct Assignments
{
    std::optional<ColorOp> nextInverseInputScaling;
    DrmAbstractColorOp *nextHardwareOp = nullptr;
    std::unordered_map<DrmAbstractColorOp *, std::deque<ColorOp::Operation>> assignments;
};

static bool backpropagateNextInputScaling(DrmAbstractColorOp *hardwareOp, DrmAbstractColorOp *previousHardwareOp, Assignments &assignments)
{
    // This potentially propagates the inverse input scaling from the next hardware block we applied
    // a color op on.
    //
    // previousHardwareOp is the hardwareOp from the last recursion step, and is used to determine
    // whether or not the current hardwareOp will have color operations assigned to it in the last
    // recursion step, or if it was bypassed instead.
    //
    // If previousHardwareOp == hardwareOp, it means that other operations will be assigned to the
    // hardware block, and the inverse input scaling can only be the last operation, i.e.
    //
    //     -> |Current/Previous block: ... [Op] [Inverse Input Scaling]| -> ... -> |Next block: [Input Scaling] [Op] ...|
    //
    // If previousHardwareOp != hardwareOp, it means that hardwareOp has been bypassed and no color
    // pipeline operations have been assigned to it. We can still try to make a use of it for the inverse
    // input scaline though, i.e.
    //
    //    -> |Previous Block: ...| -> ... -> |Current Block: [Inverse Input Scaling]| -> ... -> |Next Block: [Input Scaling] ...|

    if (!assignments.nextInverseInputScaling) {
        return true;
    }
    if (!hardwareOp->colorOpPreference(assignments.nextInverseInputScaling->operation)) {
        return false;
    }
    const bool hasOtherColorOps = previousHardwareOp == hardwareOp;
    if (hasOtherColorOps && !hardwareOp->supportsMultipleOps()) {
        return false;
    }
    // If this is the only color op being applied with this hardwareOp,
    // we need to make sure it doesn't need input scaling itself.
    if (!hasOtherColorOps && hardwareOp->inputScaling(*assignments.nextInverseInputScaling).has_value()) {
        return false;
    }
    // This is the last color op applied on this hardwareOp, if it needs
    // output scaling, we can't use it
    if (hardwareOp->outputScaling(*assignments.nextInverseInputScaling).has_value()) {
        return false;
    }
    assignments.assignments[hardwareOp].push_front(assignments.nextInverseInputScaling->operation);
    assignments.nextInverseInputScaling.reset();
    return true;
}

// for an explanation of the algorithm, see "drm color operations.md"
static std::optional<Assignments> findColorPipelineAssignments(DrmAbstractColorOp *hardwareOp, DrmAbstractColorOp *previousHardwareOp, std::span<const ColorOp> ops, const std::optional<ColorOp> &previousInverseOutputScaling)
{
    const ColorOp *currentColorOp = nullptr;
    std::span<const ColorOp> nextOps;
    const bool isInput = hardwareOp != previousHardwareOp;
    if (isInput && previousInverseOutputScaling.has_value()) {
        // we need to apply the previous operation's output scaling, prioritize it
        currentColorOp = &*previousInverseOutputScaling;
        nextOps = ops;
    } else if (ops.empty()) {
        DrmAbstractColorOp *next = hardwareOp;
        while (next) {
            if (!next->canBypass()) {
                return std::nullopt;
            }
            next = next->next();
        }
        // we're done here, the whole pipeline is assigned!
        return Assignments{
            .nextInverseInputScaling = std::nullopt,
            .nextHardwareOp = hardwareOp,
            .assignments = {},
        };
    } else {
        currentColorOp = &ops.front();
        nextOps = ops.subspan(1);
    }
    if (!hardwareOp) {
        return std::nullopt;
    }

    auto preference = hardwareOp->colorOpPreference(currentColorOp->operation);
    if (isInput && !hardwareOp->canBeInputOrOutput(currentColorOp->inputSpace)) {
        preference.reset();
    }
    if (!preference) {
        if (!hardwareOp->canBypass()) {
            return std::nullopt;
        }
        auto ret = findColorPipelineAssignments(hardwareOp->next(), hardwareOp, ops, previousInverseOutputScaling);
        if (ret) {
            backpropagateNextInputScaling(hardwareOp, previousHardwareOp, *ret);
        }
        return ret;
    }
    if (hardwareOp->canBypass() && *preference != DrmAbstractColorOp::Priority::High) {
        // attempt to find a higher priority hardware op type
        DrmAbstractColorOp *higherPreferenceOp = hardwareOp->next();
        while (higherPreferenceOp) {
            const auto higherPreference = higherPreferenceOp->colorOpPreference(currentColorOp->operation);
            if (higherPreference.has_value() && *higherPreference > *preference) {
                if (auto ret = findColorPipelineAssignments(higherPreferenceOp, hardwareOp, ops, previousInverseOutputScaling)) {
                    backpropagateNextInputScaling(hardwareOp, previousHardwareOp, *ret);
                    return ret;
                } else {
                    break;
                }
            }
            if (!higherPreferenceOp->canBypass()) {
                break;
            }
            higherPreferenceOp = higherPreferenceOp->next();
        }
    }

    // this hardware op is the best match for currentColorOp
    // before doing anything with it though, figure out if the other operations can be applied as well
    const auto outputScaling = hardwareOp->outputScaling(*currentColorOp);
    auto nextHardwareOp = hardwareOp->supportsMultipleOps() ? hardwareOp : hardwareOp->next();
    auto ret = findColorPipelineAssignments(nextHardwareOp, hardwareOp, nextOps, outputScaling ? std::optional(outputScaling->inverse) : std::nullopt);
    if (!ret) {
        if (!nextHardwareOp) {
            return std::nullopt;
        }
        // in some cases, skipping hardware ops can find working configurations
        // as some hardware ops can only operate on specific color spaces
        if (auto ret = findColorPipelineAssignments(nextHardwareOp->next(), nextHardwareOp, ops, previousInverseOutputScaling)) {
            backpropagateNextInputScaling(hardwareOp, previousHardwareOp, *ret);
            return ret;
        } else {
            return std::nullopt;
        }
    }

    const bool isOutput = ret->nextHardwareOp != hardwareOp;
    if (isOutput) {
        auto lastColorSpace = currentColorOp->outputSpace;
        if (ret->nextInverseInputScaling) {
            lastColorSpace = ret->nextInverseInputScaling->outputSpace;
        } else if (outputScaling) {
            lastColorSpace = outputScaling->inverse.inputSpace;
        }
        if (!hardwareOp->canBeInputOrOutput(lastColorSpace)) {
            return std::nullopt;
        }
    }

    // first, the pre-op scaling from the next hardware op
    // NOTE that the "previous" hardware op is hardwareOp in this case
    // as we're pushing currentColorOp before the next input scaling
    if (!backpropagateNextInputScaling(hardwareOp, hardwareOp, *ret)) {
        return std::nullopt;
    }

    // next, our own output scaling
    if (outputScaling && isOutput) {
        // need to prepend, because we recursively walk back
        // from the last color opt to the front
        ret->assignments[hardwareOp].push_front(outputScaling->scaling);
    }

    // next, the actual operation
    ret->assignments[hardwareOp].push_front(currentColorOp->operation);

    // last, our own pre-op scaling
    if (isInput) {
        const auto inputScaling = hardwareOp->inputScaling(*currentColorOp);
        if (inputScaling) {
            ret->assignments[hardwareOp].push_front(inputScaling->scaling);
            ret->nextInverseInputScaling = inputScaling->inverse;
        }
    }
    ret->nextHardwareOp = hardwareOp;
    return ret;
}

bool DrmAbstractColorOp::matchPipeline(DrmAtomicCommit *commit, const ColorPipeline &pipeline)
{
    if (pipeline.isIdentity() && s_disableAmdgpuWorkaround.value_or(!commit->gpu()->drmDevice()->isAmdgpu())) {
        // Applying this config is very simple and cheap, so do it directly
        // and avoid invalidating the cache
        DrmAbstractColorOp *currentOp = this;
        while (currentOp) {
            currentOp->bypass(commit);
            currentOp = currentOp->next();
        }
        return true;
    }
    if (m_cachedPipeline && *m_cachedPipeline == pipeline) {
        commit->merge(m_cache.get());
        return true;
    }
    if (!matchPipeline(commit->gpu(), pipeline)) {
        return false;
    }
    commit->merge(m_cache.get());
    return true;
}

bool DrmAbstractColorOp::matchPipeline(DrmGpu *gpu, const ColorPipeline &pipeline)
{
    if (pipeline.isIdentity() && s_disableAmdgpuWorkaround.value_or(!gpu->drmDevice()->isAmdgpu())) {
        return true;
    }
    if (m_cachedPipeline && *m_cachedPipeline == pipeline) {
        return true;
    }
    if (m_cachedPipelineFail && *m_cachedPipelineFail == pipeline) {
        return false;
    }

    // first, search for the most optimal pipeline, without calculating LUTs
    const auto ret = findColorPipelineAssignments(this, nullptr, pipeline.ops, std::nullopt);
    if (!ret || ret->nextInverseInputScaling) {
        m_cachedPipelineFail = pipeline;
        return false;
    }

    // now actually program the properties
    m_cache = std::make_unique<DrmAtomicCommit>(gpu);
    DrmAbstractColorOp *currentOp = this;
    while (currentOp) {
        const auto it = ret->assignments.find(currentOp);
        if (it != ret->assignments.end()) {
            const auto &[op, program] = *it;
            currentOp->program(m_cache.get(), program);
        } else {
            currentOp->bypass(m_cache.get());
        }
        currentOp = currentOp->next();
    }
    m_cachedPipeline = pipeline;
    return true;
}

std::optional<DrmAbstractColorOp::Scaling> DrmAbstractColorOp::inputScaling(const ColorOp &op) const
{
    return std::nullopt;
}

std::optional<DrmAbstractColorOp::Scaling> DrmAbstractColorOp::outputScaling(const ColorOp &op) const
{
    return std::nullopt;
}

bool DrmAbstractColorOp::canBeInputOrOutput(ColorspaceType type) const
{
    return true;
}

DrmLutColorOp16::DrmLutColorOp16(DrmAbstractColorOp *next, DrmProperty *prop, DrmEnumProperty<Lut1DInterpolation> *interpolation, uint32_t maxSize, DrmProperty *bypass)
    : DrmAbstractColorOp(next, Features{Feature::MultipleOps} | Feature::Bypass, QStringLiteral("1D LUT"))
    , m_prop(prop)
    , m_bypass(bypass)
    , m_interpolationMode(interpolation)
    , m_maxSize(maxSize)
    , m_components(m_maxSize)
{
}

std::optional<DrmAbstractColorOp::Priority> DrmLutColorOp16::colorOpPreference(const ColorOp::Operation &op)
{
    if (std::holds_alternative<ColorTransferFunction>(op) || std::holds_alternative<InverseColorTransferFunction>(op)
        || std::holds_alternative<ColorTonemapper>(op) || std::holds_alternative<std::shared_ptr<ColorTransformation>>(op)) {
        // the required resolution depends heavily on the function and on the input and output ranges / multipliers
        // but this is good enough for now
        if (m_maxSize >= 1024) {
            return Priority::Normal;
        } else if (m_maxSize >= 256) {
            return Priority::Low;
        }
        // TODO also check matrices for scaling + offset only, which can be represented just fine!
    } else if (std::holds_alternative<ColorMultiplier>(op)) {
        return Priority::Low;
    } else if (std::holds_alternative<ColorClamp>(op)) {
        return Priority::High;
    }
    return std::nullopt;
}

void DrmLutColorOp16::program(DrmAtomicCommit *commit, const std::deque<ColorOp::Operation> &operations)
{
    for (uint32_t i = 0; i < m_maxSize; i++) {
        const double input = i / double(m_maxSize - 1);
        QVector3D output(input, input, input);
        for (const auto &op : operations) {
            output = ColorOp::applyOperation(op, output);
        }
        m_components[i] = {
            .red = uint16_t(std::round(std::clamp(output.x(), 0.0f, 1.0f) * std::numeric_limits<uint16_t>::max())),
            .green = uint16_t(std::round(std::clamp(output.y(), 0.0f, 1.0f) * std::numeric_limits<uint16_t>::max())),
            .blue = uint16_t(std::round(std::clamp(output.z(), 0.0f, 1.0f) * std::numeric_limits<uint16_t>::max())),
            .reserved = 0,
        };
    }
    commit->addBlob(*m_prop, DrmBlob::create(m_prop->drmObject()->gpu(), m_components.data(), sizeof(drm_color_lut) * m_maxSize));
    if (m_bypass) {
        commit->addProperty(*m_bypass, 0);
    }
}

void DrmLutColorOp16::bypass(DrmAtomicCommit *commit)
{
    commit->addBlob(*m_prop, nullptr);
    if (m_bypass) {
        commit->addProperty(*m_bypass, 1);
    }
}

std::optional<DrmAbstractColorOp::Scaling> DrmLutColorOp16::inputScaling(const ColorOp &op) const
{
    if (op.input.min < -0.001) {
        // needs an offset to be possible
        QMatrix4x4 mat;
        mat.translate(op.input.min, op.input.min, op.input.min);
        mat.scale(op.input.max - op.input.min);
        return Scaling{
            .scaling = ColorMatrix(mat),
            .inverse = ColorOp{
                .input = op.input,
                .inputSpace = op.inputSpace,
                .operation = ColorMatrix(mat.inverted()),
                .output = ValueRange{
                    .min = 0,
                    .max = 1,
                },
                .outputSpace = op.inputSpace,
            },
        };
    } else if (op.input.max <= 1.001) {
        return std::nullopt;
    }
    return Scaling{
        .scaling = ColorMultiplier(op.input.max),
        .inverse = ColorOp{
            .input = op.input,
            .inputSpace = op.inputSpace,
            .operation = ColorMultiplier(1.0 / op.input.max),
            .output = ValueRange{
                .min = op.input.min / op.input.max,
                .max = 1.0,
            },
            .outputSpace = op.inputSpace,
        },
    };
}

std::optional<DrmAbstractColorOp::Scaling> DrmLutColorOp16::outputScaling(const ColorOp &op) const
{
    if (op.output.min < -0.001) {
        // needs an offset to be possible
        QMatrix4x4 mat;
        mat.scale(1.0 / (op.output.max - op.output.min));
        mat.translate(-op.output.min, -op.output.min, -op.output.min);
        return Scaling{
            .scaling = ColorMatrix(mat),
            .inverse = ColorOp{
                .input = ValueRange{
                    .min = 0,
                    .max = 1,
                },
                .inputSpace = op.outputSpace,
                .operation = ColorMatrix(mat.inverted()),
                .output = op.output,
                .outputSpace = op.outputSpace,
            },
        };
    } else if (op.output.max <= 1.001) {
        return std::nullopt;
    }
    return Scaling{
        .scaling = ColorMultiplier(1.0 / op.output.max),
        .inverse = ColorOp{
            .input = ValueRange{
                .min = op.output.min / op.output.max,
                .max = 1.0,
            },
            .inputSpace = op.outputSpace,
            .operation = ColorMultiplier(op.output.max),
            .output = op.output,
            .outputSpace = op.outputSpace,
        },
    };
}

DrmLutColorOp32::DrmLutColorOp32(DrmAbstractColorOp *next, DrmProperty *prop, DrmEnumProperty<Lut1DInterpolation> *interpolation, uint32_t maxSize, DrmProperty *bypass)
    : DrmAbstractColorOp(next, Features{Feature::MultipleOps} | Feature::Bypass, QStringLiteral("1D LUT"))
    , m_prop(prop)
    , m_bypass(bypass)
    , m_interpolationMode(interpolation)
    , m_maxSize(maxSize)
    , m_components(m_maxSize)
{
}

std::optional<DrmAbstractColorOp::Priority> DrmLutColorOp32::colorOpPreference(const ColorOp::Operation &op)
{
    if (std::holds_alternative<ColorTransferFunction>(op) || std::holds_alternative<InverseColorTransferFunction>(op)
        || std::holds_alternative<ColorTonemapper>(op) || std::holds_alternative<std::shared_ptr<ColorTransformation>>(op)) {
        // the required resolution depends heavily on the function and on the input and output ranges / multipliers
        // but this is good enough for now
        if (m_maxSize >= 1024) {
            return Priority::Normal;
        } else if (m_maxSize >= 256) {
            return Priority::Low;
        }
        // TODO also check matrices for scaling + offset only, which can be represented just fine!
    } else if (std::holds_alternative<ColorMultiplier>(op)) {
        return Priority::Low;
    } else if (std::holds_alternative<ColorClamp>(op)) {
        return Priority::High;
    }
    return std::nullopt;
}

void DrmLutColorOp32::program(DrmAtomicCommit *commit, const std::deque<ColorOp::Operation> &operations)
{
    for (uint32_t i = 0; i < m_maxSize; i++) {
        const double input = i / double(m_maxSize - 1);
        QVector3D output(input, input, input);
        for (const auto &op : operations) {
            output = ColorOp::applyOperation(op, output);
        }
        m_components[i] = {
            .red = uint32_t(std::round(std::clamp<double>(output.x(), 0.0, 1.0) * std::numeric_limits<uint32_t>::max())),
            .green = uint32_t(std::round(std::clamp<double>(output.y(), 0.0, 1.0) * std::numeric_limits<uint32_t>::max())),
            .blue = uint32_t(std::round(std::clamp<double>(output.z(), 0.0, 1.0) * std::numeric_limits<uint32_t>::max())),
            .reserved = 0,
        };
    }
    commit->addBlob(*m_prop, DrmBlob::create(m_prop->drmObject()->gpu(), m_components.data(), sizeof(LutComponent32) * m_maxSize));
    if (m_bypass) {
        commit->addProperty(*m_bypass, 0);
    }
}

void DrmLutColorOp32::bypass(DrmAtomicCommit *commit)
{
    commit->addBlob(*m_prop, nullptr);
    if (m_bypass) {
        commit->addProperty(*m_bypass, 1);
    }
}

std::optional<DrmAbstractColorOp::Scaling> DrmLutColorOp32::inputScaling(const ColorOp &op) const
{
    if (op.input.min < -0.001) {
        // needs an offset to be possible
        QMatrix4x4 mat;
        mat.translate(op.input.min, op.input.min, op.input.min);
        mat.scale(op.input.max - op.input.min);
        return Scaling{
            .scaling = ColorMatrix(mat),
            .inverse = ColorOp{
                .input = op.input,
                .inputSpace = op.inputSpace,
                .operation = ColorMatrix(mat.inverted()),
                .output = ValueRange{
                    .min = 0,
                    .max = 1,
                },
                .outputSpace = op.inputSpace,
            },
        };
    } else if (op.input.max <= 1.001) {
        return std::nullopt;
    }
    return Scaling{
        .scaling = ColorMultiplier(op.input.max),
        .inverse = ColorOp{
            .input = op.input,
            .inputSpace = op.inputSpace,
            .operation = ColorMultiplier(1.0 / op.input.max),
            .output = ValueRange{
                .min = op.input.min / op.input.max,
                .max = 1.0,
            },
            .outputSpace = op.inputSpace,
        },
    };
}

std::optional<DrmAbstractColorOp::Scaling> DrmLutColorOp32::outputScaling(const ColorOp &op) const
{
    if (op.output.min < -0.001) {
        // needs an offset to be possible
        QMatrix4x4 mat;
        mat.scale(1.0 / (op.output.max - op.output.min));
        mat.translate(-op.output.min, -op.output.min, -op.output.min);
        return Scaling{
            .scaling = ColorMatrix(mat),
            .inverse = ColorOp{
                .input = ValueRange{
                    .min = 0,
                    .max = 1,
                },
                .inputSpace = op.outputSpace,
                .operation = ColorMatrix(mat.inverted()),
                .output = op.output,
                .outputSpace = op.outputSpace,
            },
        };
    } else if (op.output.max <= 1.001) {
        return std::nullopt;
    }
    return Scaling{
        .scaling = ColorMultiplier(1.0 / op.output.max),
        .inverse = ColorOp{
            .input = ValueRange{
                .min = op.output.min / op.output.max,
                .max = 1.0,
            },
            .inputSpace = op.outputSpace,
            .operation = ColorMultiplier(op.output.max),
            .output = op.output,
            .outputSpace = op.outputSpace,
        },
    };
}

LegacyMatrixColorOp::LegacyMatrixColorOp(DrmAbstractColorOp *next, DrmProperty *prop)
    : DrmAbstractColorOp(next, Features{Feature::Bypass} | Feature::MultipleOps, QStringLiteral("legacy matrix"))
    , m_prop(prop)
{
}

std::optional<DrmAbstractColorOp::Priority> LegacyMatrixColorOp::colorOpPreference(const ColorOp::Operation &op)
{
    // TODO check the resolution of the matrix too?
    // -> values above abs(5) should be re-scaled with a different (previous?) operation
    if (auto matrix = std::get_if<ColorMatrix>(&op)) {
        const bool canRepresent = std::abs(matrix->mat(3, 0) - 0) < ColorPipeline::s_maxResolution
            && std::abs(matrix->mat(3, 1) - 0) < ColorPipeline::s_maxResolution
            && std::abs(matrix->mat(3, 2) - 0) < ColorPipeline::s_maxResolution
            && std::abs(matrix->mat(3, 3) - 1) < ColorPipeline::s_maxResolution
            && std::abs(matrix->mat(0, 3) - 0) < ColorPipeline::s_maxResolution
            && std::abs(matrix->mat(1, 3) - 0) < ColorPipeline::s_maxResolution
            && std::abs(matrix->mat(2, 3) - 0) < ColorPipeline::s_maxResolution;
        if (!canRepresent) {
            return std::nullopt;
        } else {
            return Priority::Normal;
        }
    } else if (std::holds_alternative<ColorMultiplier>(op)) {
        return Priority::Low;
    }
    return std::nullopt;
}

static uint64_t doubleToFixed(double value)
{
    // ctm values are in S31.32 sign-magnitude format
    uint64_t ret = std::abs(value) * (1ull << 32);
    if (value < 0) {
        ret |= 1ull << 63;
    }
    return ret;
}

void LegacyMatrixColorOp::program(DrmAtomicCommit *commit, const std::deque<ColorOp::Operation> &operations)
{
    // NOTE that matrix operations have to be added in reverse order to get the correct result!
    QMatrix4x4 result;
    for (const auto &op : operations | std::views::reverse) {
        if (auto matrix = std::get_if<ColorMatrix>(&op)) {
            result *= matrix->mat;
        } else if (auto mult = std::get_if<ColorMultiplier>(&op)) {
            result.scale(mult->factors.x(), mult->factors.y(), mult->factors.z());
        } else {
            Q_UNREACHABLE();
        }
    }
    drm_color_ctm data = {
        .matrix = {
            doubleToFixed(result(0, 0)), doubleToFixed(result(0, 1)), doubleToFixed(result(0, 2)), //
            doubleToFixed(result(1, 0)), doubleToFixed(result(1, 1)), doubleToFixed(result(1, 2)), //
            doubleToFixed(result(2, 0)), doubleToFixed(result(2, 1)), doubleToFixed(result(2, 2)), //
        },
    };
    commit->addBlob(*m_prop, DrmBlob::create(m_prop->drmObject()->gpu(), &data, sizeof(data)));
}

void LegacyMatrixColorOp::bypass(DrmAtomicCommit *commit)
{
    commit->addBlob(*m_prop, nullptr);
}

Matrix3x4ColorOp::Matrix3x4ColorOp(DrmAbstractColorOp *next, DrmProperty *prop, DrmProperty *bypass)
    : DrmAbstractColorOp(next, Features{Feature::MultipleOps} | (bypass ? Feature::Bypass : Features{}), QStringLiteral("3x4 matrix"))
    , m_prop(prop)
    , m_bypass(bypass)
{
}

std::optional<DrmAbstractColorOp::Priority> Matrix3x4ColorOp::colorOpPreference(const ColorOp::Operation &op)
{
    // TODO check the resolution of the matrix too?
    // -> values above abs(5) should be re-scaled with a different (previous?) operation
    if (auto matrix = std::get_if<ColorMatrix>(&op)) {
        const bool canRepresent = matrix->mat(3, 0) == 0
            && matrix->mat(3, 1) == 0
            && matrix->mat(3, 2) == 0
            && matrix->mat(3, 3) == 1;
        if (!canRepresent) {
            return std::nullopt;
        } else {
            return Priority::Normal;
        }
    } else if (std::holds_alternative<ColorMultiplier>(op)) {
        return Priority::Low;
    }
    return std::nullopt;
}

void Matrix3x4ColorOp::program(DrmAtomicCommit *commit, const std::deque<ColorOp::Operation> &operations)
{
    // NOTE that matrix operations have to be added in reverse order to get the correct result!
    QMatrix4x4 result;
    for (const auto &op : operations | std::views::reverse) {
        if (auto matrix = std::get_if<ColorMatrix>(&op)) {
            result *= matrix->mat;
        } else if (auto mult = std::get_if<ColorMultiplier>(&op)) {
            result.scale(mult->factors.x(), mult->factors.y(), mult->factors.z());
        } else {
            Q_UNREACHABLE();
        }
    }
    std::array<uint64_t, 12> data = {
        doubleToFixed(result(0, 0)), doubleToFixed(result(0, 1)), doubleToFixed(result(0, 2)), doubleToFixed(result(0, 3)), //
        doubleToFixed(result(1, 0)), doubleToFixed(result(1, 1)), doubleToFixed(result(1, 2)), doubleToFixed(result(1, 3)), //
        doubleToFixed(result(2, 0)), doubleToFixed(result(2, 1)), doubleToFixed(result(2, 2)), doubleToFixed(result(2, 3)), //
    };
    commit->addBlob(*m_prop, DrmBlob::create(m_prop->drmObject()->gpu(), data.data(), sizeof(uint64_t) * data.size()));
    if (m_bypass) {
        commit->addProperty(*m_bypass, 0);
    }
}

void Matrix3x4ColorOp::bypass(DrmAtomicCommit *commit)
{
    commit->addProperty(*m_bypass, 1);
}

UnknownColorOp::UnknownColorOp(DrmAbstractColorOp *next, DrmProperty *bypass)
    : DrmAbstractColorOp(next, bypass ? Feature::Bypass : Features{}, "unknown")
    , m_bypass(bypass)
{
}

std::optional<DrmAbstractColorOp::Priority> UnknownColorOp::colorOpPreference(const ColorOp::Operation &op)
{
    return std::nullopt;
}

void UnknownColorOp::program(DrmAtomicCommit *commit, const std::deque<ColorOp::Operation> &operations)
{
    Q_UNREACHABLE();
}

void UnknownColorOp::bypass(DrmAtomicCommit *commit)
{
    commit->addProperty(*m_bypass, 1);
}

DrmLut3DColorOp::DrmLut3DColorOp(DrmAbstractColorOp *next, DrmProperty *value, DrmProperty *bypass, size_t size, DrmEnumProperty<Lut3DInterpolation> *interpolation)
    : DrmAbstractColorOp(next, Features{Feature::MultipleOps} | (bypass ? Feature::Bypass : Features{}), QStringLiteral("3D LUT"))
    , m_value(value)
    , m_bypass(bypass)
    , m_size(size)
    , m_interpolation(interpolation)
    , m_components(m_size * m_size * m_size)
{
}

std::optional<DrmAbstractColorOp::Priority> DrmLut3DColorOp::colorOpPreference(const ColorOp::Operation &op)
{
    // everything can be represented as a 3D lut,
    // but resolution isn't great, so push it to other ops if possible
    return Priority::Low;
}

bool DrmLut3DColorOp::canBeInputOrOutput(ColorspaceType type) const
{
    // because resolution is so limited, input or output has to be non-linear RGB
    // linear RGB as I/O is also okay, because we can apply scaling to it
    return type == ColorspaceType::NonLinearRGB
        || type == ColorspaceType::LinearRGB;
}

void DrmLut3DColorOp::program(DrmAtomicCommit *commit, const std::deque<ColorOp::Operation> &operations)
{
    for (size_t r = 0; r < m_size; r++) {
        for (size_t g = 0; g < m_size; g++) {
            for (size_t b = 0; b < m_size; b++) {
                QVector3D output = QVector3D(r, g, b) / float(m_size - 1);
                for (const auto &op : operations) {
                    output = ColorOp::applyOperation(op, output);
                }
                const size_t index = b * m_size * m_size + g * m_size + r;
                m_components[index] = LutComponent32{
                    .red = uint32_t(std::round(std::clamp<double>(output.x(), 0.0, 1.0) * std::numeric_limits<uint32_t>::max())),
                    .green = uint32_t(std::round(std::clamp<double>(output.y(), 0.0, 1.0) * std::numeric_limits<uint32_t>::max())),
                    .blue = uint32_t(std::round(std::clamp<double>(output.z(), 0.0, 1.0) * std::numeric_limits<uint32_t>::max())),
                    .reserved = 0,
                };
            }
        }
    }
    commit->addBlob(*m_value, DrmBlob::create(m_value->drmObject()->gpu(), m_components.data(), m_components.size() * sizeof(LutComponent32)));
    if (m_bypass) {
        commit->addProperty(*m_bypass, 0);
    }
}

void DrmLut3DColorOp::bypass(DrmAtomicCommit *commit)
{
    commit->addProperty(*m_bypass, 1);
}

std::optional<DrmAbstractColorOp::Scaling> DrmLut3DColorOp::inputScaling(const ColorOp &op) const
{
    if (op.inputSpace == ColorspaceType::NonLinearRGB) {
        // The encoding is already some sort of non-linear RGB, applying PQ
        // on top of that would just make things worse rather than better
        // TODO use a custom LUT to use the limited resolution as much as is possible
        return std::nullopt;
    }
    TransferFunction tf{TransferFunction::PerceptualQuantizer, op.input.min, op.input.max};
    return Scaling{
        .scaling = ColorTransferFunction(tf),
        .inverse = ColorOp{
            .input = op.input,
            .inputSpace = op.inputSpace,
            .operation = InverseColorTransferFunction(tf),
            .output = ValueRange{
                .min = 0,
                .max = 1,
            },
            // NOTE this assumption only works because we enforce it in canBeInputOrOutput
            .outputSpace = ColorspaceType::NonLinearRGB,
        },
    };
}

std::optional<DrmAbstractColorOp::Scaling> DrmLut3DColorOp::outputScaling(const ColorOp &op) const
{
    if (op.outputSpace == ColorspaceType::NonLinearRGB) {
        // The encoding is already some sort of non-linear RGB, applying PQ
        // on top of that would just make things worse rather than better
        // TODO use a custom LUT to use the limited resolution as much as is possible
        return std::nullopt;
    }
    TransferFunction tf{TransferFunction::PerceptualQuantizer, op.output.min, op.output.max};
    return Scaling{
        .scaling = InverseColorTransferFunction(tf),
        .inverse = ColorOp{
            .input = ValueRange{
                .min = 0,
                .max = 1,
            },
            // NOTE this assumption only works because we enforce it in canBeInputOrOutput
            .inputSpace = ColorspaceType::NonLinearRGB,
            .operation = ColorTransferFunction(tf),
            .output = op.output,
            .outputSpace = op.outputSpace,
        },
    };
}

DrmMultiplier::DrmMultiplier(DrmAbstractColorOp *next, DrmProperty *value, DrmProperty *bypass)
    : DrmAbstractColorOp(next, bypass ? Feature::Bypass : Features{}, QStringLiteral("multiplier"))
    , m_value(value)
    , m_bypass(bypass)
{
}

static float commonScaling(const QMatrix4x4 &mat)
{
    return std::min({mat(0, 0), mat(1, 1), mat(2, 2)});
}

std::optional<DrmAbstractColorOp::Priority> DrmMultiplier::colorOpPreference(const ColorOp::Operation &op)
{
    if (const auto mult = std::get_if<ColorMultiplier>(&op)) {
        const float diff1 = std::abs(mult->factors.x() - mult->factors.y());
        const float diff2 = std::abs(mult->factors.y() - mult->factors.z());
        const float diff3 = std::abs(mult->factors.z() - mult->factors.x());
        const float maxDiff = std::max({diff1, diff2, diff3});
        if (maxDiff < 0.00001) {
            return Priority::High;
        }
    } else if (const auto mat = std::get_if<ColorMatrix>(&op)) {
        // we want to partially apply the matrix
        // as this has much higher resolution than matrices in hardware
        const float common = commonScaling(mat->mat);
        if (!qFuzzyIsNull(common) && !qFuzzyCompare(common, 1)) {
            return Priority::High;
        }
    }
    return std::nullopt;
}

void DrmMultiplier::program(DrmAtomicCommit *commit, const std::deque<ColorOp::Operation> &operations)
{
    double factor = 1;
    // only the first op matters, output scaling should be ignored
    const auto &op = operations.front();
    if (const auto mult = std::get_if<ColorMultiplier>(&op)) {
        factor = mult->factors.x();
    } else if (const auto mat = std::get_if<ColorMatrix>(&op)) {
        factor = commonScaling(mat->mat);
    }
    commit->addProperty(*m_value, doubleToFixed(factor));
    if (m_bypass) {
        commit->addProperty(*m_bypass, 0);
    }
}

void DrmMultiplier::bypass(DrmAtomicCommit *commit)
{
    commit->addProperty(*m_bypass, 1);
}

std::optional<DrmAbstractColorOp::Scaling> DrmMultiplier::inputScaling(const ColorOp &op) const
{
    return std::nullopt;
}

std::optional<DrmAbstractColorOp::Scaling> DrmMultiplier::outputScaling(const ColorOp &op) const
{
    if (const auto mat = std::get_if<ColorMatrix>(&op.operation)) {
        const float scaling = commonScaling(mat->mat);
        QMatrix4x4 remaining = mat->mat;
        remaining.scale(1.0 / scaling);
        if (!isFuzzyIdentity(remaining)) {
            return Scaling{
                .scaling = ColorMatrix(remaining.inverted()),
                .inverse = ColorOp{
                    .input = op.output * (1.0 / scaling),
                    .inputSpace = op.inputSpace,
                    .operation = ColorMatrix(remaining),
                    .output = op.output,
                    .outputSpace = op.outputSpace,
                },
            };
        }
    }
    return std::nullopt;
}

DrmNamed1DLut::DrmNamed1DLut(DrmAbstractColorOp *next, DrmEnumProperty<Named1DLutType> *value, DrmProperty *bypass)
    : DrmAbstractColorOp(next, bypass ? Feature::Bypass : Features{}, QStringLiteral("1D curve"))
    , m_value(value)
    , m_bypass(bypass)
{
}

std::optional<Named1DLutType> DrmNamed1DLut::getType(const ColorOp::Operation &op) const
{
    if (const auto tf = std::get_if<ColorTransferFunction>(&op)) {
        switch (tf->tf.type) {
        case TransferFunction::sRGB:
            return Named1DLutType::sRGB;
        case TransferFunction::PerceptualQuantizer:
            return Named1DLutType::PQ_125;
        case TransferFunction::gamma22:
            return Named1DLutType::gamma22;
        default:
            return std::nullopt;
        }
    } else if (const auto invTf = std::get_if<InverseColorTransferFunction>(&op)) {
        switch (invTf->tf.type) {
        case TransferFunction::sRGB:
            return Named1DLutType::sRGB_Inverse;
        case TransferFunction::PerceptualQuantizer:
            return Named1DLutType::PQ_125_Inverse;
        case TransferFunction::gamma22:
            return Named1DLutType::gamma22_Inverse;
        default:
            return std::nullopt;
        }
    } else {
        return std::nullopt;
    }
}

std::optional<DrmAbstractColorOp::Priority> DrmNamed1DLut::colorOpPreference(const ColorOp::Operation &op)
{
    const auto namedLut = getType(op);
    if (!namedLut) {
        return std::nullopt;
    }
    if (m_value->hasEnum(*namedLut)) {
        return Priority::High;
    } else {
        return std::nullopt;
    }
}

void DrmNamed1DLut::program(DrmAtomicCommit *commit, const std::deque<ColorOp::Operation> &operations)
{
    // there may be the input / output scaling matrix in here, ignore it
    const auto curve = std::ranges::find_if(operations, [](const ColorOp::Operation &op) {
        return std::holds_alternative<ColorTransferFunction>(op)
            || std::holds_alternative<InverseColorTransferFunction>(op);
    });
    Q_ASSERT(curve != operations.end());
    commit->addEnum(*m_value, *getType(*curve));
    if (m_bypass) {
        commit->addProperty(*m_bypass, 0);
    }
}

void DrmNamed1DLut::bypass(DrmAtomicCommit *commit)
{
    commit->addProperty(*m_bypass, 1);
}

std::optional<DrmAbstractColorOp::Scaling> DrmNamed1DLut::inputScaling(const ColorOp &op) const
{
    const auto evaluateMat = [](const QMatrix4x4 &mat, float value) {
        return (mat * QVector3D(value, value, value)).x();
    };
    if (const auto tf = std::get_if<InverseColorTransferFunction>(&op.operation)) {
        // PQ_125 has 125 as the maximum value (instead of 10'000 or 1)
        const double transferFunctionScale = tf->tf.type == TransferFunction::PerceptualQuantizer ? 125 : 1;
        const bool rangeMatches = op.input.min >= -0.001 && op.input.max <= transferFunctionScale + 0.001;
        if (rangeMatches && qFuzzyIsNull(tf->tf.minLuminance) && qFuzzyCompare(transferFunctionScale, tf->tf.maxLuminance)) {
            // the transfer function matches exactly, no need for scaling
            return std::nullopt;
        }

        // TODO this should ideally use an offset+scaling op, and not a generic matrix
        QMatrix4x4 mat;
        mat.scale(transferFunctionScale);
        mat.scale(1.0 / (tf->tf.maxLuminance - tf->tf.minLuminance));
        mat.translate(-tf->tf.minLuminance, -tf->tf.minLuminance, -tf->tf.minLuminance);
        return Scaling{
            .scaling = ColorMatrix(mat.inverted()),
            .inverse = ColorOp{
                .input = op.input,
                .inputSpace = op.inputSpace,
                .operation = ColorMatrix(mat),
                .output = ValueRange{
                    .min = evaluateMat(mat, op.input.min),
                    .max = evaluateMat(mat, op.input.max),
                },
                // not linear, not normal nonlinear either though
                .outputSpace = ColorspaceType::AnyNonRGB,
            },
        };
    } else {
        return std::nullopt;
    }
}

std::optional<DrmAbstractColorOp::Scaling> DrmNamed1DLut::outputScaling(const ColorOp &op) const
{
    const auto evaluateMat = [](const QMatrix4x4 &mat, float value) {
        return (mat * QVector3D(value, value, value)).x();
    };
    if (const auto tf = std::get_if<ColorTransferFunction>(&op.operation)) {
        // PQ_125 has 125 as the maximum value (instead of 10'000 or 1)
        const double transferFunctionScale = tf->tf.type == TransferFunction::PerceptualQuantizer ? 125 : 1;
        const bool rangeMatches = op.output.min >= -0.001 && op.output.max <= transferFunctionScale + 0.001;
        if (rangeMatches && qFuzzyIsNull(tf->tf.minLuminance) && qFuzzyCompare(transferFunctionScale, tf->tf.maxLuminance)) {
            // the transfer function matches exactly, no need for scaling
            return std::nullopt;
        }

        // TODO this should ideally use an offset+scaling op, and not a generic matrix
        QMatrix4x4 invMat;
        invMat.translate(tf->tf.minLuminance, tf->tf.minLuminance, tf->tf.minLuminance);
        invMat.scale(tf->tf.maxLuminance - tf->tf.minLuminance);
        invMat.scale(1.0 / transferFunctionScale);
        return Scaling{
            .scaling = ColorMatrix(invMat.inverted()),
            .inverse = ColorOp{
                .input = ValueRange{
                    .min = evaluateMat(invMat, op.output.min),
                    .max = evaluateMat(invMat, op.output.max),
                },
                // not linear, not normal nonlinear either though
                .inputSpace = ColorspaceType::AnyNonRGB,
                .operation = ColorMatrix(invMat),
                .output = op.output,
                .outputSpace = op.outputSpace,
            },
        };
    } else {
        return std::nullopt;
    }
}

DrmColorOp::DrmColorOp(DrmGpu *gpu, uint32_t objectId)
    : DrmObject(gpu, objectId, DRM_MODE_OBJECT_ANY)
    , m_next(this, QByteArrayLiteral("NEXT"))
    , m_type(this, QByteArrayLiteral("TYPE"), {
                                                  QByteArrayLiteral("1D LUT"),
                                                  QByteArrayLiteral("3x4 Matrix"),
                                                  QByteArrayLiteral("3D LUT"),
                                                  QByteArrayLiteral("Multiplier"),
                                                  QByteArrayLiteral("1D Curve"),
                                              })
    , m_data(this, QByteArrayLiteral("DATA"))
    , m_size(this, QByteArrayLiteral("SIZE"))
    , m_bypass(this, QByteArrayLiteral("BYPASS"))
    , m_multiplier(this, QByteArrayLiteral("MULTIPLIER"))
    , m_1dNamedLutType(this, QByteArrayLiteral("CURVE_1D_TYPE"), {
                                                                     QByteArrayLiteral("sRGB EOTF"),
                                                                     QByteArrayLiteral("sRGB Inverse EOTF"),
                                                                     QByteArrayLiteral("PQ 125 EOTF"),
                                                                     QByteArrayLiteral("PQ 125 Inverse EOTF"),
                                                                     QByteArrayLiteral("Gamma 2.2"),
                                                                     QByteArrayLiteral("Gamma 2.2 Inverse"),
                                                                 })
    , m_lut1dInterpolation(this, QByteArrayLiteral("LUT1D_INTERPOLATION"), {QByteArrayLiteral("Linear")})
    , m_lut3dInterpolation(this, QByteArrayLiteral("LUT3D_INTERPOLATION"), {QByteArrayLiteral("Tetrahedal")})
{
}

bool DrmColorOp::init()
{
    return updateProperties();
}

DrmAbstractColorOp *DrmColorOp::colorOp() const
{
    return m_op.get();
}

void DrmColorOp::disable(DrmAtomicCommit *commit)
{
    Q_UNREACHABLE();
}

bool DrmColorOp::updateProperties()
{
    DrmPropertyList props = queryProperties();
    m_next.update(props);
    m_type.update(props);
    m_data.update(props);
    m_size.update(props);
    m_bypass.update(props);
    m_multiplier.update(props);
    m_1dNamedLutType.update(props);
    m_lut1dInterpolation.update(props);
    m_lut3dInterpolation.update(props);

    if (!m_type.isValid()) {
        return false;
    }
    if (m_next.value() != 0) {
        if (!m_nextOp) {
            m_nextOp = std::make_unique<DrmColorOp>(gpu(), m_next.value());
            if (!m_nextOp->init()) {
                return false;
            }
        } else if (!m_nextOp->updateProperties()) {
            return false;
        }
    }

    if (m_op) {
        return true;
    }
    const auto next = m_nextOp ? m_nextOp->colorOp() : nullptr;
    const auto bypassProp = m_bypass.isValid() ? &m_bypass : nullptr;
    if (!m_type.hasEnumForValue(m_type.value())) {
        // we don't know this color operation type (yet)
        // but we have a bypass, so it's fine
        m_op = std::make_unique<UnknownColorOp>(next, bypassProp);
        qCDebug(KWIN_DRM, "skipping unknown color op %lu", m_type.value());
        return true;
    }
    switch (m_type.enumValue()) {
    case Type::Lut1D:
        if (!m_size.isValid() || !m_size.value()) {
            qCWarning(KWIN_DRM, "skipping 1D lut with invalid size property");
            m_op = std::make_unique<UnknownColorOp>(next, bypassProp);
            return true;
        }
        m_op = std::make_unique<DrmLutColorOp32>(next, &m_data, &m_lut1dInterpolation, m_size.value(), bypassProp);
        return true;
    case Type::Matrix3x4:
        m_op = std::make_unique<Matrix3x4ColorOp>(next, &m_data, bypassProp);
        return true;
    case Type::Lut3D: {
        if (!m_size.isValid() || !m_size.value()) {
            qCDebug(KWIN_DRM, "skipping 3D lut with invalid mode index, mode blob or exclusively unsupported modes");
            m_op = std::make_unique<UnknownColorOp>(next, bypassProp);
            return true;
        }
        m_op = std::make_unique<DrmLut3DColorOp>(next, &m_data, bypassProp, m_size.value(), &m_lut3dInterpolation);
        return true;
    }
    case Type::Multiplier:
        if (!m_multiplier.isValid()) {
            qCWarning(KWIN_DRM, "Skipping multiplier with invalid multiplier property");
            m_op = std::make_unique<UnknownColorOp>(next, bypassProp);
            return true;
        }
        m_op = std::make_unique<DrmMultiplier>(next, &m_multiplier, bypassProp);
        return true;
    case Type::NamedLut1D:
        if (!m_1dNamedLutType.isValid()) {
            qCWarning(KWIN_DRM, "Skipping 1D lut with invalid type property");
            m_op = std::make_unique<UnknownColorOp>(next, bypassProp);
            return true;
        }
        m_op = std::make_unique<DrmNamed1DLut>(next, &m_1dNamedLutType, bypassProp);
        return true;
    }
    Q_UNREACHABLE();
}

}
