use rshogi::records::formats::kif::{parse_ki2_bytes, parse_kif_bytes};
use rshogi::records::formats::pack;
use rshogi::records::formats::sbinpack::{chain_from_record, deserialize_file, serialize_file};
use rshogi::records::record::{GameRecord, MoveRecord};
use rshogi::types::Eval;
use std::env;
use std::fs;
use std::io;
use std::path::Path;
use std::time::{Duration, Instant};

const DEFAULT_MAX_GAMES: usize = 1000;

fn is_kif_file(path: &Path) -> bool {
    path.extension()
        .and_then(|ext| ext.to_str())
        .map_or(false, |ext| ext.eq_ignore_ascii_case("kif"))
}

fn walk_kif_files(
    dir: &Path,
    limit: usize,
    seen: &mut usize,
    handler: &mut impl FnMut(&Path) -> bool,
) -> io::Result<bool> {
    let mut entries: Vec<_> = fs::read_dir(dir)?.collect::<Result<Vec<_>, io::Error>>()?;
    entries.sort_by_key(|entry| entry.file_name());

    for entry in entries {
        if *seen >= limit {
            return Ok(false);
        }
        let path = entry.path();
        if path.is_dir() {
            if !walk_kif_files(&path, limit, seen, handler)? {
                return Ok(false);
            }
        } else if is_kif_file(&path) {
            *seen += 1;
            if !handler(&path) {
                return Ok(false);
            }
        }
    }
    Ok(true)
}

fn parse_record(bytes: &[u8]) -> Result<GameRecord, String> {
    match parse_kif_bytes(bytes) {
        Ok(record) => Ok(record),
        Err(err) => parse_ki2_bytes(bytes).map_err(|err2| format!("kif={err}; ki2={err2}")),
    }
}

fn with_filled_eval(record: &GameRecord) -> (GameRecord, usize) {
    let mut missing = 0usize;
    let mut moves: Vec<MoveRecord> = Vec::with_capacity(record.num_moves());
    for mv in record.main_moves() {
        let mut cloned = mv.clone();
        if cloned.eval().is_none() {
            cloned.set_eval(Some(Eval::from_raw(0)));
            missing += 1;
        }
        moves.push(cloned);
    }
    let mut rebuilt = GameRecord::from_main_line(
        record.init_position_sfen().to_string(),
        moves,
        record.result().clone(),
    )
    .expect("record rebuild should succeed");
    rebuilt.set_metadata(record.metadata().clone());
    (rebuilt, missing)
}

fn stem_score_from_record(record: &GameRecord) -> i16 {
    record.main_moves().next().and_then(|mv| mv.eval()).map(|eval| eval.raw()).unwrap_or(0)
}

fn format_rate(label: &str, duration: Duration, games: usize, bytes: u64) -> String {
    if duration.as_nanos() == 0 {
        return format!("  {label}: 0.00 ms");
    }
    let secs = duration.as_secs_f64();
    let games_per_sec = games as f64 / secs;
    let mb_per_sec = (bytes as f64 / (1024.0 * 1024.0)) / secs;
    format!(
        "  {label}: {:.2} ms ({:.1} games/s, {:.2} MiB/s)",
        duration.as_secs_f64() * 1000.0,
        games_per_sec,
        mb_per_sec
    )
}

