// D:\Dev\Mosis\MosisService\designer\src\testing\visual_capture.cpp
#include "visual_capture.h"
#include <glad/glad.h>
#include <png.h>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cmath>

namespace mosis::testing {

VisualCapture::VisualCapture(int width, int height)
    : m_width(width)
    , m_height(height)
{
}

std::vector<uint8_t> VisualCapture::CaptureToBuffer() {
    std::vector<uint8_t> pixels(m_width * m_height * 4);
    
    // Read pixels from framebuffer
    glReadPixels(0, 0, m_width, m_height, GL_RGBA, GL_UNSIGNED_BYTE, pixels.data());
    
    // OpenGL reads bottom-up, flip vertically
    std::vector<uint8_t> flipped(pixels.size());
    int row_size = m_width * 4;
    for (int y = 0; y < m_height; y++) {
        memcpy(
            flipped.data() + y * row_size,
            pixels.data() + (m_height - 1 - y) * row_size,
            row_size
        );
    }
    
    return flipped;
}

bool VisualCapture::CaptureScreenshot(const std::string& path) {
    auto pixels = CaptureToBuffer();
    if (pixels.empty()) return false;
    
    FILE* fp = fopen(path.c_str(), "wb");
    if (!fp) return false;
    
    png_structp png = png_create_write_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    if (!png) {
        fclose(fp);
        return false;
    }
    
    png_infop info = png_create_info_struct(png);
    if (!info) {
        png_destroy_write_struct(&png, nullptr);
        fclose(fp);
        return false;
    }
    
    if (setjmp(png_jmpbuf(png))) {
        png_destroy_write_struct(&png, &info);
        fclose(fp);
        return false;
    }
    
    png_init_io(png, fp);
    
    png_set_IHDR(
        png, info, m_width, m_height, 8,
        PNG_COLOR_TYPE_RGBA, PNG_INTERLACE_NONE,
        PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT
    );
    png_write_info(png, info);
    
    // Write rows
    std::vector<png_bytep> rows(m_height);
    for (int y = 0; y < m_height; y++) {
        rows[y] = pixels.data() + y * m_width * 4;
    }
    png_write_image(png, rows.data());
    png_write_end(png, nullptr);
    
    png_destroy_write_struct(&png, &info);
    fclose(fp);
    
    return true;
}

// Helper struct for loaded PNG data
struct PNGImage {
    int width = 0;
    int height = 0;
    std::vector<uint8_t> pixels;  // RGBA format
    bool valid = false;
};

// Load a PNG file into memory
static PNGImage LoadPNG(const std::string& path) {
    PNGImage image;

    FILE* fp = fopen(path.c_str(), "rb");
    if (!fp) return image;

    // Check PNG signature
    uint8_t header[8];
    if (fread(header, 1, 8, fp) != 8 || png_sig_cmp(header, 0, 8)) {
        fclose(fp);
        return image;
    }

    png_structp png = png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    if (!png) {
        fclose(fp);
        return image;
    }

    png_infop info = png_create_info_struct(png);
    if (!info) {
        png_destroy_read_struct(&png, nullptr, nullptr);
        fclose(fp);
        return image;
    }

    if (setjmp(png_jmpbuf(png))) {
        png_destroy_read_struct(&png, &info, nullptr);
        fclose(fp);
        return image;
    }

    png_init_io(png, fp);
    png_set_sig_bytes(png, 8);
    png_read_info(png, info);

    image.width = png_get_image_width(png, info);
    image.height = png_get_image_height(png, info);
    int color_type = png_get_color_type(png, info);
    int bit_depth = png_get_bit_depth(png, info);

    // Convert to RGBA if necessary
    if (bit_depth == 16)
        png_set_strip_16(png);
    if (color_type == PNG_COLOR_TYPE_PALETTE)
        png_set_palette_to_rgb(png);
    if (color_type == PNG_COLOR_TYPE_GRAY && bit_depth < 8)
        png_set_expand_gray_1_2_4_to_8(png);
    if (png_get_valid(png, info, PNG_INFO_tRNS))
        png_set_tRNS_to_alpha(png);
    if (color_type == PNG_COLOR_TYPE_RGB ||
        color_type == PNG_COLOR_TYPE_GRAY ||
        color_type == PNG_COLOR_TYPE_PALETTE)
        png_set_filler(png, 0xFF, PNG_FILLER_AFTER);
    if (color_type == PNG_COLOR_TYPE_GRAY ||
        color_type == PNG_COLOR_TYPE_GRAY_ALPHA)
        png_set_gray_to_rgb(png);

    png_read_update_info(png, info);

    // Allocate memory and read image
    image.pixels.resize(image.width * image.height * 4);
    std::vector<png_bytep> row_pointers(image.height);
    for (int y = 0; y < image.height; y++) {
        row_pointers[y] = image.pixels.data() + y * image.width * 4;
    }

    png_read_image(png, row_pointers.data());
    png_read_end(png, nullptr);

    png_destroy_read_struct(&png, &info, nullptr);
    fclose(fp);

    image.valid = true;
    return image;
}

// Compare two pixels with tolerance
static bool PixelsMatch(const uint8_t* p1, const uint8_t* p2, int tolerance = 2) {
    for (int c = 0; c < 4; c++) {
        if (std::abs(static_cast<int>(p1[c]) - static_cast<int>(p2[c])) > tolerance) {
            return false;
        }
    }
    return true;
}

float VisualCapture::CompareImages(const std::string& path1, const std::string& path2) {
    // Load both images
    PNGImage img1 = LoadPNG(path1);
    PNGImage img2 = LoadPNG(path2);

    // Check if both loaded successfully
    if (!img1.valid || !img2.valid) {
        return 1.0f;  // Can't compare, assume completely different
    }

    // Check dimensions match
    if (img1.width != img2.width || img1.height != img2.height) {
        return 1.0f;  // Different dimensions = completely different
    }

    // Count differing pixels
    int total_pixels = img1.width * img1.height;
    int diff_pixels = 0;

    for (int i = 0; i < total_pixels; i++) {
        const uint8_t* p1 = img1.pixels.data() + i * 4;
        const uint8_t* p2 = img2.pixels.data() + i * 4;

        if (!PixelsMatch(p1, p2)) {
            diff_pixels++;
        }
    }

    // Return difference ratio (0.0 = identical, 1.0 = completely different)
    return static_cast<float>(diff_pixels) / static_cast<float>(total_pixels);
}

} // namespace mosis::testing
