use crate::core::image::{HostImage, ImageIO};
use shared::core::color::{ColorEncoding, RGB};
use shared::core::geometry::{Point2f, Point2i, Vector2f, VectorLike};
use shared::core::image::{WrapMode, WrapMode2D};
use shared::spectra::RGBColorSpace;
use shared::utils::math::{lerp, safe_sqrt, square};
use shared::Float;
use std::hash::{Hash, Hasher};
use std::ops::{Add, Mul, Sub};
use std::path::Path;

#[cfg(feature = "cuda")]
use std::sync::OnceLock;

#[repr(C)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum FilterFunction {
    Point,
    Bilinear,
    Trilinear,
    Ewa,
}

impl std::fmt::Display for FilterFunction {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let s = match self {
            FilterFunction::Ewa => "EWA",
            FilterFunction::Trilinear => "trilinear",
            FilterFunction::Bilinear => "bilinear",
            FilterFunction::Point => "point",
        };
        write!(f, "{}", s)
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct MIPMapFilterOptions {
    pub filter: FilterFunction,
    pub max_anisotropy: Float,
}

impl Default for MIPMapFilterOptions {
    fn default() -> Self {
        Self {
            filter: FilterFunction::Ewa,
            max_anisotropy: 8.0,
        }
    }
}

impl PartialEq for MIPMapFilterOptions {
    fn eq(&self, other: &Self) -> bool {
        self.filter == other.filter
            && self.max_anisotropy.to_bits() == other.max_anisotropy.to_bits()
    }
}

impl Eq for MIPMapFilterOptions {}

impl Hash for MIPMapFilterOptions {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.filter.hash(state);
        self.max_anisotropy.to_bits().hash(state);
    }
}

pub trait MIPMapSample:
    Copy + Add<Output = Self> + Sub<Output = Self> + Mul<Float, Output = Self> + std::fmt::Debug
{
    fn zero() -> Self;
    fn sample_bilerp(image: &HostImage, st: Point2f, wrap: WrapMode2D) -> Self;
    fn sample_texel(image: &HostImage, st: Point2i, wrap: WrapMode2D) -> Self;
}

impl MIPMapSample for Float {
    fn zero() -> Self {
        0.
    }
    fn sample_bilerp(image: &HostImage, st: Point2f, wrap: WrapMode2D) -> Self {
        image.bilerp_channel_with_wrap(st, 0, wrap)
    }
    fn sample_texel(image: &HostImage, st: Point2i, wrap: WrapMode2D) -> Self {
        image.get_channel_with_wrap(st, 0, wrap)
    }
}

impl MIPMapSample for RGB {
    fn zero() -> Self {
        RGB::new(0., 0., 0.)
    }
    fn sample_bilerp(image: &HostImage, st: Point2f, wrap: WrapMode2D) -> Self {
        let nc = image.n_channels();
        if nc >= 3 {
            let r = image.bilerp_channel_with_wrap(st, 0, wrap);
            let g = image.bilerp_channel_with_wrap(st, 1, wrap);
            let b = image.bilerp_channel_with_wrap(st, 2, wrap);
            RGB::new(r, g, b)
        } else {
            let v = image.bilerp_channel_with_wrap(st, 0, wrap);
            RGB::new(v, v, v)
        }
    }
    fn sample_texel(image: &HostImage, st: Point2i, wrap: WrapMode2D) -> Self {
        let nc = image.n_channels();
        if nc >= 3 {
            let r = image.get_channel_with_wrap(st, 0, wrap);
            let g = image.get_channel_with_wrap(st, 1, wrap);
            let b = image.get_channel_with_wrap(st, 2, wrap);
            RGB::new(r, g, b)
        } else {
            let v = image.get_channel_with_wrap(st, 0, wrap);
            RGB::new(v, v, v)
        }
    }
}

#[derive(Clone, Debug)]
pub struct MIPMap {
    pub pyramid: Vec<HostImage>,
    pub color_space: Option<RGBColorSpace>,
    pub wrap_mode: WrapMode,
    pub options: MIPMapFilterOptions,
    #[cfg(feature = "cuda")]
    tex_obj: OnceLock<u64>,
}

