use crate::core::color::{ColorEncoding, ColorEncodingTrait, LINEAR};
use crate::core::geometry::{Bounds2f, Point2f, Point2fi, Point2i};
use crate::core::pbrt::Float;
use crate::utils::containers::Array2D;
use crate::utils::math::{f16_to_f32, lerp, square};
use core::hash;
use half::f16;
use smallvec::{SmallVec, smallvec};
use std::ops::{Deref, DerefMut};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum WrapMode {
    Black,
    Clamp,
    Repeat,
    OctahedralSphere,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct WrapMode2D {
    pub uv: [WrapMode; 2],
}

impl From<WrapMode> for WrapMode2D {
    fn from(w: WrapMode) -> Self {
        Self { uv: [w, w] }
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum PixelFormat {
    U8,
    F16,
    F32,
}

impl std::fmt::Display for PixelFormat {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            PixelFormat::U8 => write!(f, "U256"),
            PixelFormat::F16 => write!(f, "Half"),
            PixelFormat::F32 => write!(f, "Float"),
        }
    }
}

impl PixelFormat {
    pub fn is_8bit(&self) -> bool {
        matches!(self, PixelFormat::U8)
    }

    pub fn is_16bit(&self) -> bool {
        matches!(self, PixelFormat::F16)
    }

    pub fn is_32bit(&self) -> bool {
        matches!(self, PixelFormat::F32)
    }

    pub fn texel_bytes(&self) -> usize {
        match self {
            PixelFormat::U8 => 1,
            PixelFormat::F16 => 2,
            PixelFormat::F32 => 4,
        }
    }
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub enum Pixels {
    U8(*const u8),
    F16(*const u16),
    F32(*const f32),
}

#[repr(C)]
#[derive(Clone, Copy, Debug)]
pub struct ImageBase {
    pub format: PixelFormat,
    pub encoding: ColorEncoding,
    pub resolution: Point2i,
    pub n_channels: i32,
}

impl ImageBase {
    pub fn remap_pixel_coords(&self, p: &mut Point2i, wrap_mode: WrapMode2D) -> bool {
        let resolution = self.resolution;
        for i in 0..2 {
            if p[i] >= 0 && p[i] < resolution[i] {
                continue;
            }
            match wrap_mode.uv[i] {
                WrapMode::Black => return false,
                WrapMode::Clamp => p[i] = p[i].clamp(0, resolution[i] - 1),
                WrapMode::Repeat => p[i] = p[i].rem_euclid(resolution[i]),
                WrapMode::OctahedralSphere => {
                    p[i] = p[i].clamp(0, resolution[i] - 1);
                }
            }
        }
        true
    }
}

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct DeviceImage {
    pub base: ImageBase,
    pub pixels: Pixels,
}

impl DeviceImage {
    pub fn base(&self) -> ImageBase {
        self.base
    }

    pub fn resolution(&self) -> Point2i {
        self.base.resolution
    }

    pub fn is_valid(&self) -> bool {
        self.resolution().x() > 0 && self.resolution().y() > 0
    }

    pub fn format(&self) -> PixelFormat {
        self.base().format
    }

    pub fn n_channels(&self) -> i32 {
        self.base().n_channels
    }

    pub fn pixel_offset(&self, p: Point2i) -> u32 {
        let width = self.resolution().x() as u32;
        let idx = p.y() as u32 * width + p.x() as u32;
        idx * (self.n_channels() as u32)
    }

    pub fn bilerp_channel(&self, p: Point2f, c: i32) -> Float {
        self.bilerp_channel_with_wrap(p, c, WrapMode::Clamp.into())
    }

    pub fn bilerp_channel_with_wrap(&self, p: Point2f, c: i32, wrap_mode: WrapMode2D) -> Float {
        let x = p.x() * self.resolution().x() as Float - 0.5;
        let y = p.y() * self.resolution().y() as Float - 0.5;
        let xi = x.floor() as i32;
        let yi = y.floor() as i32;
        let dx = x - xi as Float;
        let dy = y - yi as Float;
        let v00 = self.get_channel_with_wrap(Point2i::new(xi, yi), c, wrap_mode);
        let v10 = self.get_channel_with_wrap(Point2i::new(xi + 1, yi), c, wrap_mode);
        let v01 = self.get_channel_with_wrap(Point2i::new(xi, yi + 1), c, wrap_mode);
        let v11 = self.get_channel_with_wrap(Point2i::new(xi + 1, yi + 1), c, wrap_mode);
        lerp(dy, lerp(dx, v00, v10), lerp(dx, v01, v11))
    }
}

pub trait ImageAccess {
    fn get_channel_with_wrap(&self, p: Point2i, c: i32, wrap_mode: WrapMode2D) -> Float;
    fn get_channel(&self, p: Point2i, c: i32) -> Float;
    fn lookup_nearest_channel_with_wrap(&self, p: Point2f, c: i32, wrap_mode: WrapMode2D) -> Float;
    fn lookup_nearest_channel(&self, p: Point2f, c: i32) -> Float;
}

impl ImageAccess for DeviceImage {
    fn get_channel_with_wrap(&self, mut p: Point2i, c: i32, wrap_mode: WrapMode2D) -> Float {
        if !self.base.remap_pixel_coords(&mut p, wrap_mode) {
            return 0.;
        }

        let offset = self.pixel_offset(p) + c as u32;
        unsafe {
            match self.pixels {
                Pixels::U8(ptr) => {
                    let raw_u8 = *ptr.add(offset as usize);
                    self.base().encoding.to_linear_scalar(raw_u8)
                }
                Pixels::F16(ptr) => {
                    let half_bits: u16 = *ptr.add(offset as usize) as u16;
                    f16_to_f32(half_bits)
                }
                Pixels::F32(ptr) => *ptr.add(offset as usize),
            }
        }
    }

    fn get_channel(&self, p: Point2i, c: i32) -> Float {
        self.get_channel_with_wrap(p, c, WrapMode::Clamp.into())
    }

    fn lookup_nearest_channel_with_wrap(&self, p: Point2f, c: i32, wrap_mode: WrapMode2D) -> Float {
        let pi = Point2i::new(
            p.x() as i32 * self.resolution().x(),
            p.y() as i32 * self.resolution().y(),
        );

        self.get_channel_with_wrap(pi, c, wrap_mode)
    }

    fn lookup_nearest_channel(&self, p: Point2f, c: i32) -> Float {
        self.lookup_nearest_channel_with_wrap(p, c, WrapMode::Clamp.into())
    }
}
