const std = @import("std");
const camera = @import("./camera.zig");
const color = @import("./color.zig");
const ray = @import("./ray.zig");
const vec = @import("./vec.zig");
const sphere = @import("./sphere.zig");
const hittable = @import("./hittable.zig");
const hittable_list = @import("./hittable_list.zig");
const material = @import("./material.zig");

const Thread = std.Thread;
const Queue = std.atomic.Queue;
const Camera = camera.Camera;
const Color = vec.Color;
const Hittable = hittable.Hittable;
const HittableList = hittable_list.HittableList;
const Labertian = material.Labertian;
const Metal = material.Metal;
const Dielectric = material.Dielectric;
const Material = material.Material;
const Point3 = vec.Point3;
const Ray = ray.Ray;
const Sphere = sphere.Sphere;
const Vec3 = vec.Vec3;
const info = std.log.info;
const random_in_hemisphere = vec.random_in_hemisphere;
const write_color = color.write_color;
const math = std.math;
const pi = math.pi;
const cos = math.cos;
const create_labertian = material.create_labertian;
const create_metal = material.create_metal;
const create_dielectric = material.create_dielectric;

fn ray_color(r: Ray, world: Hittable, depth: isize) Color {
    // If we've exceeded the ray bounce limit, no more light is gathered.
    if (depth <= 0) return Color.init(0, 0, 0);

    if (world.hit(r, 0.0001, std.math.inf(f32))) |hit_rec| {
        if (hit_rec.material.scatter(r, hit_rec)) |scatter_rec| {
            return scatter_rec.attenuation.mul(ray_color(scatter_rec.scattered, world, depth - 1));
        }
        return Color.init(0, 0, 0);
    }
    var unit_direction = r.direction().unit();
    const t = 0.5 * (unit_direction.y() + 1);
    return Color.init(1, 1, 1).scale(1 - t).add(Color.init(0.5, 0.7, 1.0).scale(t));
}

const Task = struct {
    cam: Camera,
    samples_per_pixel: usize,
    max_depth: isize,
    row_idx: usize,
    image_height: isize,
    image_width: isize,
    row_pixels: []Color,
    world: Hittable,
};

fn render_row(q: *Queue(Task)) void {
    var prng = std.rand.DefaultPrng.init(0);
    const rand = prng.random();
    while (true) {
        if (q.get()) |node| {
            const task = node.data;
            info("Rendering row: {}", .{task.row_idx});
            const j = task.row_idx;
            var i: usize = 0;
            while (i < task.image_width) : (i += 1) {
                var pixel_color = Color.init(0, 0, 0);
                var s: isize = 0;
                while (s < task.samples_per_pixel) : (s += 1) {
                    const u = (@intToFloat(f32, i) + rand.float(f32)) / @intToFloat(f32, task.image_width - 1);
                    const v = (@intToFloat(f32, j) + rand.float(f32)) / @intToFloat(f32, task.image_height - 1);
                    const r = task.cam.get_ray(u, v);
                    pixel_color = pixel_color.add(ray_color(r, task.world, task.max_depth));
                }
                task.row_pixels[i] = pixel_color;
            }
        } else {
            return;
        }
    }
}
fn random_scene() anyerror!Hittable {
    var prng = std.rand.DefaultPrng.init(0);
    const rand = prng.random();
    var world = HittableList.init();
    const ground_material = create_labertian(Color.init(0.5, 0.5, 0.5));
    try world.add(Hittable{ .sphere = Sphere.init(Point3.init(0, -1000, 0), 1000, ground_material) });

    var a: isize = -11;
    while (a < 11) : (a += 1) {
        var b: isize = -11;
        while (b < 11) : (b += 1) {
            const choose_mat = rand.float(f32);
            const center = Point3.init(@intToFloat(f32, a) + 0.9 * rand.float(f32), 0.2, @intToFloat(f32, b) + 0.9 * rand.float(f32));
            if (center.sub(Point3.init(4, 0.2, 0)).length() > 0.9) {
                if (choose_mat < 0.8) {
                    const albedo = Color.random_min_max(0, 1).mul(Color.random_min_max(0, 1));
                    const sphere_material = create_labertian(albedo);

                    try world.add(Hittable{ .sphere = Sphere.init(center, 0.2, sphere_material) });
                } else if (choose_mat < 0.95) {
                    const albedo = Color.init(
                        0.5 + rand.float(f32) / 2.0,
                        0.5 + rand.float(f32) / 2.0,
                        0.5 + rand.float(f32) / 2.0,
                    );
                    const fuzz = rand.float(f32) / 2.0;
                    const sphere_material = create_metal(albedo, fuzz);

                    try world.add(Hittable{ .sphere = Sphere.init(center, 0.2, sphere_material) });
                } else {
                    const sphere_material = create_dielectric(1.5);

                    try world.add(Hittable{ .sphere = Sphere.init(center, 0.2, sphere_material) });
                }
            }
        }
    }

    const material1 = create_dielectric(1.5);
    try world.add(Hittable{ .sphere = Sphere.init(Point3.init(0, 1, 0), 1.0, material1) });

    const material2 = create_labertian(Color.init(0.4, 0.2, 0.1));
    try world.add(Hittable{ .sphere = Sphere.init(Point3.init(-4, 1, 0), 1.0, material2) });

    const material3 = create_metal(Color.init(0.7, 0.6, 0.5), 0.0);
    try world.add(Hittable{ .sphere = Sphere.init(Point3.init(4, 1, 0), 1.0, material3) });

    return Hittable{ .hittable_list = world };
}