impl MIPMap {
    pub fn new(
        image: HostImage,
        color_space: Option<RGBColorSpace>,
        wrap_mode: WrapMode,
        options: MIPMapFilterOptions,
    ) -> Self {
        let pyramid = HostImage::generate_pyramid(image, wrap_mode);
        Self {
            pyramid,
            color_space,
            wrap_mode,
            options,
            #[cfg(feature = "cuda")]
            tex_obj: OnceLock::new(),
        }
    }

    pub fn level_resolution(&self, level: usize) -> Point2i {
        self.pyramid[level].resolution()
    }

    #[inline]
    pub fn is_single_channel(&self) -> bool {
        self.pyramid[0].n_channels() == 1
    }

    pub fn levels(&self) -> usize {
        self.pyramid.len()
    }

    pub fn get_rgb_colorspace(&self) -> Option<RGBColorSpace> {
        self.color_space.clone()
    }

    pub fn get_level(&self, level: usize) -> &HostImage {
        &self.pyramid[level]
    }

    pub fn base_image(&self) -> &HostImage {
        &self.pyramid[0]
    }

    pub fn filter<T: MIPMapSample>(
        &self,
        st: Point2f,
        mut dst0: Vector2f,
        mut dst1: Vector2f,
    ) -> T {
        if self.options.filter != FilterFunction::Ewa {
            let width = 2.0
                * [
                    dst0.x().abs(),
                    dst0.y().abs(),
                    dst1.x().abs(),
                    dst1.y().abs(),
                ]
                .into_iter()
                .reduce(Float::max)
                .unwrap_or(0.0);

            let n_levels = self.levels() as Float;
            let level = n_levels - 1.0 + width.max(1e-8).log2();

            if level >= n_levels - 1.0 {
                return self.texel(self.levels() - 1, Point2i::new(0, 0));
            }

            let i_level = level.floor() as usize;
            return match self.options.filter {
                FilterFunction::Point => {
                    let resolution = self.level_resolution(i_level);
                    let sti = Point2i::new(
                        (st.x() * resolution.x() as Float - 0.5).round() as i32,
                        (st.y() * resolution.y() as Float - 0.5).round() as i32,
                    );
                    self.texel(i_level, sti)
                }
                FilterFunction::Bilinear => self.bilerp(i_level, st),
                FilterFunction::Trilinear => {
                    let v0 = self.bilerp(i_level, st);
                    let v1 = self.bilerp(i_level + 1, st);
                    let t = level - i_level as Float;
                    lerp(t, v0, v1)
                }
                FilterFunction::Ewa => unreachable!(),
            };
        }

        if dst0.norm_squared() < dst1.norm_squared() {
            std::mem::swap(&mut dst0, &mut dst1);
        }
        let longer_len = dst0.norm();
        let mut shorter_len = dst1.norm();

        if shorter_len * self.options.max_anisotropy < longer_len && shorter_len > 0.0 {
            let scale = longer_len / (shorter_len * self.options.max_anisotropy);
            dst1 *= scale;
            shorter_len *= scale;
        }

        if shorter_len == 0.0 {
            return self.bilerp(0, st);
        }

        let lod = (self.levels() as Float - 1.0 + shorter_len.log2()).max(0.0);
        let ilod = lod.floor() as usize;

        let v0 = self.ewa(ilod, st, dst0, dst1);
        let v1 = self.ewa(ilod + 1, st, dst0, dst1);
        lerp(lod - ilod as Float, v0, v1)
    }

    fn texel<T: MIPMapSample>(&self, level: usize, st: Point2i) -> T {
        if level >= self.levels() {
            panic!("MIPMap level out of bounds");
        }
        let image = &self.pyramid[level];
        let wrap_2d = WrapMode2D {
            uv: [self.wrap_mode; 2],
        };
        T::sample_texel(image, st, wrap_2d)
    }

    fn bilerp<T: MIPMapSample>(&self, level: usize, st: Point2f) -> T {
        let image = &self.pyramid[level];
        let wrap_2d = WrapMode2D {
            uv: [self.wrap_mode; 2],
        };
        T::sample_bilerp(image, st, wrap_2d)
    }

