mod utils;

use clap::Parser;
use std::collections::HashMap;
use std::io::Write;
use std::{fs, io, time};
use utils::{Metadata, Zsdiff, get_hash};
use walkdir::WalkDir;

struct FileInfo {
    path: String,
    relative_path: String, // Without dir prefix
    hash: String,
}

async fn walk_dir(dir: String) -> HashMap<String, FileInfo> {
    let mut hash_list: HashMap<String, FileInfo> = HashMap::new();
    for e in WalkDir::new(&dir) {
        let e = e.unwrap();
        let path = e.path();
        if path.is_dir() {
            continue;
        }
        let content = fs::read(path).unwrap();
        let hash = get_hash(content).await;
        let path_str = path.display().to_string();
        let file_info = FileInfo {
            relative_path: path_str[dir.len() + 1..].to_string(),
            path: path_str,
            hash: hash.clone(),
        };
        hash_list.entry(hash).or_insert(file_info);
    }
    hash_list
}

async fn compare_hashes(old: HashMap<String, FileInfo>, new: HashMap<String, FileInfo>) -> Zsdiff {
    let mut diff_files: HashMap<String, Vec<u8>> = HashMap::new();
    let mut remove_files: Vec<String> = vec![];
    let mut hashes: HashMap<String, String> = HashMap::new();
    for (_, info) in &old {
        remove_files.push(info.relative_path.clone());
    }

    for (new_hash, new_fileinfo) in &new {
        let old_fileinfo = old.get(new_hash);
        remove_files.retain(|filename| !filename.eq(&new_fileinfo.relative_path));

        if old_fileinfo.is_none() {
            let path = new_fileinfo.relative_path.clone();
            diff_files.insert(path.clone(), fs::read(new_fileinfo.path.clone()).unwrap());
            hashes.insert(
                new_fileinfo.relative_path.clone(),
                new_fileinfo.hash.clone(),
            );
        }
    }

    Zsdiff {
        content: diff_files.clone(),
        metadata: Metadata {
            diff_files: diff_files.keys().cloned().collect(),
            hashes,
            remove_files,
        },
    }
}

pub async fn zsdiff(
    filename: String,
    old: String,
    new: String,
    level: i32,
) -> Result<(), io::Error> {
    let output_filename = &format!("{}.zdiff", filename);
    let old_hashes = walk_dir(old).await;
    let new_hashes = walk_dir(new).await;
    let compare_hashes = compare_hashes(old_hashes, new_hashes).await;
    let parts = compare_hashes.to_vec().await;
    let mut size_before = 0;
    for p in &parts {
        size_before += p.len();
    }
    let now = time::Instant::now();
    utils::compress_parts(parts, &fs::File::create(output_filename)?, level).await;
    let output_data = fs::read(output_filename)?;
    let size_after = output_data.len();
    let output_hash = get_hash(output_data).await;
    fs::File::create(format!("{}.md5", output_filename))?.write_all(output_hash.as_bytes())?;
    let elapsed = now.elapsed();
    println!("Zsdiff hash: {}", output_hash);
    println!("Size before: {:.1?}KB", size_before / 1024);
    println!("Size after: {:.1?}KB", size_after / 1024);
    println!(
        "Compress ratio: {:.2?}%",
        size_after as f64 / size_before as f64 * 100.0
    );
    print!("Time elapsed: {:.2?}s", elapsed);
    Ok(())
}

#[derive(Parser)]
struct Args {
    #[arg(short, long)]
    filename: String,
    #[arg(short, long, default_value_t = 11)]
    compress_level: i32,
    #[arg(short, long)]
    old: String,
    #[arg(short, long)]
    new: String,
}

#[tokio::main]
async fn main() -> io::Result<()> {
    let args = Args::parse();
    zsdiff(args.filename, args.old, args.new, args.compress_level).await
}