#include <SDL2/SDL.h>
#include <math.h>
#include <vector>

#define LIGHT_AMBIENT       0
#define LIGHT_POINT         1
#define LIGHT_DIRECTIONAL   2

SDL_Window *win;
SDL_Renderer *ren;
SDL_Texture *tex;

struct vec2 { float x, y; };
struct vec3 { 
    float x, y, z;
    vec3 operator-() { return {-x, -y, -z}; }
    vec3 operator+(vec3 a) { return {a.x+x,a.y+y,a.z+z}; }
    vec3 operator-(vec3 a) { return {x-a.x,y-a.y,z-a.z}; }
    vec3 operator*(vec3 a) { return {a.x*x,a.y*y,a.z*z}; }
    vec3 operator/(vec3 a) { return {x/a.x,y/a.y,z/a.z}; }
    vec3 operator*(float a) { return {a*x,a*y,a*z}; }
    vec3 operator/(float a) { return {x/a,y/a,z/a}; }
};
struct vec2u { uint8_t x, y; };
struct vec3u {
    uint8_t r, g, b;
    vec3u operator+(vec3u a) { return {uint8_t(r+a.r), uint8_t(g+a.g), uint8_t(b+a.b)}; }
    vec3u operator*(float a) {
        float fr = a*(float)r; float ir = 0;
        float fg = a*(float)g; float ig = 0;
        float fb = a*(float)b; float ib = 0;
        if (fr>255) { ig += (fr-255); ib += (fr-255); }
        if (fg>255) { ir += (fg-255); ib += (fg-255); }
        if (fb>255) { ir += (fb-255); ig += (fb-255); }
        return {
            uint8_t(SDL_clamp(fr+ir, 0, 255)),
            uint8_t(SDL_clamp(fg+ig, 0, 255)),
            uint8_t(SDL_clamp(fb+ib, 0, 255))
        };
    }
};

struct material_t { vec3u color; float specular; float reflective; };
struct point_t { vec3 intersection; float distance; vec3 normal; material_t material; };
struct sphere_t { vec3 center; float radius; material_t material; };
struct light_t { uint8_t type; float intensity; vec3 vector; };

struct scene_t
{
    std::vector<sphere_t> spheres;
    std::vector<light_t> lights;
};

float distance_to_viewport = 1.0f;
vec2 canvas = { 640, 640 };
vec2 viewport = { 1.0f, 1.0f };
vec3 origin = { 0, -1, 0 };
scene_t scene;

uint32_t *pixels;

void putpixel(int x, int y, vec3u color )
{
    x += canvas.x/2;
    y += canvas.y/2;
    int pos = x + y*canvas.x;
    uint32_t col = (uint32_t(color.r)<<16) + (uint32_t(color.g)<<8) + (color.b);
    pixels[pos] = col;
}

vec3 canvasToViewport(float x, float y)
{
    return { x*viewport.x/canvas.x, y*viewport.y/canvas.y, distance_to_viewport };
}

float dotProduct(vec3 a, vec3 b) { return (a.x*b.x + a.y*b.y + a.z*b.z); }
float length(vec3 a) { return sqrt(a.x*a.x+a.y*a.y+a.z*a.z); }

vec2 intersectRaySphere(vec3 origin, vec3 direction, sphere_t sphere)
{
    const float r = sphere.radius;
    const vec3 c0 = origin-sphere.center;

    float a = dotProduct(direction, direction);
    float b = 2*dotProduct(c0, direction);
    float c = dotProduct(c0, c0) - r*r;

    float discriminant = b*b - 4*a*c;
    if (discriminant < 0) return {INFINITY, INFINITY};

    float t1 = (-b + sqrt(discriminant)) / (2*a);
    float t2 = (-b - sqrt(discriminant)) / (2*a);
    return {t1, t2};
}

point_t closestIntersection(vec3 origin, vec3 direction, float t_min, float t_max)
{
    point_t point {{0,0,0}, INFINITY, {0,0,0}, {{0x87,0xce,0xfa},-1,0}};
    sphere_t *closest_sphere = nullptr;

    const float denom = dotProduct({0,1,0}, direction);
    if (denom>0.0000001f) {
        point.distance = dotProduct(-origin, {0,1,0})/denom;
        if (point.distance>=0) {
            point.intersection = origin + (direction * point.distance);
            point.normal = {0,1,0};
            int n = (int)point.intersection.x + (int)point.intersection.z;
            point.material = {{255,255,255}, 50, 0.1};
            if ((n%2)==0) point.material = {{0,0,0}, 50, 0.1};
        } else {
            point.distance = INFINITY;
        }
    }

    for (auto &sphere : scene.spheres)
    {
        vec2 t = intersectRaySphere(origin, direction, sphere);
        if (t.x >= t_min && t.x <= t_max && t.x < point.distance) {
            point.distance = t.x;
            closest_sphere = &sphere;
        }
        if (t.y >= t_min && t.y <= t_max && t.y < point.distance) {
            point.distance = t.y;
            closest_sphere = &sphere;
        }
    }

    if (closest_sphere!=nullptr)
    {
        point.intersection = origin + (direction * point.distance);
        point.normal = point.intersection - closest_sphere->center;
        point.normal = point.normal / length(point.normal);
        point.material = closest_sphere->material;
    }
    return point;
}

