use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::path::{Path, PathBuf};

use rpki::cir::decode_cir;
use serde_json::json;
use sha2::{Digest, Sha256};

#[derive(Debug, Default, PartialEq, Eq)]
struct Args {
    ours_cir: Option<PathBuf>,
    peer_cir: Option<PathBuf>,
    cache_root: Option<PathBuf>,
    rpki_client_log: Option<PathBuf>,
    out_json: Option<PathBuf>,
    sample_limit: usize,
}

fn usage() -> &'static str {
    "Usage: cir_probe_rpki_client_cache --ours-cir <path> --rpki-client-cir <path> --cache-root <path> --out-json <path> [--rpki-client-log <path>] [--sample-limit <n>]"
}

fn main() {
    if let Err(err) = real_main() {
        eprintln!("{err}");
        std::process::exit(1);
    }
}

fn real_main() -> Result<(), String> {
    run(parse_args(&std::env::args().collect::<Vec<_>>())?)
}

fn parse_args(argv: &[String]) -> Result<Args, String> {
    let mut args = Args {
        sample_limit: 20,
        ..Args::default()
    };
    let mut index = 1usize;
    while index < argv.len() {
        match argv[index].as_str() {
            "--ours-cir" => {
                index += 1;
                args.ours_cir = Some(PathBuf::from(
                    argv.get(index).ok_or("--ours-cir requires a value")?,
                ));
            }
            "--rpki-client-cir" | "--peer-cir" => {
                index += 1;
                args.peer_cir = Some(PathBuf::from(
                    argv.get(index)
                        .ok_or("--rpki-client-cir requires a value")?,
                ));
            }
            "--cache-root" => {
                index += 1;
                args.cache_root = Some(PathBuf::from(
                    argv.get(index).ok_or("--cache-root requires a value")?,
                ));
            }
            "--rpki-client-log" => {
                index += 1;
                args.rpki_client_log = Some(PathBuf::from(
                    argv.get(index)
                        .ok_or("--rpki-client-log requires a value")?,
                ));
            }
            "--out-json" => {
                index += 1;
                args.out_json = Some(PathBuf::from(
                    argv.get(index).ok_or("--out-json requires a value")?,
                ));
            }
            "--sample-limit" => {
                index += 1;
                let value = argv.get(index).ok_or("--sample-limit requires a value")?;
                args.sample_limit = value
                    .parse::<usize>()
                    .map_err(|_| format!("invalid --sample-limit: {value}"))?;
            }
            "-h" | "--help" => return Err(usage().to_string()),
            other => return Err(format!("unknown argument: {other}\n{}", usage())),
        }
        index += 1;
    }
    if args.ours_cir.is_none() {
        return Err(format!("--ours-cir is required\n{}", usage()));
    }
    if args.peer_cir.is_none() {
        return Err(format!("--rpki-client-cir is required\n{}", usage()));
    }
    if args.cache_root.is_none() {
        return Err(format!("--cache-root is required\n{}", usage()));
    }
    if args.out_json.is_none() {
        return Err(format!("--out-json is required\n{}", usage()));
    }
    Ok(args)
}

