foreign fn ppm_start(width: i32, height: i32);
foreign fn ppm_pixel(red: i32, green: i32, blue: i32);
foreign fn ppm_finish();

foreign fn sqrtf(num: f32) -> f32;
foreign fn tanf(num: f32) -> f32;
foreign fn randf() -> f32;

struct Vec3 {
    x: f32,
    y: f32,
    z: f32
}

struct Ray {
    origin: Vec3,
    dir: Vec3
}

// Added material type (0=Diffuse/Normal, 1=Glass, 2=Metal) and Refractive Index
struct Sphere {
    center: Vec3,
    radius: f32,
    color: Vec3,
    mat_type: i32,
    ref_idx: f32,
    fuzz: f32
}

struct Hit {
    hit: bool,
    t: f32,
    normal: Vec3,
    color: Vec3,
    mat_type: i32,
    ref_idx: f32,
    fuzz: f32
}

// Wrapper for refraction results
struct RefractResult {
    success: bool,
    dir: Vec3
}

fn vec3_add(a: Vec3, b: Vec3) -> Vec3 {
    return Vec3 { x: a.x + b.x, y: a.y + b.y, z: a.z + b.z };
}

fn vec3_sub(a: Vec3, b: Vec3) -> Vec3 {
    return Vec3 { x: a.x - b.x, y: a.y - b.y, z: a.z - b.z };
}

fn vec3_mul(a: Vec3, t: f32) -> Vec3 {
    return Vec3 { x: a.x * t, y: a.y * t, z: a.z * t };
}

fn vec3_dot(a: Vec3, b: Vec3) -> f32 {
    return a.x * b.x + a.y * b.y + a.z * b.z;
}

fn vec3_normalize(v: Vec3) -> Vec3 {
    let len = sqrtf(vec3_dot(v, v));
    if len == 0.0 {
        return Vec3 { x: 0.0, y: 0.0, z: 0.0 };
    }
    return vec3_mul(v, 1.0 / len);
}

// V - 2 * dot(V, N) * N
fn vec3_reflect(v: Vec3, n: Vec3) -> Vec3 {
    return vec3_sub(v, vec3_mul(n, 2.0 * vec3_dot(v, n)));
}

// Snell's Law for calculating the refracted ray direction
fn vec3_refract(v: Vec3, n: Vec3, ni_over_nt: f32) -> RefractResult {
    let uv = vec3_normalize(v);
    let dt = vec3_dot(uv, n);
    let discriminant = 1.0 - ni_over_nt * ni_over_nt * (1.0 - dt * dt);

    if discriminant > 0.0 {
        let temp1 = vec3_mul(vec3_sub(uv, vec3_mul(n, dt)), ni_over_nt);
        let temp2 = vec3_mul(n, sqrtf(discriminant));
        return RefractResult {
            success: true,
            dir: vec3_sub(temp1, temp2)
        };
    }

    // Total internal reflection occurred
    return RefractResult {
        success: false,
        dir: Vec3 { x: 0.0, y: 0.0, z: 0.0 }
    };
}

// The Cross Product generates a perpendicular vector to a and b
fn vec3_cross(a: Vec3, b: Vec3) -> Vec3 {
    return Vec3 {
        x: a.y * b.z - a.z * b.y,
        y: a.z * b.x - a.x * b.z,
        z: a.x * b.y - a.y * b.x
    };
}

// Generates a random point inside a 2D unit disk (used for depth of field blur)
fn random_in_unit_disk() -> Vec3 {
    let p = Vec3 { x: 1.0, y: 1.0, z: 0.0 };
    let found = false;
    while !found {
        p = Vec3 {
            x: randf() * 2.0 - 1.0,
            y: randf() * 2.0 - 1.0,
            z: 0.0
        };
        // Check if the point is inside the radius of 1.0
        if vec3_dot(p, p) < 1.0 {
            found = true;
        }
    }
    return p;
}

fn intersect_sphere(ray: *Ray, sphere: *Sphere, t_min: f32, t_max: f32) -> Hit {
    let oc = vec3_sub((*sphere).center, (*ray).origin);
    let a = vec3_dot((*ray).dir, (*ray).dir);
    let h = vec3_dot((*ray).dir, oc);
    let c = vec3_dot(oc, oc) - (*sphere).radius * (*sphere).radius;
    let discriminant = h * h - a * c;

    let miss = Hit {
        hit: false, t: 0.0, normal: Vec3 { x: 0.0, y: 0.0, z: 0.0 },
        color: Vec3 { x: 0.0, y: 0.0, z: 0.0 }, mat_type: 0, ref_idx: 0.0,
        fuzz: 0.0 // Ensure this is present
    };

    if discriminant < 0.0 { return miss; }

    let sqrtd = sqrtf(discriminant);
    let root = (h - sqrtd) / a;

    if root <= t_min {
        root = (h + sqrtd) / a;
        if root <= t_min { return miss; }
    }
    if root >= t_max { return miss; }

    let hit_point = vec3_add((*ray).origin, vec3_mul((*ray).dir, root));
    let normal = vec3_mul(vec3_sub(hit_point, (*sphere).center), 1.0 / (*sphere).radius);

    return Hit {
        hit: true, t: root, normal: normal,
        color: (*sphere).color, mat_type: (*sphere).mat_type,
        ref_idx: (*sphere).ref_idx,
        fuzz: (*sphere).fuzz // Copy fuzz from the sphere
    };
}