fn main() -> io::Result<()> {
    let mut args = env::args().skip(1);
    let dir = args.next().unwrap_or_else(|| "_refs/data/kif/Suisho2mkif-1100k/2mkif".to_string());
    let max_games =
        args.next().and_then(|value| value.parse::<usize>().ok()).unwrap_or(DEFAULT_MAX_GAMES);

    let mut records: Vec<GameRecord> = Vec::new();
    let mut attempted = 0usize;
    let mut parsed_ok = 0usize;
    let mut filled_eval = 0usize;
    let mut parse_errors = 0usize;

    let mut seen = 0usize;
    let mut handler = |path: &Path| -> bool {
        attempted += 1;
        let bytes = match fs::read(path) {
            Ok(bytes) => bytes,
            Err(_) => {
                parse_errors += 1;
                return true;
            }
        };
        let record = match parse_record(&bytes) {
            Ok(record) => record,
            Err(_) => {
                parse_errors += 1;
                return true;
            }
        };
        parsed_ok += 1;

        let (record, missing) = with_filled_eval(&record);
        filled_eval += missing;
        records.push(record);
        true
    };

    let _ = walk_kif_files(Path::new(&dir), max_games, &mut seen, &mut handler)?;

    let mut pack_encode_time = Duration::ZERO;
    let mut sbin_encode_time = Duration::ZERO;
    let mut pack_decode_time = Duration::ZERO;
    let mut sbin_decode_time = Duration::ZERO;

    let mut pack_bytes_all: Vec<Vec<u8>> = Vec::new();
    let mut sbin_bytes_all: Vec<Vec<u8>> = Vec::new();

    let mut pack_ok = 0usize;
    let mut sbin_ok = 0usize;
    let mut both_ok = 0usize;
    let mut pack_errors = 0usize;
    let mut sbin_errors = 0usize;
    let mut total_pack_bytes: u64 = 0;
    let mut total_sbin_bytes: u64 = 0;

    for record in &records {
        let start = Instant::now();
        let pack_bytes = pack::serialize(record);
        pack_encode_time += start.elapsed();

        let start = Instant::now();
        let sbin_bytes = chain_from_record(record, stem_score_from_record(record))
            .and_then(|chain| serialize_file(&[chain]));
        sbin_encode_time += start.elapsed();

        let pack_bytes = match pack_bytes {
            Ok(bytes) => {
                pack_ok += 1;
                Some(bytes)
            }
            Err(_) => {
                pack_errors += 1;
                None
            }
        };

        let sbin_bytes = match sbin_bytes {
            Ok(bytes) => {
                sbin_ok += 1;
                Some(bytes)
            }
            Err(_) => {
                sbin_errors += 1;
                None
            }
        };

        if let (Some(pack_bytes), Some(sbin_bytes)) = (pack_bytes, sbin_bytes) {
            total_pack_bytes += pack_bytes.len() as u64;
            total_sbin_bytes += sbin_bytes.len() as u64;
            pack_bytes_all.push(pack_bytes);
            sbin_bytes_all.push(sbin_bytes);
            both_ok += 1;
        }
    }

    for bytes in &pack_bytes_all {
        let start = Instant::now();
        let _ = pack::deserialize(bytes);
        pack_decode_time += start.elapsed();
    }

    for bytes in &sbin_bytes_all {
        let start = Instant::now();
        let _ = deserialize_file(bytes);
        sbin_decode_time += start.elapsed();
    }

    println!("pack vs sbinpack benchmark");
    println!("  dir: {dir}");
    println!("  files scanned: {seen}");
    println!("  attempted: {attempted}");
    println!("  parsed: {parsed_ok} (errors: {parse_errors})");
    println!("  filled eval: {filled_eval}");
    println!("  pack ok: {pack_ok} (errors: {pack_errors})");
    println!("  sbinpack ok: {sbin_ok} (errors: {sbin_errors})");
    println!("  both ok: {both_ok}");
    println!("  total pack bytes (both ok): {total_pack_bytes}");
    println!("  total sbinpack bytes (both ok): {total_sbin_bytes}");
    if total_pack_bytes > 0 {
        let ratio = total_sbin_bytes as f64 / total_pack_bytes as f64;
        let savings = (1.0 - ratio) * 100.0;
        println!("  ratio (sbinpack/pack): {ratio:.4}");
        println!("  savings vs pack: {savings:.2}%");
    }
    println!();
    println!("{}", format_rate("pack encode", pack_encode_time, records.len(), total_pack_bytes));
    println!(
        "{}",
        format_rate("sbinpack encode", sbin_encode_time, records.len(), total_sbin_bytes)
    );
    println!("{}", format_rate("pack decode", pack_decode_time, both_ok, total_pack_bytes));
    println!("{}", format_rate("sbinpack decode", sbin_decode_time, both_ok, total_sbin_bytes));

    Ok(())
}
