#include <defocus/defocus.h>

#include <math.h>
#include <stdlib.h>
#include <string.h>

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

/* *********** *
 *             *
 *    Vec 3    *
 *             *
 * *********** */

/* We can later use this to store 4 vecs for simd */
typedef struct
{
    float x;
    float y;
    float z;
} df_v3;

static df_v3 normalize(df_v3 v)
{
    float l = sqrtf(v.x * v.x + v.y * v.y + v.z * v.z);
    return (df_v3){v.x / l, v.y / l, v.z / l};
}

static float dot(df_v3 a, df_v3 b) { return a.x * b.x + a.y * b.y + a.z * b.z; }

/* ****************** *
 *                    *
 *    List of hits    *
 *                    *
 * ****************** */

typedef struct
{
    df_v3 at;
    df_v3 normal;
    float ray_t;
    int front_face; /* 1 if we hit the front face */
} df_hit_record;

static void set_face_normal(df_hit_record *record, df_v3 ray_dir, df_v3 outward_normal)
{
    record->front_face = dot(ray_dir, outward_normal) < 0;
    record->normal =
        record->front_face ? outward_normal : (df_v3){-outward_normal.x, -outward_normal.y, -outward_normal.z};
}

typedef struct
{
    df_hit_record *hits;
    size_t count;
    size_t capacity;
} df_hit_list;

static void emit_hit(df_hit_list *list, df_hit_record record)
{
    if (list->count == list->capacity) {
        size_t cap2 = (list->capacity > 0) ? 2 * list->capacity : 128;
        df_hit_record *tmp = realloc(list->hits, sizeof(df_hit_record) * cap2);
        if (!tmp)
            return;
        list->hits = tmp;
        list->capacity = cap2;
    }
    list->hits[list->count++] = record;
}

static df_hit_list merge_hit_lists(const df_hit_list **lists, unsigned int count)
{
    size_t total_size = 0;
    for (unsigned int i = 0; i < count; ++i)
        total_size += lists[i]->count;

    df_hit_list out;
    out.count = total_size;
    out.hits = malloc(sizeof(df_hit_record) * total_size);
    if (!out.hits)
        return out;

    df_hit_record *dst = out.hits;
    for (unsigned int i = 0; i < count; ++i) {
        memcpy(dst, lists[i]->hits, sizeof(df_hit_record) * lists[i]->count);
        dst += lists[i]->count;
    }

    return out;
}

/* ********************* *
 *                       *
 *    Image functions    *
 *                       *
 * ********************* */

#define MAX_IMAGES 1024

struct
{
    struct
    {
        stbi_uc *pixels;
        int w;
        int h;
    } images[MAX_IMAGES];
    unsigned int image_count;
} _image_table;

DF_API df_image_handle df_load_image(const char *path, int *w, int *h)
{
    if (_image_table.image_count == MAX_IMAGES)
        return 0;

    int width, height, c;
    stbi_uc *pixels = stbi_load(path, &width, &height, &c, 4);
    if (!pixels) {
        return 0;
    }
    df_image_handle handle = (_image_table.image_count + 1);
    _image_table.image_count++;
    _image_table.images[handle].pixels = pixels;
    _image_table.images[handle].w = width;
    _image_table.images[handle].h = height;

    if (w) *w = width;
    if (h) *h = height;
    return handle;
}

/* ********************************* *
 *                                   *
 *    Intersection test functions    *
 *
 * ********************************* */

typedef struct
{
    float ray_t;
    float img_u;
    float img_v;
    df_image_handle image;
} df_hit; 

