#include "effects/overlay.hpp"

#include <algorithm>
#include <cmath>
#include <cstdint>

#include "core/parallel.hpp"
#include "effects/source_fill.hpp"

namespace effector {

const std::vector<ParamSpec>& OverlayEffect::params() const {
    static const std::vector<ParamSpec> specs = {
        ParamSpec::makePath("overlay_image", "Image",
                            "The picture laid over the input: a logo, a texture, a frame."),
        ParamSpec::makeChoice("overlay_fit", "Fitting", "fit",
                              {"fit", "fill", "stretch", "original", "tile"},
                              "fit keeps the proportions and shows the whole image; fill keeps "
                              "them and covers the canvas, cropping what does not fit; stretch "
                              "distorts it to the canvas exactly; original leaves it at its own "
                              "size; tile repeats it at its own size over the whole canvas."),
        ParamSpec::makeDouble("overlay_scale", "Scale", 1.0, 0.05, 8.0,
                              "Resizes the fitted image (the tile, for tiling) by this factor.",
                              0.05),
        ParamSpec::makeBool("overlay_center", "Centre", true,
                            "Places the image in the middle of the canvas rather than at its top "
                            "left corner; the offset is counted from there either way."),
        ParamSpec::makeInt("overlay_offset_x", "Offset X", 0, -20000, 20000,
                           "Moves the image right by this many pixels (left when negative)."),
        ParamSpec::makeInt("overlay_offset_y", "Offset Y", 0, -20000, 20000,
                           "Moves the image down by this many pixels (up when negative)."),
        ParamSpec::makeDouble("overlay_opacity", "Opacity", 1.0, 0.0, 1.0,
                              "How strongly the image is mixed over the input; its own "
                              "transparency is taken into account as well."),
    };
    return specs;
}

namespace {

// Where the overlay lands on the canvas, in canvas pixels.
struct Placement {
    int x = 0;
    int y = 0;
    int width = 1;
    int height = 1;
};

Placement placeOverlay(const std::string& fit, double scale, bool centre, int offsetX, int offsetY,
                       int srcWidth, int srcHeight, int width, int height) {
    double drawWidth = srcWidth;
    double drawHeight = srcHeight;
    if (fit == "stretch") {
        drawWidth = width;
        drawHeight = height;
    } else if (fit == "fit" || fit == "fill") {
        const double byWidth = static_cast<double>(width) / srcWidth;
        const double byHeight = static_cast<double>(height) / srcHeight;
        const double factor = fit == "fit" ? std::min(byWidth, byHeight) : std::max(byWidth, byHeight);
        drawWidth = srcWidth * factor;
        drawHeight = srcHeight * factor;
    }
    // "original" and "tile" keep the file's own size; the scale below applies
    // to every mode, so a fitted image can still be nudged a little larger.

    Placement placement;
    placement.width = std::max(1, static_cast<int>(std::lround(drawWidth * scale)));
    placement.height = std::max(1, static_cast<int>(std::lround(drawHeight * scale)));
    placement.x = centre ? (width - placement.width) / 2 : 0;
    placement.y = centre ? (height - placement.height) / 2 : 0;
    placement.x += offsetX;
    placement.y += offsetY;
    return placement;
}

int wrapIndex(int value, int limit) {
    const int wrapped = value % limit;
    return wrapped < 0 ? wrapped + limit : wrapped;
}

// Bilinear sample of the overlay at a point in its own pixel coordinates,
// clamped at the edges or wrapped around when the image is being tiled.
void sampleOverlay(const Image& src, const std::uint8_t* alpha, double u, double v, bool wrap,
                   int rgb[3], int& outAlpha) {
    const int w = src.width();
    const int h = src.height();
    const double fu = std::floor(u);
    const double fv = std::floor(v);
    const int u0 = static_cast<int>(fu);
    const int v0 = static_cast<int>(fv);
    const double du = u - fu;
    const double dv = v - fv;

    int xs[2];
    int ys[2];
    for (int i = 0; i < 2; ++i) {
        xs[i] = wrap ? wrapIndex(u0 + i, w) : std::clamp(u0 + i, 0, w - 1);
        ys[i] = wrap ? wrapIndex(v0 + i, h) : std::clamp(v0 + i, 0, h - 1);
    }

    double accum[3] = {0.0, 0.0, 0.0};
    double accumAlpha = 0.0;
    for (int j = 0; j < 2; ++j) {
        const double wy = j == 0 ? 1.0 - dv : dv;
        const std::uint8_t* row = src.row(ys[j]);
        const std::uint8_t* alphaRow = alpha + static_cast<std::size_t>(ys[j]) * w;
        for (int i = 0; i < 2; ++i) {
            const double weight = wy * (i == 0 ? 1.0 - du : du);
            const std::uint8_t* px = row + static_cast<std::size_t>(xs[i]) * 3;
            const double alphaWeight = weight * alphaRow[xs[i]];
            // Colour weighted by alpha: a transparent pixel in a PNG is usually
            // black, and averaging it in unweighted draws a dark fringe around
            // everything the image cuts out.
            for (int c = 0; c < 3; ++c) accum[c] += px[c] * alphaWeight;
            accumAlpha += alphaWeight;
        }
    }
    outAlpha = std::clamp(static_cast<int>(accumAlpha + 0.5), 0, 255);
    for (int c = 0; c < 3; ++c) {
        rgb[c] = accumAlpha <= 0.0 ? 0
                                   : std::clamp(static_cast<int>(accum[c] / accumAlpha + 0.5), 0,
                                                255);
    }
}

}  // namespace

void OverlayEffect::loadOverlay(const std::string& path, Image& image,
                                std::vector<std::uint8_t>& alpha) const {
    std::lock_guard<std::mutex> lock(cacheMutex_);
    if (cachedPath_ != path || cached_.empty()) {
        std::vector<std::uint8_t> loadedAlpha;
        Image loaded;
        try {
            loaded = Image::load(path, &loadedAlpha);
        } catch (const std::exception& error) {
            throw EffectError("Cannot read the overlay image '" + path + "': " + error.what());
        }
        if (loaded.empty()) throw EffectError("The overlay image '" + path + "' is empty.");
        cached_ = std::move(loaded);
        cachedAlpha_ = std::move(loadedAlpha);
        cachedPath_ = path;
    }
    image = cached_;
    alpha = cachedAlpha_;
}

Image OverlayEffect::render(const std::vector<Image>& sources, int width, int height, Rng& rng,
                            const ParamValues& values) const {
    if (width < 1 || height < 1) throw EffectError("Canvas size must be at least 1x1 pixels.");

    Image canvas(width, height);
    if (!sources.empty()) {
        // Drawn before any threading, so that a seed always gives the same image.
        const std::size_t index = rng.below(sources.size() - 1);
        const Image& src = sources[index];
        const int sx = static_cast<int>(rng.inRange(0, std::max(0, src.width() - width)));
        const int sy = static_cast<int>(rng.inRange(0, std::max(0, src.height() - height)));
        fillFromSource(canvas, src, sx, sy);
    }

    const std::string path = paramText(values, "overlay_image", "");
    const double opacity = std::clamp(paramDouble(values, "overlay_opacity", 1.0), 0.0, 1.0);
    // No image chosen yet is the state the interface starts in, not an error:
    // the stage simply passes its input through.
    if (path.empty() || opacity <= 0.0) return canvas;

    Image overlay;
    std::vector<std::uint8_t> alpha;
    loadOverlay(path, overlay, alpha);

    const std::string fit = paramChoice(values, "overlay_fit", "fit");
    const Placement place = placeOverlay(
        fit, std::clamp(paramDouble(values, "overlay_scale", 1.0), 0.05, 8.0),
        paramBool(values, "overlay_center", true),
        static_cast<int>(paramInt(values, "overlay_offset_x", 0)),
        static_cast<int>(paramInt(values, "overlay_offset_y", 0)), overlay.width(),
        overlay.height(), width, height);

    const bool tiling = fit == "tile";
    // Source pixels per canvas pixel, the same both ways for every mode but
    // stretch, where the two axes are scaled apart on purpose.
    const double stepX = static_cast<double>(overlay.width()) / place.width;
    const double stepY = static_cast<double>(overlay.height()) / place.height;

    const int firstRow = tiling ? 0 : std::max(0, place.y);
    const int lastRow = tiling ? height : std::min(height, place.y + place.height);
    const int firstCol = tiling ? 0 : std::max(0, place.x);
    const int lastCol = tiling ? width : std::min(width, place.x + place.width);
    if (firstRow >= lastRow || firstCol >= lastCol) return canvas;  // entirely off the canvas

    const int mix = static_cast<int>(opacity * 255.0 + 0.5);
    const std::uint8_t* alphaData = alpha.data();

    parallelFor(static_cast<std::size_t>(lastRow - firstRow), [&](std::size_t begin,
                                                                 std::size_t end) {
        int rgb[3];
        int pixelAlpha = 255;
        for (std::size_t i = begin; i < end; ++i) {
            const int y = firstRow + static_cast<int>(i);
            const double v = (y - place.y + 0.5) * stepY - 0.5;
            std::uint8_t* out = canvas.row(y);
            for (int x = firstCol; x < lastCol; ++x) {
                const double u = (x - place.x + 0.5) * stepX - 0.5;
                sampleOverlay(overlay, alphaData, u, v, tiling, rgb, pixelAlpha);

                const int weight = pixelAlpha * mix / 255;
                if (weight == 0) continue;
                std::uint8_t* px = out + static_cast<std::size_t>(x) * 3;
                for (int c = 0; c < 3; ++c) {
                    px[c] = static_cast<std::uint8_t>(px[c] + (rgb[c] - px[c]) * weight / 255);
                }
            }
        }
    });

    return canvas;
}

}  // namespace effector