use anyhow::Context;
use clap::Parser;
use indicatif::ProgressBar;
use libreschmux::{
    codec::{
        compress::compress_block,
        decode::{decode_block, DecodeConfig},
        encode::{encode_block, EncodeConfig, EncodeMode},
    },
    format::{header::Header, ser::Sink},
    frame::Frame,
    helpers::vector::Vec2,
};
use log::info;
use rayon::{ThreadPool, ThreadPoolBuilder};
use std::io::{BufReader, BufWriter};

#[derive(Parser)]
#[clap(about, version)]
pub struct EncodeArgs {
    #[arg(short = 'W', long)]
    width: usize,
    #[arg(short = 'H', long)]
    height: usize,

    #[arg(short, long)]
    mode: EncodeMode,

    #[arg(long)]
    no_linear_transform: bool,
    #[arg(long)]
    no_value_scale: bool,
    #[arg(long)]
    no_translate: bool,

    #[arg[short = 'L', long]]
    length: Option<usize>,

    #[arg(short, long, default_value = "8")]
    min_block_size: isize,

    #[arg(short = 't', long, default_value = "200")]
    ref_thres: f64,
}

fn main() -> anyhow::Result<()> {
    env_logger::init_from_env("LOG");
    let args = EncodeArgs::parse();

    let mut input = BufReader::new(std::io::stdin());
    let mut output = BufWriter::new(std::io::stdout());

    let config = EncodeConfig {
        mode: args.mode,
        ref_thres: args.ref_thres,
        weight_factor: 50.0,
        max_diff_area: 10_000,
        min_block_size: args.min_block_size,
        do_translate: !args.no_translate,
        do_linear_transform: !args.no_linear_transform,
        do_value_scale: !args.no_value_scale,
    };

    let size = Vec2 {
        x: args.width as isize,
        y: args.height as isize,
    };
    output
        .put(Header {
            resolution: size,
            frame_count: args.length.unwrap_or(usize::MAX),
        })
        .context("writing header")?;

    let progress_bar = args.length.map(|len| ProgressBar::new(len as u64));
    let mut prev_frame = Frame::new(size);

    for i in 0.. {
        info!("encode frame {i}");
        let mut frame = Frame::read(&mut input, size).context("reading raw frame")?;

        let v1 = frame.view();
        let v2 = prev_frame.view();

        let (error, root) = encode_block(v1, v2, &config);

        // compress_block(
        //     &mut root,
        //     Vec2 {
        //         x: size.x as usize,
        //         y: size.y as usize,
        //     },
        // );

        root.write(&mut output, size)
            .context("writing encoded frame")?;

        info!(
            "cumulative error: {error} ({} per pixel)",
            error / frame.view().area() as f64
        );

        if let Some(progress_bar) = &progress_bar {
            progress_bar.inc(1);
        }

        decode_block(&root, frame.view_mut(), prev_frame.view(), &DecodeConfig {});

        prev_frame = frame;
    }

    Ok(())
}