/* -1 => does not hit, >= 0 => hit (sphere index) */
static float sphere_test(float ray_origin_x,
                         float ray_origin_y,
                         float ray_origin_z,
                         float ray_dx,
                         float ray_dy,
                         float ray_dz,
                         const df_sphere *spheres,
                         unsigned int sphere_count)
{
    float result = -1.f;
    for (unsigned int i = 0; i < sphere_count; ++i) {
        float delta_x = ray_origin_x - spheres[i].center_x;
        float delta_y = ray_origin_y - spheres[i].center_y;
        float delta_z = ray_origin_z - spheres[i].center_z;

        float a = ray_dx * ray_dx + ray_dy * ray_dy + ray_dz * ray_dz;
        float b = (delta_x * ray_dx + delta_y * ray_dy + delta_z * ray_dz);
        float c = (delta_x * delta_x + delta_y * delta_y + delta_z * delta_z) - (spheres[i].radius * spheres[i].radius);
        float discriminant = b * b - a * c;

        /* we can get the hit location t as: (-b - sqrtf(disciminant)) / (2.f * a) */
        if (discriminant > 0) {
            float t = (-b - sqrtf(discriminant)) / (2.f * a);
#if 0
            df_hit_record hit;
            hit.at.x = ray_origin_x + t * ray_dx;
            hit.at.y = ray_origin_y + t * ray_dy;
            hit.at.z = ray_origin_z + t * ray_dz;
            hit.ray_t = t;
            df_v3 outward_normal = {
                hit.at.x - spheres[i].center_x,
                hit.at.y - spheres[i].center_y,
                hit.at.z - spheres[i].center_z
            };
            set_face_normal(&hit, (df_v3){ray_dx, ray_dy, ray_dz}, outward_normal);
#endif

            if (t > result)
                result = t;

        }
    }
    return result;
}

static df_hit plane_test(float ray_origin_x,
                        float ray_origin_y,
                        float ray_origin_z,
                        float ray_dx,
                        float ray_dy,
                        float ray_dz,
                        const df_plane *planes,
                        unsigned int plane_count)
{
    df_hit result = {
        .ray_t = -1.f,
    };
    for (unsigned int i = 0; i < plane_count; ++i) {
        float dot = (ray_dx * planes[i].normal_x + ray_dy * planes[i].normal_y + ray_dz * planes[i].normal_z);
        if (dot > DF_EPSF32 || dot < -DF_EPSF32) {
            float delta_x = planes[i].base_x - ray_origin_x;
            float delta_y = planes[i].base_y - ray_origin_y;
            float delta_z = planes[i].base_z - ray_origin_z;

            float num = delta_x * planes[i].normal_x + delta_y * planes[i].normal_y + delta_z * planes[i].normal_z;
            float t = num / dot;

            if (t > result.ray_t) {
                result.ray_t = t;
                /* Project point on plane to image coordinate system */
                float px = ray_origin_x + t * ray_dx;
                float py = ray_origin_y + t * ray_dy;
                float pz = ray_origin_z + t * ray_dz;

                float img_p3_x = px - planes[i].img_p0_x;
                float img_p3_y = py - planes[i].img_p0_y;
                float img_p3_z = pz - planes[i].img_p0_z;

                /* FIXME(Kevin): We would need to take plane rotation into account.
                 * Alternatively, just pass w & h into the plane and let the user
                 * (i.e. higher level code) worry about that */
                float w = planes[i].img_p1_x - planes[i].img_p0_x;
                float h = planes[i].img_p1_y - planes[i].img_p0_y;

                result.img_u =
                    img_p3_x * planes[i].img_ax0_x + img_p3_y * planes[i].img_ax0_y + img_p3_z * planes[i].img_ax0_z;
                result.img_v =
                    img_p3_x * planes[i].img_ax1_x + img_p3_y * planes[i].img_ax1_y + img_p3_z * planes[i].img_ax1_z;
                result.img_u /= w;
                result.img_v /= h;
                result.image = planes[i].image;
            }
        }
        /* Line is parallel to the plane, either contained or not. Do we care? */
    }
    return result;
}

