use bv1::{
    debug::draw_debug,
    decode::decode,
    encode::{encode, EncodeConfig},
    huff::{read_huff, write_huff},
    Block, Frame, Pixel, PixelValue, View, P2,
};
use clap::{Parser, Subcommand};
use std::{
    io::{stdin, stdout, BufReader, BufWriter, Read, Write},
    time::Instant,
};

#[derive(Parser)]
#[clap(about, version)]
struct Args {
    // Width of the video signal
    #[arg(short = 'W', long)]
    width: u16,
    // Height of the video signal
    #[arg(short = 'H', long)]
    height: u16,
    #[clap(subcommand)]
    action: Action,
}

#[derive(Clone, Subcommand)]
enum Action {
    // Compress video
    Encode {
        #[arg(short, long, default_value_t = 800)]
        max_block_size: usize,
        #[arg(short, long, default_value_t = 10_000)]
        attention_split: u32,
        #[arg(short, long, default_value_t = 10.)]
        threshold: f32,
        #[arg(short, long, default_value_t = 10)]
        keyframe_interval: usize,
    },
    // Decompress video
    Decode {
        #[arg(short, long)]
        debug: bool,
    },
}

fn main() {
    let args = Args::parse();

    let size = P2 {
        x: args.width as i32,
        y: args.height as i32,
    };
    match args.action {
        Action::Encode {
            max_block_size,
            threshold,
            attention_split,
            keyframe_interval,
        } => {
            let config = EncodeConfig {
                threshold,
                max_block_size,
                attention_split,
            };
            let mut last_frame = Frame::new(size);

            let mut stdin = BufReader::new(stdin());
            let mut stdout = BufWriter::new(stdout());

            let huff = true;

            for frame_number in 0.. {
                let mut frame = read_frame(&mut stdin, size);

                let mut config = config.clone();
                if frame_number % keyframe_interval != 0 {
                    config.threshold = std::f32::INFINITY;
                }

                let t = Instant::now();
                let b: Block = encode(&last_frame, &frame, View::all(size), &config);
                let time_encode = t.elapsed();

                let t = Instant::now();
                decode(&last_frame, &mut frame, View::all(size), &b);
                last_frame = frame;
                let time_decode = t.elapsed();

                if huff {
                    let mut buf = vec![];
                    let mut bufw = std::io::Cursor::new(&mut buf);
                    b.write(&mut bufw).unwrap();
                    drop(bufw);
                    let t = Instant::now();
                    let bits_raw = buf.len() * 8;
                    let bits_huff = write_huff(&buf, &mut stdout).unwrap();
                    let time_huff = t.elapsed();
                    drop(buf);

                    eprintln!(
                        "frame {frame_number}: {:?}",
                        time_decode + time_huff + time_encode
                    );
                    eprintln!(
                        "\tencode {time_encode:?} ({:.2}%)",
                        (bits_raw as f32 / (size.area() * 24) as f32) * 100.0
                    );
                    eprintln!(
                        "\thuff {time_huff:?} ({:.2}%)",
                        (bits_huff as f32 / bits_raw as f32) * 100.0
                    );
                    eprintln!("\tdecode {time_decode:?}");
                } else {
                    b.write(&mut stdout).unwrap();
                }
            }
        }
        Action::Decode { debug } => {
            let mut frame = Frame::new(size);
            let mut last_frame = Frame::new(size);
            let mut debug_frame = if debug { Some(Frame::new(size)) } else { None };

            let mut stdin = BufReader::new(stdin());
            let mut stdout = BufWriter::new(stdout());

            let huff = true;

            loop {
                let b = if huff {
                    let mut buf = vec![];
                    read_huff(&mut stdin, &mut buf).unwrap();
                    let mut buf = std::io::Cursor::new(&mut buf);
                    Block::read(&mut buf, View::all(size)).unwrap()
                } else {
                    Block::read(&mut stdin, View::all(size)).unwrap()
                };

                decode(&last_frame, &mut frame, View::all(size), &b);

                if let Some(debug_frame) = &mut debug_frame {
                    debug_frame.pixels.copy_from_slice(&frame.pixels);
                    draw_debug(debug_frame, View::all(size), &b);
                    write_frame(&mut stdout, &debug_frame);
                } else {
                    write_frame(&mut stdout, &frame);
                }

                last_frame.pixels.copy_from_slice(&frame.pixels); // TODO use mem::swap
                frame.pixels.iter_mut().for_each(|e| *e = Pixel::BLACK);
            }
        }
    }
}

fn read_frame(inp: &mut impl Read, size: P2) -> Frame {
    let mut f = Frame::new(size);

    for y in 0..size.y {
        for x in 0..size.x {
            let mut cc = [0u8; 3];
            inp.read_exact(&mut cc).unwrap();
            f[P2 { x, y }] = Pixel {
                r: cc[0] as PixelValue,
                g: cc[1] as PixelValue,
                b: cc[2] as PixelValue,
            };
        }
    }
    f
}

fn write_frame(out: &mut impl Write, frame: &Frame) {
    for y in 0..frame.size.y {
        for x in 0..frame.size.x {
            let p = frame[P2 { x, y }];
            let mut cc = [0u8; 3];
            cc[0] = p.r.clamp(0, 255) as u8;
            cc[1] = p.g.clamp(0, 255) as u8;
            cc[2] = p.b.clamp(0, 255) as u8;
            out.write_all(&mut cc).unwrap();
        }
    }
}