fn schlick(cosine: f32, ref_idx: f32) -> f32 {
    let r0 = (1.0 - ref_idx) / (1.0 + ref_idx);
    let r0_sq = r0 * r0;
    let comp = 1.0 - cosine;
    // Approximating (1 - cosine)^5
    return r0_sq + (1.0 - r0_sq) * comp * comp * comp * comp * comp;
}

fn random_unit_vector() -> Vec3 {
    // Generate a random vector in a cube
    let v = Vec3 {
        x: randf() * 2.0 - 1.0,
        y: randf() * 2.0 - 1.0,
        z: randf() * 2.0 - 1.0
    };
    // Normalize it to project it onto the unit sphere
    return vec3_normalize(v);
}

fn ray_color(ray: *Ray, spheres: *[Sphere; 500], sphere_count: i32, depth: i32) -> Vec3 {
    if depth <= 0 {
        return Vec3 { x: 0.0, y: 0.0, z: 0.0 };
    }

    let closest_so_far = 100000.0; // Increased to ensure far spheres render
    let hit_anything = false;
    let final_hit = Hit {
        hit: false, t: 0.0, normal: Vec3 { x: 0.0, y: 0.0, z: 0.0 },
        color: Vec3 { x: 0.0, y: 0.0, z: 0.0 }, mat_type: 0, ref_idx: 0.0, fuzz: 0.0
    };

    // Iterate up to the active sphere_count, not a hardcoded 4
    let i = 0;
    while i < sphere_count {
        let hit = intersect_sphere(ray, &(*spheres)[i], 0.001, closest_so_far);
        if hit.hit {
            hit_anything = true;
            closest_so_far = hit.t;
            final_hit = hit;
        }
        i = i + 1;
    }

    if hit_anything {
        let hit_point = vec3_add((*ray).origin, vec3_mul((*ray).dir, final_hit.t));
        let unit_dir = vec3_normalize((*ray).dir);

        if final_hit.mat_type == 1 { // Glass
            let outward_normal: Vec3;
            let ni_over_nt: f32;
            let dir_dot_normal = vec3_dot(unit_dir, final_hit.normal);

            if dir_dot_normal > 0.0 {
                outward_normal = vec3_mul(final_hit.normal, -1.0);
                ni_over_nt = final_hit.ref_idx;
            } else {
                outward_normal = final_hit.normal;
                ni_over_nt = 1.0 / final_hit.ref_idx;
            }

            let cos_theta = vec3_dot(vec3_mul(unit_dir, -1.0), outward_normal);
            let refract_result = vec3_refract((*ray).dir, outward_normal, ni_over_nt);

            let reflect_prob = 1.0;
            if refract_result.success {
                reflect_prob = schlick(cos_theta, final_hit.ref_idx);
            }

            let scattered_ray: Ray;
            if randf() < reflect_prob {
                // Use normalized unit_dir to prevent self-intersection
                let reflected = vec3_reflect(unit_dir, final_hit.normal);
                scattered_ray = Ray { origin: hit_point, dir: reflected };
            } else {
                scattered_ray = Ray { origin: hit_point, dir: refract_result.dir };
            }

            let bounce_color = ray_color(&scattered_ray, spheres, sphere_count, depth - 1);
            return Vec3 {
                x: final_hit.color.x * bounce_color.x,
                y: final_hit.color.y * bounce_color.y,
                z: final_hit.color.z * bounce_color.z
            };
        }

        if final_hit.mat_type == 2 { // Metal
            let reflected = vec3_reflect(unit_dir, final_hit.normal);

            // Apply fuzz by randomly perturbing the reflected vector
            let fuzzed_dir = vec3_add(reflected, vec3_mul(random_unit_vector(), final_hit.fuzz));
            let scattered_ray = Ray { origin: hit_point, dir: vec3_normalize(fuzzed_dir) };

            // Only scatter if the fuzzed ray doesn't bounce inward
            if vec3_dot(scattered_ray.dir, final_hit.normal) > 0.0 {
                let bounce_color = ray_color(&scattered_ray, spheres, sphere_count, depth - 1);
                return Vec3 {
                    x: final_hit.color.x * bounce_color.x,
                    y: final_hit.color.y * bounce_color.y,
                    z: final_hit.color.z * bounce_color.z
                };
            }
            return Vec3 { x: 0.0, y: 0.0, z: 0.0 };
        }

        // True Lambertian Diffuse
        let scatter_dir = vec3_add(final_hit.normal, random_unit_vector());

        // Catch degenerate scatter direction (if random vector exactly opposes normal)
        if vec3_dot(scatter_dir, scatter_dir) < 0.001 {
            scatter_dir = final_hit.normal;
        }

        // Ensure the outbound ray is perfectly normalized
        let scattered_ray = Ray { origin: hit_point, dir: vec3_normalize(scatter_dir) };
        let bounce_color = ray_color(&scattered_ray, spheres, sphere_count, depth - 1);

        return Vec3 {
            x: final_hit.color.x * bounce_color.x,
            y: final_hit.color.y * bounce_color.y,
            z: final_hit.color.z * bounce_color.z
        };
    }

    // Skybox gradient
    let unit_dir = vec3_normalize((*ray).dir);
    let t = 0.5 * (unit_dir.y + 1.0);
    return vec3_add(
        vec3_mul(Vec3 { x: 1.0, y: 1.0, z: 1.0 }, 1.0 - t),
        vec3_mul(Vec3 { x: 0.5, y: 0.7, z: 1.0 }, t)
    );
}