fn run(args: Args) -> Result<(), String> {
    let ours_cir_path = args.ours_cir.as_ref().expect("validated");
    let peer_cir_path = args.peer_cir.as_ref().expect("validated");
    let cache_root = args.cache_root.as_ref().expect("validated");
    let out_json = args.out_json.as_ref().expect("validated");

    let ours = decode_cir(&read_file(ours_cir_path)?)
        .map_err(|e| format!("decode ours CIR failed: {e}"))?;
    let peer = decode_cir(&read_file(peer_cir_path)?)
        .map_err(|e| format!("decode rpki-client CIR failed: {e}"))?;
    let peer_objects = peer
        .validated_objects()
        .map(|item| item.rsync_uri.as_str())
        .collect::<BTreeSet<_>>();
    let only_in_ours = ours
        .validated_objects()
        .filter(|item| !peer_objects.contains(item.rsync_uri.as_str()))
        .map(|item| ProbeObject {
            uri: item.rsync_uri.clone(),
            sha256_hex: hex::encode(&item.sha256),
        })
        .collect::<Vec<_>>();

    let rrdp_dirs = list_rrdp_dirs(cache_root)?;
    let log_mentions = match args.rpki_client_log.as_ref() {
        Some(path) => build_log_mentions(path, &only_in_ours)?,
        None => LogMentions::default(),
    };

    let mut cache_hash_match_count = 0usize;
    let mut cache_hash_mismatch_count = 0usize;
    let mut cache_missing_count = 0usize;
    let mut valid_path_exists_count = 0usize;
    let mut rrdp_temp_path_exists_count = 0usize;
    let mut rsync_temp_path_exists_count = 0usize;
    let mut log_mention_count = 0usize;
    let mut samples = Vec::new();
    let mut missing_uris = Vec::new();
    let mut matched_uris = Vec::new();
    let mut mismatched_uris = Vec::new();
    let mut log_mentioned_uris = Vec::new();

    for object in &only_in_ours {
        let probe = probe_cache(cache_root, &rrdp_dirs, object)?;
        if probe.hash_match {
            cache_hash_match_count += 1;
            matched_uris.push(object.uri.as_str());
        } else if probe.exists {
            cache_hash_mismatch_count += 1;
            mismatched_uris.push(object.uri.as_str());
        } else {
            cache_missing_count += 1;
            missing_uris.push(object.uri.as_str());
        }
        if probe.valid_path_exists {
            valid_path_exists_count += 1;
        }
        if probe.rrdp_temp_path_exists {
            rrdp_temp_path_exists_count += 1;
        }
        if probe.rsync_temp_path_exists {
            rsync_temp_path_exists_count += 1;
        }
        let log_mention = log_mentions.mentioned.contains(object.uri.as_str());
        if log_mention {
            log_mention_count += 1;
            log_mentioned_uris.push(object.uri.as_str());
        }
        if samples.len() < args.sample_limit {
            samples.push(json!({
                "uri": object.uri,
                "sha256": object.sha256_hex,
                "cacheStatus": probe.status(),
                "cacheLocations": probe.locations,
                "logMentionedAsSuperfluousOrDeleted": log_mention,
            }));
        }
    }

    let summary = json!({
        "onlyInOursTotal": only_in_ours.len(),
        "cacheRoot": cache_root,
        "rpkiClientLog": args.rpki_client_log,
        "rrdpCacheDirs": rrdp_dirs.len(),
        "cacheProbe": {
            "hashMatchCount": cache_hash_match_count,
            "hashMismatchCount": cache_hash_mismatch_count,
            "missingCount": cache_missing_count,
            "validPathExistsCount": valid_path_exists_count,
            "rrdpTempPathExistsCount": rrdp_temp_path_exists_count,
            "rsyncTempPathExistsCount": rsync_temp_path_exists_count,
            "hashMatchByExtension": group_by_extension(matched_uris.iter().copied()),
            "missingByExtension": group_by_extension(missing_uris.iter().copied()),
            "hashMismatchByExtension": group_by_extension(mismatched_uris.iter().copied()),
            "hashMatchByHostTop": top_hosts(matched_uris.iter().copied(), args.sample_limit),
            "missingByHostTop": top_hosts(missing_uris.iter().copied(), args.sample_limit),
            "hashMismatchByHostTop": top_hosts(mismatched_uris.iter().copied(), args.sample_limit),
        },
        "logProbe": {
            "enabled": args.rpki_client_log.is_some(),
            "relevantLineCount": log_mentions.relevant_line_count,
            "mentionedCount": log_mention_count,
            "mentionedByExtension": group_by_extension(log_mentioned_uris.iter().copied()),
            "mentionedByHostTop": top_hosts(log_mentioned_uris.iter().copied(), args.sample_limit),
        },
        "samples": samples,
    });

    write_json(out_json, &summary)?;
    println!("{}", out_json.display());
    Ok(())
}

#[derive(Clone, Debug)]
struct ProbeObject {
    uri: String,
    sha256_hex: String,
}

#[derive(Default, Debug)]
struct CacheProbe {
    exists: bool,
    hash_match: bool,
    valid_path_exists: bool,
    rrdp_temp_path_exists: bool,
    rsync_temp_path_exists: bool,
    locations: Vec<serde_json::Value>,
}

