const std = @import("std");
const vec = @import("./vec.zig");
const Ray = @import("./ray.zig").Ray;
const HitRecord = @import("./hittable.zig").HitRecord;

const Vec3 = vec.Vec3;
const Color = vec.Color;
const random_unit_vector = vec.random_unit_vector;
const random_in_unit_sphere = vec.random_in_unit_sphere;
const pow = std.math.pow;

pub fn create_labertian(albedo: Vec3) Material {
    return Material{ .labertian = Labertian{ .albedo = albedo } };
}

pub fn create_metal(albedo: Vec3, fuzz: f32) Material {
    return Material{ .metal = Metal{ .albedo = albedo, .fuzz = fuzz } };
}

pub fn create_dielectric(ir: f32) Material {
    return Material{ .dielectric = Dielectric{ .ir = ir } };
}

pub const Material = union(enum) {
    labertian: Labertian,
    metal: Metal,
    dielectric: Dielectric,

    pub fn scatter(material: Material, r_in: Ray, rec: HitRecord) ?ScatterRec {
        return switch (material) {
            .labertian => |labertian| labertian.scatter(r_in, rec),
            .metal => |metal| metal.scatter(r_in, rec),
            .dielectric => |dielectric| dielectric.scatter(r_in, rec),
        };
    }
};

pub const ScatterRec = struct { attenuation: Color, scattered: Ray };

fn reflect(v: Vec3, n: Vec3) Vec3 {
    return v.sub(n.scale(v.dot(n) * 2));
}
fn refract(uv: Vec3, n: Vec3, etai_over_etat: f32) Vec3 {
    const cos_theta = @minimum(-uv.dot(n), 1.0);

    const r_out_perp = uv.add(n.scale(cos_theta)).scale(etai_over_etat);
    const r_out_parallel = n.scale(-@sqrt(@fabs(1.0 - r_out_perp.length_squared())));
    return r_out_perp.add(r_out_parallel);
}

pub const Labertian = struct {
    albedo: Vec3,
    pub fn scatter(labertian: Labertian, r_in: Ray, rec: HitRecord) ?ScatterRec {
        _ = r_in;
        var scatter_direction = rec.normal.add(random_unit_vector());
        if (scatter_direction.near_zero()) {
            scatter_direction = rec.normal;
        }
        return ScatterRec{
            .scattered = Ray.init(rec.p, scatter_direction),
            .attenuation = labertian.albedo,
        };
    }
};

pub const Metal = struct {
    albedo: Vec3,
    fuzz: f32,

    pub fn scatter(metal: Metal, r_in: Ray, rec: HitRecord) ?ScatterRec {
        const reflected = reflect(r_in.direction().unit(), rec.normal);
        const scattered = Ray.init(rec.p, reflected.add(random_in_unit_sphere().scale(metal.fuzz)));
        const attenuation = metal.albedo;
        if (scattered.direction().dot(rec.normal) > 0) {
            return ScatterRec{
                .scattered = scattered,
                .attenuation = attenuation,
            };
        } else {
            return null;
        }
    }
};

var prng = std.rand.DefaultPrng.init(0);
const rand = prng.random();
pub const Dielectric = struct {
    ir: f32, // Index of Refraction

    pub fn scatter(dielectric: Dielectric, r_in: Ray, rec: HitRecord) ?ScatterRec {
        const attenuation = Color.init(1, 1, 1);
        const refraction_ratio = if (rec.front_face) 1.0 / dielectric.ir else dielectric.ir;

        const unit_direction = r_in.direction().unit();
        const cos_theta = @minimum(-unit_direction.dot(rec.normal), 1.0);
        const sin_theta = @sqrt(1 - cos_theta * cos_theta);

        const cannot_refract = refraction_ratio * sin_theta > 1;
        const direction = if (cannot_refract or reflectance(cos_theta, refraction_ratio) > rand.float(f32))
            reflect(unit_direction, rec.normal)
        else
            refract(unit_direction, rec.normal, refraction_ratio);

        const scattered = Ray.init(rec.p, direction);
        return ScatterRec{
            .attenuation = attenuation,
            .scattered = scattered,
        };
    }
    fn reflectance(cosine: f32, ref_idx: f32) f32 {
        // Use Schlick's approximation for reflectance.
        var r0 = (1 - ref_idx) / (1 + ref_idx);
        r0 = r0 * r0;
        return r0 + (1 - r0) * pow(f32, (1 - cosine), 5);
    }
};