pub struct RgbImage {
    pub data: Vec<u8>,
    pub stride: usize,
    pub width: usize,
    pub height: usize,
}

impl RgbImage {
    pub fn get(&self, x: usize, y: usize) -> Option<(u8, u8, u8)> {
        if x >= self.width || y >= self.height {
            return None;
        }
        let i = y * self.stride + x * 3;
        let rgb = &self.data[i..];
        Some((rgb[0], rgb[1], rgb[2]))
    }
}

pub struct YCbCrImage {
    pub luma: Vec<u8>,
    pub chroma_b: Vec<u8>,
    pub chroma_r: Vec<u8>,
    pub width: usize,
    pub height: usize,
}

impl YCbCrImage {
    pub fn to_rgb(&self) -> RgbImage {
        ycbcr_to_rgb(self)
    }
}

pub fn ycbcr_to_rgb(ycbcr: &YCbCrImage) -> RgbImage {
    let mut rgb = vec![];
    for i in 0..ycbcr.luma.len() {
        let y = ycbcr.luma[i] as f64;
        let cb = ycbcr.chroma_b[i] as f64 - 128.0;
        let cr = ycbcr.chroma_r[i] as f64 - 128.0;

        let r = y + 1.40200 * cr;
        let g = y - 0.34414 * cb - 0.71414 * cr;
        let b = y + 1.77200 * cb;

        rgb.push(r.clamp(0.0, 255.0) as u8);
        rgb.push(g.clamp(0.0, 255.0) as u8);
        rgb.push(b.clamp(0.0, 255.0) as u8);
    }

    RgbImage {
        data: rgb,
        stride: 3 * ycbcr.width,
        width: ycbcr.width,
        height: ycbcr.height,
    }
}

#[cfg(test)]
mod tests {
    use crate::{rgb_to_ycbcr, RgbImageRef};

    #[test]
    fn test_rgb_to_ycbcr() {
        let rgb = [86, 86, 0];
        let rgb = RgbImageRef {
            data: &rgb,
            stride: 3,
            width: 1,
            height: 1,
        };
        let ycbcr = rgb_to_ycbcr(&rgb);
        assert_eq!(&ycbcr.luma, &[76]);
        assert_eq!(&ycbcr.chroma_b, &[85]);
        assert_eq!(&ycbcr.chroma_r, &[134]);
    }
}
