use framework::{
    common::huffman::encode_huff,
    vector::{UVec2, Vec2},
    Frame, Framework, Pixel,
};
use rayon::join;

fn main() {
    let (mut framework, params) = Framework::init();
    let root = Area {
        x1: 0,
        y1: 0,
        x2: params.width,
        y2: params.height,
    };

    match params.mode {
        framework::CodecMode::Encode => {
            let mut oframe = Frame::new(params.width, params.height);
            let mut dframe = Frame::new(params.width, params.height);
            let mut out = Vec::<u8>::new();
            while let Some(iframe) = framework.next_frame() {
                let t = encode(&oframe, &iframe, root);
                decode(&mut oframe, root, &t);
                write(&mut out, &t);

                if params.debug == 1 {
                    dframe.pixels.copy_from_slice(&oframe.pixels);
                    debug(&mut dframe, root, &t);
                    framework.decode_done(&dframe);
                } else if params.debug == 2 {
                    debug_diff(&mut dframe, root, &t);
                    framework.decode_done(&dframe);
                } else {
                    framework.decode_done(&oframe)
                }

                let huff = encode_huff(&out);
                framework.encode_done(&huff);
                out.clear();
            }
        }
        framework::CodecMode::Decode => todo!(),
    }
}

pub fn write(w: &mut Vec<u8>, t: &DiffTree) {
    match t {
        DiffTree::Split([a, b]) => {
            w.push(0);
            write(w, a);
            write(w, b);
        }
        DiffTree::Diff(d) => w.extend([1, d.r as u8, d.g as u8, d.b as u8]),
    }
}

pub fn decode(f: &mut Frame, area: Area, tree: &DiffTree) {
    match tree {
        DiffTree::Split([ta, tb]) => {
            let (aa, ab) = area.split();
            decode(f, aa, ta);
            decode(f, ab, tb);
        }
        DiffTree::Diff(diff) => {
            apply_area_diff(f, area, *diff);
        }
    }
}
pub fn debug(f: &mut Frame, area: Area, tree: &DiffTree) {
    match tree {
        DiffTree::Split([ta, tb]) => {
            let (aa, ab) = area.split();
            debug(f, aa, ta);
            debug(f, ab, tb);
        }
        DiffTree::Diff(_diff) => {
            let Area { x1, y1, x2, y2 } = area;
            for x in x1..x2 {
                for y in y1..y2 {
                    f[(x, y)] = Pixel { r: 0, g: 0, b: 255 }
                }
            }
        }
    }
}
pub fn debug_diff(f: &mut Frame, area: Area, tree: &DiffTree) {
    match tree {
        DiffTree::Split([ta, tb]) => {
            let (aa, ab) = area.split();
            debug_diff(f, aa, ta);
            debug_diff(f, ab, tb);
        }
        DiffTree::Diff(diff) => {
            let Area { x1, y1, x2, y2 } = area;
            for x in x1..x2 {
                for y in y1..y2 {
                    f[(x, y)] = Pixel {
                        r: 127u8.saturating_add_signed(diff.r),
                        b: 127u8.saturating_add_signed(diff.b),
                        g: 127u8.saturating_add_signed(diff.g),
                    }
                }
            }
        }
    }
}

pub fn encode(a: &Frame, b: &Frame, area: Area) -> DiffTree {
    if area.area() == 1 {
        DiffTree::Diff(diff_pixel(
            a,
            b,
            Vec2 {
                x: area.x1,
                y: area.y1,
            },
        ))
    } else {
        let (aa, ba) = area.split();
        let (at, bt) = join(|| encode(a, b, aa), || encode(a, b, ba));
        // let (at, bt) = (encode(a, b, aa), encode(a, b, ba));

        match (&at, &bt) {
            (DiffTree::Diff(ad), DiffTree::Diff(bd)) => {
                let d_r = ad.r.abs_diff(bd.r);
                let d_g = ad.g.abs_diff(bd.g);
                let d_b = ad.b.abs_diff(bd.b);

                let visdiff = (d_r as usize + d_g as usize + d_b as usize) * aa.area();
                if visdiff < 100 {
                    return DiffTree::Diff(Pixel {
                        r: ((ad.r as i16 + bd.r as i16) / 2) as i8,
                        g: ((ad.g as i16 + bd.g as i16) / 2) as i8,
                        b: ((ad.b as i16 + bd.b as i16) / 2) as i8,
                    });
                }
            }
            _ => (),
        }
        DiffTree::Split([Box::new(at), Box::new(bt)])
    }
}

pub enum DiffTree {
    Split([Box<DiffTree>; 2]),
    Diff(Pixel<i8>),
}

#[derive(Debug, Clone, Copy)]
pub struct Area {
    x1: usize,
    y1: usize,
    x2: usize,
    y2: usize,
}

#[inline]
fn diff_clamp(x: u8, y: u8) -> i8 {
    if y >= x {
        (y - x).min(127) as i8
    } else {
        -((x - y).min(128) as i8)
    }
}
pub fn diff_pixel(a: &Frame, b: &Frame, p: UVec2) -> Pixel<i8> {
    let (ap, bp) = (a[p], b[p]);
    Pixel {
        r: diff_clamp(ap.r, bp.r),
        g: diff_clamp(ap.g, bp.g),
        b: diff_clamp(ap.b, bp.b),
    }
}

pub fn apply_area_diff(frame: &mut Frame, Area { x1, y1, x2, y2 }: Area, diff: Pixel<i8>) {
    for x in x1..x2 {
        for y in y1..y2 {
            let p = &mut frame[(x, y)];
            p.r = p.r.saturating_add_signed(diff.r);
            p.g = p.g.saturating_add_signed(diff.g);
            p.b = p.b.saturating_add_signed(diff.b);
        }
    }
}

impl Area {
    pub fn area(&self) -> usize {
        self.width() as usize * self.height() as usize
    }
    pub fn width(&self) -> usize {
        self.x2 - self.x1
    }
    pub fn height(&self) -> usize {
        self.y2 - self.y1
    }
    pub fn split(&self) -> (Self, Self) {
        let Area { x1, y1, x2, y2 } = *self;
        if self.width() > self.height() {
            let xm = (self.x1 + self.x2) / 2;
            (Self { x1, x2: xm, y1, y2 }, Self { x1: xm, x2, y1, y2 })
        } else {
            let ym = (self.y1 + self.y2) / 2;
            (Self { x1, x2, y1, y2: ym }, Self { x1, x2, y1: ym, y2 })
        }
    }
}