    fn ewa<T: MIPMapSample>(
        &self,
        level: usize,
        mut st: Point2f,
        mut dst0: Vector2f,
        mut dst1: Vector2f,
    ) -> T {
        if level > self.levels() {
            return self.texel(self.levels() - 1, Point2i::new(0, 0));
        }

        let level_res = self.level_resolution(level);
        st[0] = st[0] * level_res[0] as Float - 0.5;
        st[1] = st[1] * level_res[1] as Float - 0.5;
        dst0[0] *= level_res[0] as Float;
        dst0[1] *= level_res[1] as Float;
        dst1[0] *= level_res[0] as Float;
        dst1[1] *= level_res[1] as Float;

        let mut a = square(dst0[1]) + square(dst1[1]) + 1.;
        let mut b = -2. * (dst0[0] + dst0[1] + dst1[1]);
        let mut c = square(dst0[0]) + square(dst1[0]) + 1.;
        let inv_f = 1. / (a * c - square(b) * 0.25);
        a *= inv_f;
        b *= inv_f;
        c *= inv_f;

        let det = -square(b) + 4. * a * c;
        let inv_det = 1. / det;
        let u_sqrt = safe_sqrt(det * c);
        let v_sqrt = safe_sqrt(det * a);

        let s0: i32 = (st[0] - 2. * inv_det * u_sqrt).ceil() as i32;
        let s1: i32 = (st[0] + 2. * inv_det * u_sqrt).floor() as i32;
        let t0: i32 = (st[1] - 2. * inv_det * v_sqrt).ceil() as i32;
        let t1: i32 = (st[1] + 2. * inv_det * v_sqrt).floor() as i32;

        let mut sum = T::zero();
        let mut sum_wts = 0.;

        for it in t0..=t1 {
            let tt = it as Float - st[1];
            for is in s0..=s1 {
                let ss = is as Float - st[0];
                let r2 = a * square(ss) + b * ss * tt + c * square(tt);
                if r2 < 1.0 {
                    let index = (r2 * MIP_FILTER_LUT_SIZE as Float)
                        .min((MIP_FILTER_LUT_SIZE - 1) as Float)
                        as usize;
                    let weight = MIP_FILTER_LUT[index];
                    sum = sum + self.texel::<T>(level, Point2i::new(is, it)) * weight;
                    sum_wts += weight;
                }
            }
        }
        sum * (1. / sum_wts)
    }

    pub fn create_from_file(
        filename: &Path,
        options: MIPMapFilterOptions,
        wrap_mode: WrapMode,
        encoding: ColorEncoding,
    ) -> Result<MIPMap, ()> {
        let image_and_metadata = HostImage::read(filename, Some(encoding)).unwrap();
        let image = image_and_metadata.image;
        Ok(MIPMap::new(
            image,
            image_and_metadata.metadata.colorspace,
            wrap_mode,
            options,
        ))
    }

    #[cfg(feature = "cuda")]
    pub fn texture_object(&self) -> u64 {
        *self
            .tex_obj
            .get_or_init(|| create_cuda_texture(&self.pyramid, self.wrap_mode))
    }

    #[cfg(not(feature = "cuda"))]
    pub fn texture_object(&self) -> u64 {
        0
    }
}