DF_API int df_trace_rays(df_trace_rays_settings settings,
                         const df_sphere *spheres,
                         unsigned int sphere_count,
                         const df_plane *planes,
                         unsigned int plane_count,
                         uint8_t **result)
{
    int image_width = settings.image_width;
    int image_height = settings.image_height;

    float aspect_ratio = (float)image_width / (float)image_height;

    float viewport_height = 2.f;
    float viewport_width = aspect_ratio * viewport_height;

    float focal_length = settings.focal_length;

    /* Simple perspective projection.
     * The lens is placed at (0, 0, 0)
     */

    float lower_left_x = -viewport_width / 2.f;
    float lower_left_y = -viewport_height / 2.f;
    float lower_left_z = -focal_length;

    uint8_t *pixels = malloc(image_width * image_height * 3);
    if (!pixels)
        return 0;

    float max_img_u = 0.f;
    float max_img_v = 0.f;

    for (int y = 0; y < image_height; ++y) {
        /* TODO(Kevin): SIMD */
        uint8_t *row = pixels + y * image_width * 3;
        for (int x = 0; x < image_width; ++x) {
            float u = (float)x / (float)(image_width - 1);
            float v = (float)y / (float)(image_height - 1);

            /* Target = Delta because origin is (0, 0, 0) */
            float target_x = lower_left_x + u * viewport_width;
            float target_y = lower_left_y + v * viewport_height;
            float target_z = lower_left_z;

            /* Raycast against all spheres */
            float sphere_hit_t = sphere_test(0, 0, 0, target_x, target_y, target_z, spheres, sphere_count);
            df_hit plane_hit = plane_test(0, 0, 0, target_x, target_y, target_z, planes, plane_count);


            #if 0
            float hits[] = {sphere_hit_t, plane_hit.ray_t};
            float hit_t = df_max_f32(hits, DF_ARRAY_COUNT(hits));

            if (hit_t >= 0) {
                df_v3 hit_p = {target_x * hit_t, target_y * hit_t, target_z * hit_t};
                hit_p.z -= -1.f;
                df_v3 normal = normalize(hit_p);

                row[x * 3 + 0] = (uint8_t)((.5f * (normal.x + 1.f)) * 255);
                row[x * 3 + 1] = (uint8_t)((.5f * (normal.y + 1.f)) * 255);
                row[x * 3 + 2] = (uint8_t)((.5f * (normal.z + 1.f)) * 255);
            }
            #endif
            if (plane_hit.ray_t >= 0) {
                float img_u = plane_hit.img_u;
                float img_v = plane_hit.img_v;
                
                // Temporary, handle arbitrary scale (esp. resulting from different orientation)
                float view_aspect = viewport_width / viewport_height;
                
                if (img_u >= 0 && img_v >= 0 && img_u <= 1.f && img_v <= 1.f) {

                    int pixelx = (int)floorf(img_u * (_image_table.images[plane_hit.image].w - 1));
                    int pixely = (int)floorf(img_v * (_image_table.images[plane_hit.image].h - 1));

                    stbi_uc *pixel = _image_table.images[plane_hit.image].pixels +
                                     4 * (pixely * _image_table.images[plane_hit.image].w + pixelx);

                    row[x * 3 + 0] = pixel[0];
                    row[x * 3 + 1] = pixel[1];
                    row[x * 3 + 2] = pixel[2];
                }
            }
            else {
                /* Gradient background color */
                float len = sqrtf(target_x * target_x + target_y * target_y + target_z * target_z);
                float t = .5f * (target_y / len + 1.f);
                float r = (1.f - t) + t * 0.5f;
                float g = (1.f - t) + t * 0.7f;
                float b = (1.f - t) + t * 1.0f;

                row[x * 3 + 0] = (uint8_t)(r * 255);
                row[x * 3 + 1] = (uint8_t)(g * 255);
                row[x * 3 + 2] = (uint8_t)(b * 255);
            }
        }
    }

    *result = pixels;
    return 1;
}