use crate::mesh::Vec3;
use crate::bvh::BVH;
use crate::mesh::TriangleMesh;

/// An SDF oracle provides signed distance and gradient queries.
/// CPU implementation uses BVH; GPU would sample a 3D texture.
pub trait SdfOracle: Send + Sync {
    fn sdf(&self, p: Vec3) -> f64;

    fn gradient(&self, p: Vec3) -> Vec3 {
        let e = 1e-4;
        let dx = self.sdf(Vec3::new(p.x + e, p.y, p.z))
               - self.sdf(Vec3::new(p.x - e, p.y, p.z));
        let dy = self.sdf(Vec3::new(p.x, p.y + e, p.z))
               - self.sdf(Vec3::new(p.x, p.y - e, p.z));
        let dz = self.sdf(Vec3::new(p.x, p.y, p.z + e))
               - self.sdf(Vec3::new(p.x, p.y, p.z - e));
        let inv = 1.0 / (2.0 * e);
        Vec3::new(dx * inv, dy * inv, dz * inv).normalized()
    }

    fn project_to_surface(&self, p: Vec3, max_iter: u32) -> Option<(Vec3, Vec3)> {
        let mut pos = p;
        for _ in 0..max_iter {
            let d = self.sdf(pos);
            if d.abs() < 1e-5 {
                let n = self.gradient(pos);
                return Some((pos, n));
            }
            let g = self.gradient(pos);
            pos = pos - g * d;
        }
        None
    }
}

/// BVH-backed oracle for CPU path.
pub struct BvhOracle<'a> {
    pub mesh: &'a TriangleMesh,
    pub bvh: &'a BVH,
}

impl<'a> SdfOracle for BvhOracle<'a> {
    fn sdf(&self, p: Vec3) -> f64 {
        self.bvh.signed_distance(self.mesh, p)
    }
}

/// Voxelized SDF for GPU-friendly queries (also usable on CPU).
/// Stores a regular grid of signed distances.
pub struct VoxelOracle {
    pub data: Vec<f32>,
    pub resolution: usize,
    pub origin: Vec3,
    pub cell_size: f64,
}

impl VoxelOracle {
    pub fn from_mesh(mesh: &TriangleMesh, bvh: &BVH, resolution: usize) -> Self {
        let pad = mesh.bounds.diagonal() * 0.1;
        let origin = Vec3::new(
            mesh.bounds.min.x - pad,
            mesh.bounds.min.y - pad,
            mesh.bounds.min.z - pad,
        );
        let extent = mesh.bounds.diagonal() + 2.0 * pad;
        let cell_size = extent / resolution as f64;

        let n3 = resolution * resolution * resolution;
        let mut data = vec![0.0f32; n3];

        for iz in 0..resolution {
            for iy in 0..resolution {
                for ix in 0..resolution {
                    let p = Vec3::new(
                        origin.x + (ix as f64 + 0.5) * cell_size,
                        origin.y + (iy as f64 + 0.5) * cell_size,
                        origin.z + (iz as f64 + 0.5) * cell_size,
                    );
                    let idx = iz * resolution * resolution + iy * resolution + ix;
                    data[idx] = bvh.signed_distance(mesh, p) as f32;
                }
            }
        }

        Self { data, resolution, origin, cell_size }
    }

    fn sample(&self, p: Vec3) -> f64 {
        let r = self.resolution;
        let fx = (p.x - self.origin.x) / self.cell_size - 0.5;
        let fy = (p.y - self.origin.y) / self.cell_size - 0.5;
        let fz = (p.z - self.origin.z) / self.cell_size - 0.5;

        let ix = fx.floor() as isize;
        let iy = fy.floor() as isize;
        let iz = fz.floor() as isize;

        let tx = fx - fx.floor();
        let ty = fy - fy.floor();
        let tz = fz - fz.floor();

        let get = |x: isize, y: isize, z: isize| -> f64 {
            let cx = x.clamp(0, r as isize - 1) as usize;
            let cy = y.clamp(0, r as isize - 1) as usize;
            let cz = z.clamp(0, r as isize - 1) as usize;
            self.data[cz * r * r + cy * r + cx] as f64
        };

        // Trilinear interpolation
        let c000 = get(ix, iy, iz);
        let c100 = get(ix + 1, iy, iz);
        let c010 = get(ix, iy + 1, iz);
        let c110 = get(ix + 1, iy + 1, iz);
        let c001 = get(ix, iy, iz + 1);
        let c101 = get(ix + 1, iy, iz + 1);
        let c011 = get(ix, iy + 1, iz + 1);
        let c111 = get(ix + 1, iy + 1, iz + 1);

        let c00 = c000 * (1.0 - tx) + c100 * tx;
        let c01 = c001 * (1.0 - tx) + c101 * tx;
        let c10 = c010 * (1.0 - tx) + c110 * tx;
        let c11 = c011 * (1.0 - tx) + c111 * tx;

        let c0 = c00 * (1.0 - ty) + c10 * ty;
        let c1 = c01 * (1.0 - ty) + c11 * ty;

        c0 * (1.0 - tz) + c1 * tz
    }
}

impl SdfOracle for VoxelOracle {
    fn sdf(&self, p: Vec3) -> f64 {
        self.sample(p)
    }
}