#[cfg(feature = "cuda")]
fn create_cuda_texture(pyramid: &[HostImage], wrap_mode: WrapMode) -> u64 {
    use cuda_runtime_sys::*;

    let base = &pyramid[0];
    let (width, height) = (
        base.resolution().x() as usize,
        base.resolution().y() as usize,
    );
    let channels = base.n_channels() as usize;

    unsafe {
        let channel_desc = cudaCreateChannelDesc(
            32,
            if channels > 1 { 32 } else { 0 },
            if channels > 2 { 32 } else { 0 },
            if channels > 3 { 32 } else { 0 },
            cudaChannelFormatKind::cudaChannelFormatKindFloat,
        );

        let mut array: cudaArray_t = std::ptr::null_mut();
        cudaMallocArray(&mut array, &channel_desc, width, height, 0);

        // Get raw pixel data from the base image
        let pixel_data = base
            .as_f32_slice()
            .expect("GPU upload requires Float encoded image");
        let row_bytes = width * channels * std::mem::size_of::<f32>();

        cudaMemcpy2DToArray(
            array,
            0,
            0,
            pixel_data.as_ptr() as *const _,
            row_bytes,
            row_bytes,
            height,
            cudaMemcpyKind::cudaMemcpyHostToDevice,
        );

        let res_desc = cudaResourceDesc {
            resType: cudaResourceType::cudaResourceTypeArray,
            res: cudaResourceDesc__bindgen_ty_1 {
                array: cudaResourceDesc__bindgen_ty_1__bindgen_ty_1 { array },
            },
        };

        let address_mode = match wrap_mode {
            WrapMode::Repeat => cudaTextureAddressMode::cudaAddressModeWrap,
            WrapMode::Clamp => cudaTextureAddressMode::cudaAddressModeClamp,
            WrapMode::Black => cudaTextureAddressMode::cudaAddressModeBorder,
            WrapMode::OctahedralSphere => cudaTextureAddressMode::cudaAddressModeBorder,
        };

        let tex_desc = cudaTextureDesc {
            addressMode: [address_mode; 3],
            filterMode: cudaTextureFilterMode::cudaFilterModeLinear,
            readMode: cudaTextureReadMode::cudaReadModeElementType,
            normalizedCoords: 1,
            ..std::mem::zeroed()
        };

        let mut tex_obj: cudaTextureObject_t = 0;
        cudaCreateTextureObject(&mut tex_obj, &res_desc, &tex_desc, std::ptr::null());
        tex_obj
    }
}

static MIP_FILTER_LUT_SIZE: usize = 128;

static MIP_FILTER_LUT: [Float; MIP_FILTER_LUT_SIZE] = [
    0.864664733,
    0.849040031,
    0.83365953,
    0.818519294,
    0.80361563,
    0.788944781,
    0.774503231,
    0.760287285,
    0.746293485,
    0.732518315,
    0.718958378,
    0.705610275,
    0.692470789,
    0.679536581,
    0.666804492,
    0.654271305,
    0.641933978,
    0.629789352,
    0.617834508,
    0.606066525,
    0.594482362,
    0.583079159,
    0.571854174,
    0.560804546,
    0.549927592,
    0.539220572,
    0.528680861,
    0.518305838,
    0.50809288,
    0.498039544,
    0.488143265,
    0.478401601,
    0.468812168,
    0.45937258,
    0.450080454,
    0.440933526,
    0.431929469,
    0.423066139,
    0.414341331,
    0.405752778,
    0.397298455,
    0.388976216,
    0.380784035,
    0.372719884,
    0.364781618,
    0.356967449,
    0.34927541,
    0.341703475,
    0.334249914,
    0.32691282,
    0.319690347,
    0.312580705,
    0.305582166,
    0.298692942,
    0.291911423,
    0.285235822,
    0.278664529,
    0.272195935,
    0.265828371,
    0.259560347,
    0.253390193,
    0.247316495,
    0.241337672,
    0.235452279,
    0.229658857,
    0.223955944,
    0.21834214,
    0.212816045,
    0.207376286,
    0.202021524,
    0.196750447,
    0.191561714,
    0.186454013,
    0.181426153,
    0.176476851,
    0.171604887,
    0.166809067,
    0.162088141,
    0.157441005,
    0.152866468,
    0.148363426,
    0.143930718,
    0.139567271,
    0.135272011,
    0.131043866,
    0.126881793,
    0.122784719,
    0.11875169,
    0.114781633,
    0.11087364,
    0.107026696,
    0.103239879,
    0.0995122194,
    0.0958427936,
    0.0922307223,
    0.0886750817,
    0.0851749927,
    0.0817295909,
    0.0783380121,
    0.0749994367,
    0.0717130303,
    0.0684779733,
    0.0652934611,
    0.0621587038,
    0.0590728968,
    0.0560353249,
    0.0530452281,
    0.0501018465,
    0.0472044498,
    0.0443523228,
    0.0415447652,
    0.0387810767,
    0.0360605568,
    0.0333825648,
    0.0307464004,
    0.0281514227,
    0.0255970061,
    0.0230824798,
    0.0206072628,
    0.0181707144,
    0.0157722086,
    0.013411209,
    0.0110870898,
    0.0087992847,
    0.0065472275,
    0.00433036685,
    0.0021481365,
    0.,
];