pub fn main() anyerror!void {
    // Image
    const aspect_ratio: f32 = 3.0 / 2.0;
    const image_width = 1200;
    const image_height = @floatToInt(isize, @intToFloat(f32, image_width) / aspect_ratio);
    const samples_per_pixel = 500;
    const max_depth = 50;

    // World
    const world = try random_scene();

    // Camera
    const lookfrom = Point3.init(13, 2, 3);
    const lookat = Point3.init(0, 0, 0);
    const vup = Vec3.init(0, 1, 0);
    const dist_to_focus = 10;
    const aperature = 0.1;

    const cam = Camera.init(lookfrom, lookat, vup, 20, aspect_ratio, aperature, dist_to_focus);

    // Render
    const Node = Queue(Task).Node;
    const allocator = std.heap.page_allocator;
    const cpus = try Thread.getCpuCount();
    var pixels = try allocator.create([image_width * image_height]Color);
    var q = Queue(Task).init();
    var threads: std.ArrayList(std.Thread) = std.ArrayList(std.Thread).init(allocator);
    var row: usize = 0;
    while (row < image_height) : (row += 1) {
        const node = try allocator.create(Node);
        const start: usize = row * image_width;
        const end: usize = (row + 1) * image_width;
        node.* = .{
            .prev = undefined,
            .next = undefined,
            .data = Task{
                .cam = cam,
                .samples_per_pixel = samples_per_pixel,
                .max_depth = max_depth,
                .row_idx = row,
                .image_width = image_width,
                .image_height = image_height,
                .row_pixels = pixels[start..end],
                .world = world,
            },
        };
        q.put(node);
    }

    var t: usize = 0;
    while (t < cpus) : (t += 1) {
        try threads.append(try Thread.spawn(.{}, render_row, .{&q}));
    }

    // Wait for all threads to finish.
    for (threads.items) |thread| {
        Thread.join(thread);
    }

    const stdout = std.io.getStdOut();

    info("Writing image {d}x{d}", .{ image_width, image_height });
    try stdout.writer().print("P3\n{d} {d}\n255\n", .{ image_width, image_height });

    var j: usize = 0;
    while (j < image_height) : (j += 1) {
        var i: usize = 0;
        while (i < image_width) : (i += 1) {
            // Flip image upside down.
            const j_idx = @as(usize, image_height) - j - 1;
            const pixel_color = pixels[i + j_idx * image_width];
            try write_color(pixel_color, samples_per_pixel);
        }
    }
    info("Done", .{});
}

test {
    // Run tests in imported source.
    std.testing.refAllDecls(@This());
}
