pub mod decode;
pub mod encode;

use decode::Dec;
use encode::Enc;
use framework::bitstream_filter_main;
use glam::IVec2;
use std::env::args;

pub const BLOCK_SIZE: i32 = 8;

fn main() -> Result<(), std::io::Error> {
    match args().nth(1).unwrap().as_str() {
        "enc" => bitstream_filter_main::<Enc>(),
        "dec" => bitstream_filter_main::<Dec>(),
        _ => panic!("unknown mode"),
    }
}

#[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 export_block(&self, off: IVec2, buf: &mut Vec<u8>) {
        // Luma
        for y in 0..BLOCK_SIZE {
            let y_off = off.x + (y + off.y) * self.res.x;
            buf.extend(&self.data[y_off as usize..(y_off + BLOCK_SIZE) as usize]);
        }
        // Chroma
        let uvplane_off = self.res.x * self.res.y;
        for y in 0..BLOCK_SIZE / 2 {
            let y_off = uvplane_off + (off.x & !1) + (y + off.y / 2) * self.res.x;
            buf.extend(&self.data[y_off as usize..(y_off + BLOCK_SIZE) as usize]);
        }
    }
    fn import_block(&mut self, off: IVec2, buf: &[u8]) -> usize {
        let mut p = 0;
        // Luma
        for y in 0..BLOCK_SIZE {
            let y_off = off.x + (y + off.y) * self.res.x;
            self.data[y_off as usize..(y_off + BLOCK_SIZE) as usize]
                .copy_from_slice(&buf[p..p + BLOCK_SIZE as usize]);
            p += BLOCK_SIZE as usize;
        }
        // Chroma
        let uvplane_off = self.res.x * self.res.y;
        for y in 0..BLOCK_SIZE / 2 {
            let y_off = uvplane_off + (off.x & !1) + (y + off.y / 2) * self.res.x;
            self.data[y_off as usize..(y_off + BLOCK_SIZE) as usize]
                .copy_from_slice(&buf[p..p + BLOCK_SIZE as usize]);
            p += BLOCK_SIZE as usize;
        }
        p
    }
    fn import_block_diff(&mut self, off: IVec2, buf: &[u8]) -> usize {
        // Luma
        for y in 0..BLOCK_SIZE {
            let y_off = off.x + (y + off.y) * self.res.x;
            let i = y * BLOCK_SIZE;
            for x in 0..BLOCK_SIZE {
                self.data[(y_off + x) as usize] = self.data[(y_off + x) as usize]
                    .saturating_add_signed(buf[(i + x) as usize] as i8 - 127);
            }
        }
        // Chroma
        let uvplane_off = self.res.x * self.res.y;
        for y in 0..BLOCK_SIZE / 2 {
            let y_off = uvplane_off + (off.x & !1) + (y + off.y / 2) * self.res.x;
            let i = BLOCK_SIZE * BLOCK_SIZE + y * BLOCK_SIZE;
            for x in 0..BLOCK_SIZE {
                self.data[(y_off + x) as usize] = self.data[(y_off + x) as usize]
                    .saturating_add_signed(buf[(i + x) as usize] as i8 - 127);
            }
        }
        (BLOCK_SIZE * BLOCK_SIZE + BLOCK_SIZE * BLOCK_SIZE / 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
    }
}