use crate::BLOCK_SIZE;
use framework::BitstreamFilter;
use glam::{IVec2, ivec2};

pub struct Enc {
    res: IVec2,
    last: Frame,
}
impl BitstreamFilter for Enc {
    const INPUT_CODEC_ID: &str = "V_UNCOMPRESSED";
    const OUTPUT_CODEC_ID: &str = "V_VCETEST2";

    fn new(width: u32, height: u32) -> Self {
        let res = ivec2(width as i32, height as i32);
        Self {
            res,
            last: Frame::new(res),
        }
    }

    fn process_block(&mut self, a: Vec<u8>) -> Vec<u8> {
        let mut frame = Frame {
            data: a,
            res: self.res,
        };

        for by in 0..frame.res.y / BLOCK_SIZE {
            for bx in 0..frame.res.x / BLOCK_SIZE {
                let boff = ivec2(bx * BLOCK_SIZE, by * BLOCK_SIZE);

                let mut best_d = Frame::compare_block(&frame, &self.last, boff, boff);
                let mut best_off = ivec2(0, 0);

                for granularity in [8, 4, 2, 1] {
                    for dir in [
                        ivec2(1, 0),
                        ivec2(1, 1),
                        ivec2(0, 1),
                        ivec2(-1, 1),
                        ivec2(-1, 0),
                        ivec2(-1, -1),
                        ivec2(0, -1),
                        ivec2(1, -1),
                    ] {
                        let roff = boff + dir * granularity;
                        if roff.x >= 0
                            && roff.y >= 0
                            && roff.x + BLOCK_SIZE < frame.res.x
                            && roff.y + BLOCK_SIZE < frame.res.y
                        {
                            let d = Frame::compare_block(&frame, &self.last, boff, roff);
                            if d < best_d {
                                best_d = d;
                                best_off = roff;
                            }
                        }
                    }
                }

                if best_d < 8 * 8 * 20 {
                    Frame::copy_block(&self.last, &mut frame, best_off, boff);
                }
            }
        }

        self.last = frame.clone();
        frame.data
    }
}

#[derive(Clone)]
struct Frame {
    res: IVec2,
    data: Vec<u8>,
}

impl Frame {
    fn new(res: IVec2) -> Self {
        Frame {
            res,
            data: vec![127; (res.x * res.y + res.x * res.y / 2) as usize],
        }
    }
    fn copy_block(aframe: &Frame, bframe: &mut Frame, aoff: IVec2, boff: IVec2) {
        assert_eq!(aframe.res, bframe.res);
        let res = aframe.res;

        // Luma
        for y in 0..BLOCK_SIZE {
            let ay_off = aoff.x + (y + aoff.y) * res.x;
            let by_off = boff.x + (y + boff.y) * res.x;

            bframe.data[by_off as usize..(by_off + BLOCK_SIZE) as usize]
                .copy_from_slice(&aframe.data[ay_off as usize..(ay_off + BLOCK_SIZE) as usize]);
        }

        // Chroma
        let uvplane_off = res.x * res.y;

        for y in 0..BLOCK_SIZE / 2 {
            let ay_off = uvplane_off + (aoff.x & !1) + (y + aoff.y / 2) * res.x;
            let by_off = uvplane_off + (boff.x & !1) + (y + boff.y / 2) * res.x;

            bframe.data[by_off as usize..(by_off + BLOCK_SIZE) as usize]
                .copy_from_slice(&aframe.data[ay_off as usize..(ay_off + BLOCK_SIZE) as usize]);
        }
    }
    fn compare_block(aframe: &Frame, bframe: &Frame, aoff: IVec2, boff: IVec2) -> u32 {
        assert_eq!(aframe.res, bframe.res);
        let res = aframe.res;
        let mut diff = 0;

        // Luma
        for y in 0..BLOCK_SIZE {
            let ay_off = aoff.x + (y + aoff.y) * res.x;
            let by_off = boff.x + (y + boff.y) * res.x;

            for x in 0..BLOCK_SIZE {
                diff += aframe.data[(ay_off + x) as usize]
                    .abs_diff(bframe.data[(by_off + x) as usize]) as u32
            }
        }

        // Chroma
        let uvplane_off = res.x * res.y;

        for y in 0..BLOCK_SIZE / 2 {
            let ay_off = uvplane_off + (aoff.x & !1) + (y + aoff.y / 2) * res.x;
            let by_off = uvplane_off + (boff.x & !1) + (y + boff.y / 2) * res.x;

            for x in 0..BLOCK_SIZE {
                diff += aframe.data[(ay_off + x) as usize]
                    .abs_diff(bframe.data[(by_off + x) as usize]) as u32;
            }
        }

        diff
    }
}