#pragma once
#include <SDL3/SDL.h>
#include <SDL3_image/SDL_image.h>
#include <iostream>
#include <vector>
#include <random>
#include <algorithm>

class PixelShuffler {
private:
    SDL_Renderer* renderer;
    SDL_Surface* surface;
    SDL_Texture* texture;
    // Store each shuffle step as a list of swaps
    struct Swap {
        int x1, y1, x2, y2;
    };
    std::vector<std::vector<Swap>> shuffleHistory; // Each shuffle creates a list of swaps
    std::mt19937 rng;


public:
    PixelShuffler(SDL_Renderer* r, SDL_Surface* surf) : renderer(r), surface(surf), rng(std::random_device{}()) {
    }

    // Get pixel at position (x, y)
    Uint32 getPixel(int x, int y) {
        if (x < 0 || x >= surface->w || y < 0 || y >= surface->h) {
            return 0;
        }

        Uint8* pixels = (Uint8*)surface->pixels;
        int bpp = SDL_BYTESPERPIXEL(surface->format);
        Uint8* pixel = pixels + y * surface->pitch + x * bpp;

        switch (bpp) {
            case 1:
                return *pixel;
            case 2:
                return *(Uint16*)pixel;
            case 3:
                if (SDL_BYTEORDER == SDL_BIG_ENDIAN)
                    return ((Uint32)pixel[0] << 16)
                        | ((Uint32)pixel[1] << 8)
                        |  (Uint32)pixel[2];
                else
                    return  (Uint32)pixel[0]
                        | ((Uint32)pixel[1] << 8)
                        | ((Uint32)pixel[2] << 16);
            case 4:
                return *(Uint32*)pixel;
            default:
                return 0;
        }
    }

    // Set pixel at position (x, y)
    void setPixel(int x, int y, Uint32 color) {
        if (x < 0 || x >= surface->w || y < 0 || y >= surface->h) {
            return;
        }

        Uint8* pixels = (Uint8*)surface->pixels;
        int bpp = SDL_BYTESPERPIXEL(surface->format);
        Uint8* pixel = pixels + y * surface->pitch + x * bpp;

        switch (bpp) {
            case 1:
                *pixel = color;
                break;
            case 2:
                *(Uint16*)pixel = color;
                break;
            case 3:
                if (SDL_BYTEORDER == SDL_BIG_ENDIAN) {
                    pixel[0] = (color >> 16) & 0xff;
                    pixel[1] = (color >> 8) & 0xff;
                    pixel[2] = color & 0xff;
                } else {
                    pixel[0] = color & 0xff;
                    pixel[1] = (color >> 8) & 0xff;
                    pixel[2] = (color >> 16) & 0xff;
                }
                break;
            case 4:
                *(Uint32*)pixel = color;
                break;
        }
    }

    // Shuffle pixels by swapping each pixel with a random neighbor
    void shufflePixels() {
        SDL_LockSurface(surface);

        int width = surface->w;
        int height = surface->h;

        std::vector<Swap> currentShuffleSwaps;

        // For each pixel, swap with a random neighbor
        for (int y = 0; y < height; y++) {
            for (int x = 0; x < width; x++) {
                // Get random neighbor direction (0=right, 1=down, 2=left, 3=up)
                int direction = rng() % 4;
                int nx = x, ny = y;

                switch (direction) {
                    case 0: nx = (x + 1) % width; break;
                    case 1: ny = (y + 1) % height; break;
                    case 2: nx = (x - 1 + width) % width; break;
                    case 3: ny = (y - 1 + height) % height; break;
                }

                // Only swap if neighbor is valid
                if (nx >= 0 && nx < width && ny >= 0 && ny < height) {
                    // Swap pixels
                    Uint32 pixel1 = getPixel(x, y);
                    Uint32 pixel2 = getPixel(nx, ny);
                    setPixel(x, y, pixel2);
                    setPixel(nx, ny, pixel1);

                    // Record this swap
                    currentShuffleSwaps.push_back({x, y, nx, ny});
                }
            }
        }

        // Add this shuffle to history
        shuffleHistory.push_back(currentShuffleSwaps);

        SDL_UnlockSurface(surface);
        texture = SDL_CreateTextureFromSurface(renderer, getSurface());
    }

    // Unshuffle by reversing the last shuffle step
    bool unshuffleStep() {
        if (shuffleHistory.empty()) {
            return false; // Nothing to unshuffle
        }

        SDL_LockSurface(surface);

        // Get the last shuffle
        std::vector<Swap>& lastShuffle = shuffleHistory.back();

        // Reverse the swaps in reverse order
        for (auto it = lastShuffle.rbegin(); it != lastShuffle.rend(); ++it) {
            const Swap& swap = *it;

            // Swap back
            Uint32 pixel1 = getPixel(swap.x1, swap.y1);
            Uint32 pixel2 = getPixel(swap.x2, swap.y2);
            setPixel(swap.x1, swap.y1, pixel2);
            setPixel(swap.x2, swap.y2, pixel1);
        }

        // Remove this shuffle from history
        shuffleHistory.pop_back();

        SDL_UnlockSurface(surface);
        texture = SDL_CreateTextureFromSurface(renderer, getSurface());
        return true;
    }

    // Check if there are shuffles to undo
    bool canUnshuffle() const {
        return !shuffleHistory.empty();
    }

    // Get number of shuffle steps in history
    int getShuffleCount() const {
        return shuffleHistory.size();
    }

    SDL_Surface* getSurface() { return surface; }


    // Get the working texture to render
    SDL_Texture* getTexture() {
        return texture;
    }
};