impl CacheProbe {
    fn status(&self) -> &'static str {
        if self.hash_match {
            "hash_match"
        } else if self.exists {
            "hash_mismatch"
        } else {
            "missing"
        }
    }
}

fn probe_cache(
    cache_root: &Path,
    rrdp_dirs: &[PathBuf],
    object: &ProbeObject,
) -> Result<CacheProbe, String> {
    let Some(stripped) = strip_rsync_uri(&object.uri) else {
        return Ok(CacheProbe::default());
    };
    let mut probe = CacheProbe::default();
    let mut candidates = vec![
        ("valid", cache_root.join(stripped)),
        ("rsync_temp", cache_root.join(".rsync").join(stripped)),
    ];
    for dir in rrdp_dirs {
        candidates.push(("rrdp_temp", dir.join(stripped)));
    }

    for (kind, path) in candidates {
        if !path.is_file() {
            continue;
        }
        let sha256_hex = sha256_file_hex(&path)?;
        let hash_matches = sha256_hex == object.sha256_hex;
        probe.exists = true;
        probe.hash_match |= hash_matches;
        match kind {
            "valid" => probe.valid_path_exists = true,
            "rsync_temp" => probe.rsync_temp_path_exists = true,
            "rrdp_temp" => probe.rrdp_temp_path_exists = true,
            _ => {}
        }
        if probe.locations.len() < 4 {
            probe.locations.push(json!({
                "kind": kind,
                "path": path,
                "sha256": sha256_hex,
                "hashMatches": hash_matches,
            }));
        }
    }

    Ok(probe)
}

#[derive(Default, Debug)]
struct LogMentions {
    mentioned: BTreeSet<String>,
    relevant_line_count: usize,
}

fn build_log_mentions(path: &Path, objects: &[ProbeObject]) -> Result<LogMentions, String> {
    let content = std::fs::read_to_string(path)
        .map_err(|e| format!("read log failed: {}: {e}", path.display()))?;
    let mut by_file_name: HashMap<String, Vec<(&str, &str)>> = HashMap::new();
    for object in objects {
        let Some(stripped) = strip_rsync_uri(&object.uri) else {
            continue;
        };
        by_file_name
            .entry(file_name(stripped).to_string())
            .or_default()
            .push((object.uri.as_str(), stripped));
    }

    let mut mentioned = BTreeSet::new();
    let mut relevant_line_count = 0usize;
    for line in content.lines() {
        if !is_relevant_cleanup_line(line) {
            continue;
        }
        relevant_line_count += 1;
        let Some(name) = line.split('/').next_back().map(trim_log_token) else {
            continue;
        };
        let Some(candidates) = by_file_name.get(name) else {
            continue;
        };
        for (uri, stripped) in candidates {
            if line.contains(stripped) {
                mentioned.insert((*uri).to_string());
            }
        }
    }

    Ok(LogMentions {
        mentioned,
        relevant_line_count,
    })
}

fn is_relevant_cleanup_line(line: &str) -> bool {
    line.contains("superfluous")
        || line.contains("deleted ")
        || line.contains("deleted superfluous")
        || line.contains("bad message digest")
        || line.contains("referenced file supposed to be deleted")
}

fn trim_log_token(token: &str) -> &str {
    token.trim_matches(|ch: char| ch == ':' || ch == ',' || ch == ')' || ch == '(')
}

fn list_rrdp_dirs(cache_root: &Path) -> Result<Vec<PathBuf>, String> {
    let rrdp = cache_root.join(".rrdp");
    let mut dirs = Vec::new();
    let Ok(read_dir) = std::fs::read_dir(&rrdp) else {
        return Ok(dirs);
    };
    for entry in read_dir {
        let entry = entry.map_err(|e| format!("read .rrdp entry failed: {e}"))?;
        let path = entry.path();
        if path.is_dir() {
            dirs.push(path);
        }
    }
    dirs.sort();
    Ok(dirs)
}