vec3 reflectRay(vec3 R, vec3 N)
{
    return N * 2 * dotProduct(N, R) - R;
}

float computeLighting(vec3 P, vec3 normal, vec3 camera, float s) {
    float i = 0.0f;
    float t_max;
    vec3 light_direction;
    for (auto light : scene.lights)
    {
        if (light.type == LIGHT_AMBIENT) {
            i += light.intensity;
        } else {
            if (light.type == LIGHT_POINT) {
                light_direction = light.vector - P;
                t_max = 1;
            } else {
                light_direction = light.vector;
                t_max = INFINITY;
            }

            // Shadow check
            point_t object_inbetween = closestIntersection(P, light_direction, 0.001f, t_max);
            if (object_inbetween.distance != INFINITY) continue;

            // Difuse
            float n_dot_l = dotProduct(normal, light_direction);
            if (n_dot_l > 0) {
                i += light.intensity * n_dot_l/(length(normal) * length(light_direction));
            }

            // Specular
            if (s != -1) {
                vec3 R = reflectRay(light_direction, normal);
                float r_dot_v = dotProduct(R, camera);
                if (r_dot_v > 0) {
                    i += light.intensity * pow(r_dot_v/(length(R) * length(camera)), s);
                }
            }
        }
    }
    return i;
}

vec3u traceRay(vec3 origin, vec3 direction, float t_min, float t_max, int rdepth)
{
    point_t point = closestIntersection(origin, direction, t_min, t_max);

    float lighting = point.distance==INFINITY ? 1.0f : computeLighting(point.intersection, point.normal, -direction, point.material.specular);
    vec3u local_color = point.material.color * lighting;

    float r = point.material.reflective;
    if (rdepth <= 0 || r <= 0) return local_color;

    vec3 R = reflectRay(-direction, point.normal);
    vec3u reflected_color = traceRay(point.intersection, R, 0.001, INFINITY, rdepth - 1);

    return local_color * (1 - r) + reflected_color * r;
}

void raytrace()
{
    int pitch;
    SDL_LockTexture(tex, NULL, (void**)&pixels, &pitch);

    for (int x = -canvas.x/2; x <= canvas.x/2; x++)
    {
        for (int y = -canvas.y/2; y <= canvas.y/2; y++)
        {
            vec3 D = canvasToViewport(x, y);
            vec3u color = traceRay(origin, D, 1, INFINITY, 3);
            putpixel(x, y, color);
        }
    }
    SDL_UnlockTexture(tex);
    SDL_RenderCopy(ren, tex, NULL, NULL);
    SDL_RenderPresent(ren);
}

int main(int argc, char *argv[])
{
    SDL_Init(SDL_INIT_EVENTS|SDL_INIT_VIDEO);
    win = SDL_CreateWindow("Raytracer", SDL_WINDOWPOS_UNDEFINED, SDL_WINDOWPOS_UNDEFINED, canvas.x, canvas.y, 0);
    ren = SDL_CreateRenderer(win, -1, 0);
    tex = SDL_CreateTexture(ren, SDL_PIXELFORMAT_ARGB8888, SDL_TEXTUREACCESS_STREAMING, canvas.x, canvas.y);

    scene.spheres.push_back({ {0,-1.2,5}, 1, { {0,0,0}, 100, 0.9 } });
    scene.spheres.push_back({ {-3,-3,1}, 1, { {255,2,2}, 100, 0.2 } });
    scene.spheres.push_back({ {2,-1.3,4}, 1, { {2,255,2}, 500, 0.1 } });
    scene.spheres.push_back({ {-1,0,3.5}, 1, { {2,2,255}, 10, 0.1 } });

    scene.lights.push_back({ LIGHT_AMBIENT, 0.2f });
    scene.lights.push_back({ LIGHT_DIRECTIONAL, 0.6f, {1,-8,4} });
    scene.lights.push_back({ LIGHT_POINT, 0.9f, {1,-2,3} });

    uint32_t t = SDL_GetTicks();
    raytrace();
    printf("temps: %ims\n", SDL_GetTicks()-t);

    bool should_exit = false;
    SDL_Event e;
    while (!should_exit)
    {
        while (SDL_PollEvent(&e))
        {
            if ( (e.type==SDL_KEYDOWN) && (e.key.keysym.scancode ==SDL_SCANCODE_ESCAPE) ) { should_exit = true; break; }
        }
    }
    return 0;
}