#include "effects/bitcrush.hpp"

#include <algorithm>
#include <array>
#include <vector>

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

namespace effector {

const std::vector<ParamSpec>& BitCrushEffect::params() const {
    static const std::vector<ParamSpec> specs = {
        ParamSpec::makeChoice("crush_op", "Operation", "mixed",
                              {"mixed", "posterize", "xor", "bitrotate", "channel_swap"},
                              "posterize drops the low bits, xor flips random low bits, "
                              "bitrotate turns each byte's bits round, channel_swap shuffles "
                              "red, green and blue; mixed picks one for every band."),
        ParamSpec::makeInt("crush_bits", "Bits", 3, 1, 7,
                           "How many bits the operation reaches: dropped by posterize, flipped "
                           "at most by xor, turned at most by bitrotate."),
        ParamSpec::makeDouble("crush_coverage", "Coverage", 0.5, 0.0, 1.0,
                              "Chance of each band being crushed."),
        ParamSpec::makeInt("crush_band_height", "Band height", 32, 1, 4096,
                           "Thickness in pixels of each band."),
    };
    return specs;
}

namespace {

enum class Op : std::uint8_t { None, Posterize, Xor, Rotate, Swap };

struct Band {
    Op op = Op::None;
    int rotate = 0;
    std::array<std::uint8_t, 3> mask{};
    std::array<int, 3> order{0, 1, 2};
};

Op opNamed(const std::string& name) {
    if (name == "posterize") return Op::Posterize;
    if (name == "xor") return Op::Xor;
    if (name == "bitrotate") return Op::Rotate;
    if (name == "channel_swap") return Op::Swap;
    return Op::None;  // mixed: drawn per band
}

}  // namespace

Image BitCrushEffect::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 Op fixed = opNamed(paramChoice(values, "crush_op", "mixed"));
    const int bits = static_cast<int>(std::clamp(paramInt(values, "crush_bits", 3), 1L, 7L));
    const double coverage = std::clamp(paramDouble(values, "crush_coverage", 0.5), 0.0, 1.0);
    const int bandHeight = static_cast<int>(std::max(1L, paramInt(values, "crush_band_height", 32)));

    Image canvas = canvasFromSources(sources, width, height, rng);

    const int bandCount = (height + bandHeight - 1) / bandHeight;
    std::vector<Band> plan(static_cast<std::size_t>(bandCount));
    for (auto& band : plan) {
        if (!chance(rng.next(), coverage)) continue;
        band.op = fixed != Op::None ? fixed : static_cast<Op>(1 + rng.below(3));
        switch (band.op) {
            case Op::Xor:
                for (auto& mask : band.mask) {
                    mask = static_cast<std::uint8_t>(rng.inRange(1, (1L << bits) - 1));
                }
                break;
            case Op::Rotate:
                band.rotate = static_cast<int>(rng.inRange(1, bits));
                break;
            case Op::Swap:
                // Fisher-Yates, never leaving the order as it was.
                do {
                    band.order = {0, 1, 2};
                    for (int i = 2; i > 0; --i) {
                        std::swap(band.order[static_cast<std::size_t>(i)],
                                  band.order[rng.below(static_cast<std::uint64_t>(i))]);
                    }
                } while (band.order[0] == 0 && band.order[1] == 1);
                break;
            default:
                break;
        }
    }

    const std::uint8_t keep = static_cast<std::uint8_t>(0xFF << bits);

    parallelFor(static_cast<std::size_t>(height), [&](std::size_t begin, std::size_t end) {
        for (std::size_t y = begin; y < end; ++y) {
            const Band& band = plan[y / static_cast<std::size_t>(bandHeight)];
            if (band.op == Op::None) continue;
            std::uint8_t* row = canvas.row(static_cast<int>(y));
            for (int x = 0; x < width; ++x) {
                std::uint8_t* px = row + static_cast<std::size_t>(x) * 3;
                switch (band.op) {
                    case Op::Posterize:
                        for (int c = 0; c < 3; ++c) px[c] &= keep;
                        break;
                    case Op::Xor:
                        for (int c = 0; c < 3; ++c) px[c] ^= band.mask[static_cast<std::size_t>(c)];
                        break;
                    case Op::Rotate:
                        for (int c = 0; c < 3; ++c) {
                            px[c] = static_cast<std::uint8_t>((px[c] << band.rotate) |
                                                              (px[c] >> (8 - band.rotate)));
                        }
                        break;
                    case Op::Swap: {
                        const std::uint8_t in[3] = {px[0], px[1], px[2]};
                        for (int c = 0; c < 3; ++c) px[c] = in[band.order[static_cast<std::size_t>(c)]];
                        break;
                    }
                    case Op::None:
                        break;
                }
            }
        }
    });

    return canvas;
}

}  // namespace effector