fn sha256_file_hex(path: &Path) -> Result<String, String> {
    let bytes = std::fs::read(path)
        .map_err(|e| format!("read cache file failed: {}: {e}", path.display()))?;
    Ok(hex::encode(Sha256::digest(&bytes)))
}

fn read_file(path: &Path) -> Result<Vec<u8>, String> {
    std::fs::read(path).map_err(|e| format!("read file failed: {}: {e}", path.display()))
}

fn write_json(path: &Path, value: &serde_json::Value) -> Result<(), String> {
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)
            .map_err(|e| format!("create parent dirs failed: {}: {e}", parent.display()))?;
    }
    std::fs::write(
        path,
        serde_json::to_vec_pretty(value).map_err(|e| e.to_string())?,
    )
    .map_err(|e| format!("write json failed: {}: {e}", path.display()))
}

fn strip_rsync_uri(uri: &str) -> Option<&str> {
    uri.strip_prefix("rsync://")
}

fn file_name(path: &str) -> &str {
    path.rsplit('/').next().unwrap_or(path)
}

fn group_by_extension<'a>(uris: impl IntoIterator<Item = &'a str>) -> BTreeMap<String, usize> {
    let mut counts = BTreeMap::new();
    for uri in uris {
        *counts.entry(uri_extension(uri)).or_insert(0) += 1;
    }
    counts
}

fn top_hosts<'a>(uris: impl IntoIterator<Item = &'a str>, limit: usize) -> Vec<serde_json::Value> {
    let mut counts = BTreeMap::new();
    for uri in uris {
        *counts.entry(uri_host(uri)).or_insert(0usize) += 1;
    }
    let mut rows = counts.into_iter().collect::<Vec<_>>();
    rows.sort_by(|(host_a, count_a), (host_b, count_b)| {
        count_b.cmp(count_a).then_with(|| host_a.cmp(host_b))
    });
    rows.into_iter()
        .take(limit)
        .map(|(host, count)| {
            json!({
                "host": host,
                "count": count,
            })
        })
        .collect()
}

fn uri_host(uri: &str) -> String {
    let without_scheme = uri.split_once("://").map(|(_, rest)| rest).unwrap_or(uri);
    without_scheme
        .split('/')
        .next()
        .filter(|host| !host.is_empty())
        .unwrap_or("<unknown>")
        .to_string()
}

fn uri_extension(uri: &str) -> String {
    let path = uri.split_once("://").map(|(_, rest)| rest).unwrap_or(uri);
    let path = path.split_once('/').map(|(_, path)| path).unwrap_or(path);
    let file = path.rsplit('/').next().unwrap_or(path);
    if let Some((_, ext)) = file.rsplit_once('.')
        && !ext.is_empty()
    {
        return format!(".{}", ext.to_ascii_lowercase());
    }
    "<none>".to_string()
}

#[cfg(test)]
mod tests {
    use super::*;
    use rpki::cir::{CanonicalInputRepresentation, CirObject, CirTrustAnchor, encode_cir};

    #[test]
    fn parse_args_accepts_required_flags() {
        let args = parse_args(&[
            "cir_probe_rpki_client_cache".to_string(),
            "--ours-cir".to_string(),
            "ours.cir".to_string(),
            "--rpki-client-cir".to_string(),
            "peer.cir".to_string(),
            "--cache-root".to_string(),
            "cache".to_string(),
            "--out-json".to_string(),
            "probe.json".to_string(),
            "--sample-limit".to_string(),
            "7".to_string(),
        ])
        .expect("parse args");
        assert_eq!(args.ours_cir.as_deref(), Some(Path::new("ours.cir")));
        assert_eq!(args.peer_cir.as_deref(), Some(Path::new("peer.cir")));
        assert_eq!(args.cache_root.as_deref(), Some(Path::new("cache")));
        assert_eq!(args.out_json.as_deref(), Some(Path::new("probe.json")));
        assert_eq!(args.sample_limit, 7);
    }

