pub mod linear_rgb;
pub mod srgb;
pub use linear_rgb::*;
pub use srgb::*;

pub mod farbfeld;
pub use farbfeld::*;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Extent {
    pub x: usize,
    pub y: usize,
}

#[derive(Debug)]
pub struct Image<P> {
    pub extent: Extent,
    pub pixels: Vec<P>,
}

#[derive(Debug)]
pub enum Error {
    Invalid,
    Io(std::io::Error),
}

pub trait Format<P> {
    fn read<I>(input: &mut I) -> Result<Image<P>, Error>
    where
        I: std::io::Read;

    fn write<O>(output: &mut O, img: &Image<P>) -> std::io::Result<()>
    where
        O: std::io::Write;
}

pub trait Pixel:
    std::ops::AddAssign
    + std::ops::SubAssign
    + std::ops::MulAssign<f32>
    + std::ops::DivAssign<f32>
    + std::ops::Add<Output = Self>
    + std::ops::Sub<Output = Self>
    + std::ops::Mul<f32, Output = Self>
    + std::ops::Div<f32, Output = Self>
    + Clone
    + Copy
    + Default
    + Send
    + Sync
    + Sized
{
}

impl<P> Pixel for P where
    P: std::ops::AddAssign
        + std::ops::SubAssign
        + std::ops::MulAssign<f32>
        + std::ops::DivAssign<f32>
        + std::ops::Add<Output = Self>
        + std::ops::Sub<Output = Self>
        + std::ops::Mul<f32, Output = Self>
        + std::ops::Div<f32, Output = Self>
        + Clone
        + Copy
        + Default
        + Send
        + Sync
        + Sized
{
}

impl Extent {
    pub fn size(&self) -> usize {
        self.x * self.y
    }
}

impl<P: Pixel> Image<P> {
    pub fn load<F: Format<P>>(path: &str) -> Result<Image<P>, Error> {
        let file = std::fs::File::open(path)?;
        F::read(&mut std::io::BufReader::new(file))
    }

    pub fn save<F: Format<P>>(&self, path: &str) -> std::io::Result<()> {
        let file = std::fs::File::create(path)?;
        F::write(&mut std::io::BufWriter::new(file), self)
    }

    pub fn with_extent(extent: Extent) -> Self {
        let mut pixels = Vec::with_capacity(extent.size());
        pixels.resize(extent.size(), P::default());
        Self { extent, pixels }
    }

    pub fn for_each<F>(&mut self, f: F)
    where
        F: Fn(usize, usize, &mut P) + Send + Sync,
    {
        if let Ok(count) = std::thread::available_parallelism() {
            self.for_each_parallel(f, count.get())
        } else {
            Self::for_each_thread(
                f,
                self.pixels.as_mut_slice(),
                0,
                self.extent.x,
            )
        }
    }

    pub fn map<F, Q>(&self, f: F) -> Image<Q>
    where
        F: Fn(P) -> Q + Send + Sync,
        Q: Pixel,
    {
        let mut img = Image::with_extent(self.extent);
        img.for_each(|x, y, p| {
            *p = f(self.pixels[self.extent.x * y + x]);
        });
        img
    }

    fn for_each_parallel<F>(&mut self, f: F, n: usize)
    where
        F: Fn(usize, usize, &mut P) + Send + Sync,
    {
        let row = self.extent.x;
        let per_thread = self.extent.size() / n;
        let rem = self.extent.size() % n;

        std::thread::scope(|scope| {
            let mut rest = self.pixels.as_mut_slice();
            let mut buf: &mut [P];
            {
                let count = per_thread + rem;
                (buf, rest) = rest.split_at_mut(count);
                let fref = &f;
                scope.spawn(move || Self::for_each_thread(fref, buf, 0, row));
            }
            for i in 1..n {
                (buf, rest) = rest.split_at_mut(per_thread);
                let start = i * per_thread + rem;
                let fref = &f;
                scope.spawn(move || {
                    Self::for_each_thread(fref, buf, start, row)
                });
            }
        });
    }

    fn for_each_thread<F>(f: F, pixels: &mut [P], start: usize, row: usize)
    where
        F: Fn(usize, usize, &mut P),
    {
        for i in 0..pixels.len() {
            let x = (start + i) % row;
            let y = (start + i) / row;
            f(x, y, &mut pixels[i]);
        }
    }
}

impl From<std::io::Error> for Error {
    fn from(error: std::io::Error) -> Self {
        Error::Io(error)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn f32_image_created_with_extent_has_all_zeros() {
        let img = Image::<f32>::with_extent(Extent { x: 10, y: 10 });
        assert_eq!(img.pixels.len(), 100);
        for pixel in img.pixels {
            assert_eq!(pixel, 0.0);
        }
    }

    #[test]
    fn f32_image_for_each() {
        let mut img = Image::<f32>::with_extent(Extent { x: 10, y: 10 });
        img.for_each(|x, y, p| *p = (x + y) as f32);
        for x in 0..img.extent.x {
            for y in 0..img.extent.y {
                assert_eq!(img.pixels[y * img.extent.x + x], (x + y) as f32);
            }
        }
    }

    #[test]
    fn f32_image_map() {
        let mut img = Image::<f32>::with_extent(Extent { x: 10, y: 10 });
        for i in 0..img.extent.size() {
            img.pixels[i] = i as f32;
        }
        let res = img.map(|p| p * 2.0);
        assert_eq!(res.extent, img.extent);
        for i in 0..res.extent.size() {
            assert_eq!(res.pixels[i], 2.0 * i as f32);
        }
    }
}