#include "effects/corrupt.hpp"

#include <algorithm>
#include <cstring>
#include <vector>

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

namespace effector {

const std::vector<ParamSpec>& CorruptEffect::params() const {
    static const std::vector<ParamSpec> specs = {
        ParamSpec::makeInt("corrupt_block", "Block size", 16, 2, 256,
                           "Width and height in pixels of one block, as a codec would cut them."),
        ParamSpec::makeDouble("corrupt_rate", "Rate", 0.08, 0.0, 1.0,
                              "Chance of each block being damaged."),
        ParamSpec::makeChoice("corrupt_kinds", "Damage", "mixed",
                              {"mixed", "repeat", "displace", "colour"},
                              "repeat freezes a block into a copy of the one to its left, "
                              "displace pulls one in from elsewhere in the image, colour knocks "
                              "it off colour; mixed picks one of them for every block."),
        ParamSpec::makeBool("corrupt_bleed", "Bleed", false,
                            "Carry each damaged block's fault on to the end of its row, the way "
                            "one broken bit spoils the rest of a stream."),
    };
    return specs;
}

namespace {

enum class Damage : std::uint8_t { None, Repeat, Displace, Colour };

struct Block {
    Damage damage = Damage::None;
    int sx = 0;  // displace: where the block is read from
    int sy = 0;
    int shift[3] = {0, 0, 0};  // colour: added to each channel
};

}  // namespace

Image CorruptEffect::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.");

    const int size = static_cast<int>(std::max(2L, paramInt(values, "corrupt_block", 16)));
    const double rate = std::clamp(paramDouble(values, "corrupt_rate", 0.08), 0.0, 1.0);
    const std::string kinds = paramChoice(values, "corrupt_kinds", "mixed");
    const bool bleed = paramBool(values, "corrupt_bleed", false);

    const Image input = canvasFromSources(sources, width, height, rng);

    const int cols = (width + size - 1) / size;
    const int rows = (height + size - 1) / size;

    auto drawBlock = [&](int bx) {
        Block block;
        if (kinds == "repeat") block.damage = Damage::Repeat;
        else if (kinds == "displace") block.damage = Damage::Displace;
        else if (kinds == "colour") block.damage = Damage::Colour;
        else block.damage = static_cast<Damage>(1 + rng.below(2));
        // The first block of a row has nothing to its left to repeat.
        if (block.damage == Damage::Repeat && bx == 0) block.damage = Damage::Displace;
        if (block.damage == Damage::Displace) {
            block.sx = static_cast<int>(rng.inRange(0, std::max(0, width - size)));
            block.sy = static_cast<int>(rng.inRange(0, std::max(0, height - size)));
        } else if (block.damage == Damage::Colour) {
            for (int& shift : block.shift) shift = static_cast<int>(rng.inRange(-128, 128));
        }
        return block;
    };

    // The whole damage map is drawn up front, block by block in reading order.
    std::vector<Block> plan(static_cast<std::size_t>(cols) * rows);
    for (int by = 0; by < rows; ++by) {
        const Block* carried = nullptr;
        for (int bx = 0; bx < cols; ++bx) {
            Block& block = plan[static_cast<std::size_t>(by) * cols + bx];
            if (carried != nullptr) {
                block = *carried;
                // A displaced run keeps moving along with the row, so the
                // stretch reads as one shifted strip rather than one block
                // over and over.
                if (block.damage == Damage::Displace) block.sx += size;
                carried = &block;
                continue;
            }
            if (!chance(rng.next(), rate)) continue;
            block = drawBlock(bx);
            if (bleed) carried = &block;
        }
    }

    Image canvas = input;
    // Block rows are independent (a repeat only ever reads to its left), so
    // they run in parallel and each walks its row in order.
    parallelFor(static_cast<std::size_t>(rows), [&](std::size_t begin, std::size_t end) {
        for (std::size_t by = begin; by < end; ++by) {
            const int y0 = static_cast<int>(by) * size;
            const int h = std::min(size, height - y0);
            for (int bx = 0; bx < cols; ++bx) {
                const Block& block = plan[by * static_cast<std::size_t>(cols) + bx];
                if (block.damage == Damage::None) continue;
                const int x0 = bx * size;
                const int w = std::min(size, width - x0);
                const std::size_t bytes = static_cast<std::size_t>(w) * 3;

                for (int line = 0; line < h; ++line) {
                    std::uint8_t* out = canvas.row(y0 + line) + static_cast<std::size_t>(x0) * 3;
                    switch (block.damage) {
                        case Damage::Repeat:
                            std::memmove(out, out - static_cast<std::size_t>(size) * 3, bytes);
                            break;
                        case Damage::Displace: {
                            const int sy = std::min(block.sy + line, height - 1);
                            const std::uint8_t* in = input.row(sy);
                            for (int x = 0; x < w; ++x) {
                                const int sx = (block.sx + x) % width;
                                std::memcpy(out + static_cast<std::size_t>(x) * 3,
                                            in + static_cast<std::size_t>(sx) * 3, 3);
                            }
                            break;
                        }
                        case Damage::Colour:
                            for (std::size_t i = 0; i < bytes; ++i) {
                                out[i] = static_cast<std::uint8_t>(
                                    std::clamp(out[i] + block.shift[i % 3], 0, 255));
                            }
                            break;
                        case Damage::None:
                            break;
                    }
                }
            }
        }
    });

    return canvas;
}

}  // namespace effector