fn main() -> i32 {
    let aspect_ratio = 16.0 / 9.0;
    let image_width = 1920;
    let image_height = ((image_width as f32) / aspect_ratio) as i32;

    let samples_per_pixel = 10; // Increase this for higher quality
    let max_depth = 32;

    // Allocate our static maximum array and track how many we've placed
    let spheres: [Sphere; 500];
    let sphere_count = 0;

    // Ground sphere
    spheres[sphere_count] = Sphere {
        center: Vec3 { x: 0.0, y: -1000.0, z: 0.0 }, radius: 1000.0,
        color: Vec3 { x: 0.5, y: 0.5, z: 0.5 }, mat_type: 0, ref_idx: 0.0, fuzz: 0.0
    };
    sphere_count = sphere_count + 1;

    // Generate random tiny spheres
    let a = -11;
    while a < 11 {
        let b = -11;
        while b < 11 {
            let choose_mat = randf();
            let center = Vec3 {
                x: (a as f32) + 0.9 * randf(),
                y: 0.2,
                z: (b as f32) + 0.9 * randf()
            };

            // Avoid spawning inside the giant main spheres
            let center_offset = vec3_sub(center, Vec3 { x: 4.0, y: 0.2, z: 0.0 });
            let dist_sq = vec3_dot(center_offset, center_offset);

            if dist_sq > (0.9 * 0.9) {
                if choose_mat < 0.8 {
                    // Diffuse (albedo = color::random() * color::random())
                    spheres[sphere_count] = Sphere {
                        center: center, radius: 0.2,
                        color: Vec3 { x: randf() * randf(), y: randf() * randf(), z: randf() * randf() },
                        mat_type: 0, ref_idx: 0.0, fuzz: 0.0
                    };
                } else if choose_mat < 0.95 {
                    // Metal (albedo = random(0.5, 1), fuzz = random(0, 0.5))
                    spheres[sphere_count] = Sphere {
                        center: center, radius: 0.2,
                        color: Vec3 { x: 0.5 + 0.5 * randf(), y: 0.5 + 0.5 * randf(), z: 0.5 + 0.5 * randf() },
                        mat_type: 2, ref_idx: 0.0, fuzz: randf() * 0.5
                    };
                } else {
                    // Glass
                    spheres[sphere_count] = Sphere {
                        center: center, radius: 0.2,
                        color: Vec3 { x: 1.0, y: 1.0, z: 1.0 }, // Clear
                        mat_type: 1, ref_idx: 1.5, fuzz: 0.0
                    };
                }
                sphere_count = sphere_count + 1;
            }
            b = b + 1;
        }
        a = a + 1;
    }

    // Material 1: Giant Central Glass Sphere
    spheres[sphere_count] = Sphere {
        center: Vec3 { x: 0.0, y: 1.0, z: 0.0 }, radius: 1.0,
        color: Vec3 { x: 1.0, y: 1.0, z: 1.0 }, mat_type: 1, ref_idx: 1.5, fuzz: 0.0
    };
    sphere_count = sphere_count + 1;

    // Material 2: Giant Left Diffuse Sphere
    spheres[sphere_count] = Sphere {
        center: Vec3 { x: -4.0, y: 1.0, z: 0.0 }, radius: 1.0,
        color: Vec3 { x: 0.4, y: 0.2, z: 0.1 }, mat_type: 0, ref_idx: 0.0, fuzz: 0.0
    };
    sphere_count = sphere_count + 1;

    // Material 3: Giant Right Metal Sphere
    spheres[sphere_count] = Sphere {
        center: Vec3 { x: 4.0, y: 1.0, z: 0.0 }, radius: 1.0,
        color: Vec3 { x: 0.7, y: 0.6, z: 0.5 }, mat_type: 2, ref_idx: 0.0, fuzz: 0.0
    };
    sphere_count = sphere_count + 1;

    ppm_start(image_width, image_height);

    // --- Camera Settings ---
    let vfov = 20.0;
    let lookfrom = Vec3 { x: 13.0, y: 2.0, z: 3.0 };
    let lookat = Vec3 { x: 0.0, y: 0.0, z: 0.0 };
    let vup = Vec3 { x: 0.0, y: 1.0, z: 0.0 };
    let defocus_angle = 0.6;
    let focus_dist = 10.0;

    let pi = 3.14159265;
    let theta = vfov * pi / 180.0;
    let h = tanf(theta / 2.0);

    let viewport_height = 2.0 * h * focus_dist;
    let viewport_width = viewport_height * ((image_width as f32) / (image_height as f32));

    // Calculate the u,v,w unit basis vectors for the camera coordinate frame.
    let w = vec3_normalize(vec3_sub(lookfrom, lookat));
    let u = vec3_normalize(vec3_cross(vup, w));
    let v = vec3_cross(w, u);

    // Calculate the vectors across the horizontal and down the vertical viewport edges.
    let viewport_u = vec3_mul(u, viewport_width);
    let viewport_v = vec3_mul(v, -viewport_height); // Negative to point down the image

    // Calculate the horizontal and vertical delta vectors from pixel to pixel.
    let pixel_delta_u = vec3_mul(viewport_u, 1.0 / (image_width as f32));
    let pixel_delta_v = vec3_mul(viewport_v, 1.0 / (image_height as f32));

    // Calculate the location of the upper left pixel.
    let half_u = vec3_mul(viewport_u, 0.5);
    let half_v = vec3_mul(viewport_v, 0.5);
    let focus_w = vec3_mul(w, focus_dist);
    let viewport_upper_left = vec3_sub(vec3_sub(vec3_sub(lookfrom, focus_w), half_u), half_v);

    let pixel00_loc = vec3_add(viewport_upper_left, vec3_mul(vec3_add(pixel_delta_u, pixel_delta_v), 0.5));

    // Calculate the camera defocus disk basis vectors.
    let defocus_radius = focus_dist * tanf((defocus_angle / 2.0) * pi / 180.0);
    let defocus_disk_u = vec3_mul(u, defocus_radius);
    let defocus_disk_v = vec3_mul(v, defocus_radius);

    // --- Render ---
    let j = 0;
    while j < image_height {
        let i = 0;
        while i < image_width {
            let pixel_color = Vec3 { x: 0.0, y: 0.0, z: 0.0 };

            let s = 0;
            while s < samples_per_pixel {
                let px = -0.5 + randf();
                let py = -0.5 + randf();

                let pixel_center = vec3_add(pixel00_loc, vec3_add(vec3_mul(pixel_delta_u, i as f32), vec3_mul(pixel_delta_v, j as f32)));
                let sample_offset = vec3_add(vec3_mul(pixel_delta_u, px), vec3_mul(pixel_delta_v, py));
                let sample_loc = vec3_add(pixel_center, sample_offset);

                // Apply Depth of Field logic
                let ray_origin = lookfrom;
                if defocus_angle > 0.0 {
                    let p = random_in_unit_disk();
                    let defocus_offset = vec3_add(vec3_mul(defocus_disk_u, p.x), vec3_mul(defocus_disk_v, p.y));
                    ray_origin = vec3_add(lookfrom, defocus_offset);
                }

                let ray_dir = vec3_sub(sample_loc, ray_origin);
                let ray = Ray { origin: ray_origin, dir: ray_dir };

                let sample_color = ray_color(&ray, &spheres, sphere_count, max_depth);
                pixel_color = vec3_add(pixel_color, sample_color);

                s = s + 1;
            }

            let scale = 1.0 / (samples_per_pixel as f32);
            pixel_color = vec3_mul(pixel_color, scale);

            pixel_color.x = sqrtf(pixel_color.x);
            pixel_color.y = sqrtf(pixel_color.y);
            pixel_color.z = sqrtf(pixel_color.z);

            let ir = (255.999 * pixel_color.x) as i32;
            let ig = (255.999 * pixel_color.y) as i32;
            let ib = (255.999 * pixel_color.z) as i32;

            ppm_pixel(ir, ig, ib);

            i = i + 1;
        }
        j = j + 1;
    }

    ppm_finish();

    return 0;
}