    #[test]
    fn parse_args_rejects_missing_invalid_and_unknown_flags() {
        let missing = parse_args(&[
            "cir_probe_rpki_client_cache".to_string(),
            "--ours-cir".to_string(),
            "ours.cir".to_string(),
            "--peer-cir".to_string(),
            "peer.cir".to_string(),
            "--cache-root".to_string(),
            "cache".to_string(),
        ])
        .unwrap_err();
        assert!(missing.contains("--out-json is required"), "{missing}");

        let invalid_limit = parse_args(&[
            "cir_probe_rpki_client_cache".to_string(),
            "--ours-cir".to_string(),
            "ours.cir".to_string(),
            "--peer-cir".to_string(),
            "peer.cir".to_string(),
            "--cache-root".to_string(),
            "cache".to_string(),
            "--out-json".to_string(),
            "probe.json".to_string(),
            "--sample-limit".to_string(),
            "NaN".to_string(),
        ])
        .unwrap_err();
        assert!(
            invalid_limit.contains("invalid --sample-limit"),
            "{invalid_limit}"
        );

        let unknown = parse_args(&[
            "cir_probe_rpki_client_cache".to_string(),
            "--unexpected".to_string(),
        ])
        .unwrap_err();
        assert!(unknown.contains("unknown argument"), "{unknown}");
    }

    #[test]
    fn run_probes_valid_temp_rrdp_missing_and_log_mentions() {
        let temp = tempfile::tempdir().expect("tempdir");
        let cache_root = temp.path().join("cache");
        let valid_bytes = b"valid object bytes";
        let rrdp_bytes = b"rrdp object bytes";
        let expected_mismatch_bytes = b"expected mismatch bytes";

        write_cache_file(
            &cache_root.join("rpki.example.test/repo/a.roa"),
            valid_bytes,
        );
        write_cache_file(
            &cache_root.join(".rsync/rpki.example.test/repo/b.mft"),
            b"different bytes",
        );
        write_cache_file(
            &cache_root.join(".rrdp/session-1/rpki.example.test/repo/c.cer"),
            rrdp_bytes,
        );

        let ours_cir = sample_cir(&[
            (
                "rsync://missing.example.test/repo/e.crl",
                b"missing bytes".as_slice(),
            ),
            (
                "rsync://rpki.example.test/repo/a.roa",
                valid_bytes.as_slice(),
            ),
            (
                "rsync://rpki.example.test/repo/b.mft",
                expected_mismatch_bytes.as_slice(),
            ),
            (
                "rsync://rpki.example.test/repo/c.cer",
                rrdp_bytes.as_slice(),
            ),
            (
                "rsync://shared.example.test/repo/d.roa",
                b"shared".as_slice(),
            ),
        ]);
        let peer_cir = sample_cir(&[(
            "rsync://shared.example.test/repo/d.roa",
            b"shared".as_slice(),
        )]);
        let ours = temp.path().join("ours.cir");
        let peer = temp.path().join("peer.cir");
        write_cir(&ours, &ours_cir);
        write_cir(&peer, &peer_cir);

        let log = temp.path().join("rpki-client.log");
        std::fs::write(
            &log,
            "rpki-client: deleted superfluous /cache/rpki.example.test/repo/b.mft\n",
        )
        .expect("write log");
        let out_json = temp.path().join("nested/probe.json");

        run(Args {
            ours_cir: Some(ours),
            peer_cir: Some(peer),
            cache_root: Some(cache_root),
            rpki_client_log: Some(log),
            out_json: Some(out_json.clone()),
            sample_limit: 10,
        })
        .expect("run");

        let summary: serde_json::Value =
            serde_json::from_slice(&std::fs::read(out_json).expect("read summary"))
                .expect("summary json");
        assert_eq!(summary["onlyInOursTotal"], 4);
        assert_eq!(summary["rrdpCacheDirs"], 1);
        assert_eq!(summary["cacheProbe"]["hashMatchCount"], 2);
        assert_eq!(summary["cacheProbe"]["hashMismatchCount"], 1);
        assert_eq!(summary["cacheProbe"]["missingCount"], 1);
        assert_eq!(summary["cacheProbe"]["validPathExistsCount"], 1);
        assert_eq!(summary["cacheProbe"]["rsyncTempPathExistsCount"], 1);
        assert_eq!(summary["cacheProbe"]["rrdpTempPathExistsCount"], 1);
        assert_eq!(summary["logProbe"]["enabled"], true);
        assert_eq!(summary["logProbe"]["relevantLineCount"], 1);
        assert_eq!(summary["logProbe"]["mentionedCount"], 1);
        assert_eq!(summary["samples"].as_array().unwrap().len(), 4);
    }

