#include <raylib.h>
#include <stdlib.h>
#include <math.h>
#include <omp.h>

bool is_inside_shape(Color pixel)
{
    return pixel.r == 255 && pixel.g == 255 && pixel.b == 255;
}

Image generate_distance_field(Image input_texture, int search_radius)
{
    int width = input_texture.width;
    int height = input_texture.height;

    Image distance_field = GenImageColor(width, height, (Color){128, 128, 128, 255});
    Color *input_pixels = LoadImageColors(input_texture);
    Color *output_pixels = LoadImageColors(distance_field);

    TraceLog(LOG_INFO, "Using %d threads for parallel processing", omp_get_max_threads());

#pragma omp parallel for schedule(dynamic, 64)
    for (int y = 0; y < height; y++)
    {
        if (omp_get_thread_num() == 0 && y % (height / 10) == 0)
        {
            TraceLog(LOG_INFO, "Progress: %d%%", (y * 100) / height);
        }

        for (int x = 0; x < width; x++)
        {
            int index = y * width + x;
            bool is_inside = is_inside_shape(input_pixels[index]);

            float min_distance_sq = search_radius * search_radius;
            bool found_boundary = false;

            for (int dy = -search_radius; dy <= search_radius; dy += 1)
            {
                for (int dx = -search_radius; dx <= search_radius; dx += 1)
                {
                    int sample_x = x + dx;
                    int sample_y = y + dy;

                    if (sample_x < 0 || sample_x >= width || sample_y < 0 || sample_y >= height) continue;

                    int sample_index = sample_y * width + sample_x;
                    bool sample_is_inside = is_inside_shape(input_pixels[sample_index]);

                    if (sample_is_inside != is_inside)
                    {
                        float distance_sq = dx * dx + dy * dy;
                        if (distance_sq < min_distance_sq)
                        {
                            min_distance_sq = distance_sq;
                            found_boundary = true;
                        }
                    }
                }
            }

            float distance = found_boundary ? sqrtf(min_distance_sq) : (float)search_radius;
            float normalized_distance = distance / (float)search_radius;

            unsigned char distance_value;
            if (is_inside)
            {
                distance_value = (unsigned char)(127.0f - 127.0f * normalized_distance);
            }
            else
            {
                distance_value = (unsigned char)(128.0f + 127.0f * normalized_distance);
            }

            output_pixels[index] = (Color){distance_value, distance_value, distance_value, 255};
        }
    }

    for (int i = 0; i < width * height; i++)
    {
        ImageDrawPixel(&distance_field, i % width, i / width, output_pixels[i]);
    }

    UnloadImageColors(input_pixels);
    UnloadImageColors(output_pixels);

    return distance_field;
}

int main(int argc, char *argv[])
{
    if (argc < 3)
    {
        TraceLog(LOG_ERROR, "Usage: %s <input_image> <output_image> [search_radius]", argv[0]);
        return 1;
    }

    const char *input_path = argv[1];
    const char *output_path = argv[2];
    int search_radius = (argc >= 4) ? atoi(argv[3]) : 100;

    Image input_image = LoadImage(input_path);
    if (input_image.data == NULL)
    {
        TraceLog(LOG_ERROR, "Error: Failed to load input image '%s'", input_path);
        return 1;
    }

    Image distance_field_image = generate_distance_field(input_image, search_radius);
    if (!ExportImage(distance_field_image, output_path))
    {
        TraceLog(LOG_ERROR, "Failed to save output image '%s'", output_path);
        UnloadImage(input_image);
        UnloadImage(distance_field_image);
        return 1;
    }

    TraceLog(LOG_INFO, "Distance field saved to '%s'", output_path);

    UnloadImage(input_image);
    UnloadImage(distance_field_image);

    return 0;
}