    #[test]
    fn run_reports_decode_failure_with_side_label() {
        let temp = tempfile::tempdir().expect("tempdir");
        let ours = temp.path().join("ours.cir");
        let peer = temp.path().join("peer.cir");
        std::fs::write(&ours, b"not a cir").expect("write invalid");
        write_cir(
            &peer,
            &sample_cir(&[(
                "rsync://shared.example.test/repo/d.roa",
                b"shared".as_slice(),
            )]),
        );
        let err = run(Args {
            ours_cir: Some(ours),
            peer_cir: Some(peer),
            cache_root: Some(temp.path().join("cache")),
            rpki_client_log: None,
            out_json: Some(temp.path().join("probe.json")),
            sample_limit: 20,
        })
        .unwrap_err();
        assert!(err.contains("decode ours CIR failed"), "{err}");
    }

    #[test]
    fn uri_helpers_extract_rsync_path_host_and_extension() {
        let uri = "rsync://rpki.example.test/repo/a/b.ROA";
        assert_eq!(strip_rsync_uri(uri), Some("rpki.example.test/repo/a/b.ROA"));
        assert_eq!(uri_host(uri), "rpki.example.test");
        assert_eq!(uri_extension(uri), ".roa");
        assert_eq!(file_name("rpki.example.test/repo/a/b.ROA"), "b.ROA");
    }

    #[test]
    fn log_mentions_match_cleanup_lines_by_uri_suffix() {
        let temp = tempfile::tempdir().expect("tempdir");
        let log = temp.path().join("run.log");
        std::fs::write(
            &log,
            "rpki-client: deleted superfluous .rrdp/abc/rpki.example.test/repo/a.roa\n",
        )
        .expect("write log");
        let objects = vec![ProbeObject {
            uri: "rsync://rpki.example.test/repo/a.roa".to_string(),
            sha256_hex: "00".repeat(32),
        }];
        let mentions = build_log_mentions(&log, &objects).expect("mentions");
        assert_eq!(mentions.relevant_line_count, 1);
        assert!(
            mentions
                .mentioned
                .contains("rsync://rpki.example.test/repo/a.roa")
        );
    }

    fn sample_cir(objects: &[(&str, &[u8])]) -> CanonicalInputRepresentation {
        let mut objects = objects
            .iter()
            .map(|(rsync_uri, bytes)| CirObject {
                rsync_uri: (*rsync_uri).to_string(),
                sha256: Sha256::digest(bytes).to_vec(),
            })
            .collect::<Vec<_>>();
        objects.sort_by(|left, right| left.rsync_uri.cmp(&right.rsync_uri));
        CanonicalInputRepresentation::new_v4(
            time::OffsetDateTime::UNIX_EPOCH,
            objects,
            Vec::new(),
            vec![sample_trust_anchor()],
            Vec::new(),
            Vec::new(),
        )
    }

    fn sample_trust_anchor() -> CirTrustAnchor {
        let ta_rsync_uri = "rsync://example.test/ta.cer";
        let ta_certificate_der = b"ta-der".to_vec();
        CirTrustAnchor {
            ta_rsync_uri: ta_rsync_uri.to_string(),
            tal_uri: "https://tal.example.test/apnic.tal".to_string(),
            tal_bytes: format!("{ta_rsync_uri}\n\nAQID\n").into_bytes(),
            ta_certificate_sha256: Sha256::digest(&ta_certificate_der).to_vec(),
            ta_certificate_der,
        }
    }

    fn write_cir(path: &Path, cir: &CanonicalInputRepresentation) {
        std::fs::write(path, encode_cir(cir).expect("encode cir")).expect("write cir");
    }

    fn write_cache_file(path: &Path, bytes: &[u8]) {
        std::fs::create_dir_all(path.parent().expect("parent")).expect("mkdir");
        std::fs::write(path, bytes).expect("write cache file");
    }
}