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

use rpki::ccr::{
    CcrContentInfo, VapCompareRow, VrpCompareRow, compare_state_digests, decode_ccr_compare_views,
    decode_content_info, write_vap_csv, write_vrp_csv,
};
use serde_json::json;

#[derive(Debug, Default, PartialEq, Eq)]
struct Args {
    ours_ccr: Option<PathBuf>,
    peer_ccr: Option<PathBuf>,
    out_json: Option<PathBuf>,
    out_md: Option<PathBuf>,
    out_dir: Option<PathBuf>,
    trust_anchor: String,
    fallback_compare_views: bool,
}

fn usage() -> &'static str {
    "Usage: ccr_state_compare --ours-ccr <path> --rpki-client-ccr <path> --out-json <path> [--out-md <path>] [--out-dir <path>] [--trust-anchor <name>] [--fallback-compare-views]"
}

fn parse_args(argv: &[String]) -> Result<Args, String> {
    let mut args = Args {
        trust_anchor: "unknown".to_string(),
        ..Args::default()
    };
    let mut i = 1usize;
    while i < argv.len() {
        match argv[i].as_str() {
            "--ours-ccr" => {
                i += 1;
                args.ours_ccr = Some(argv.get(i).ok_or("--ours-ccr requires a value")?.into());
            }
            "--rpki-client-ccr" | "--peer-ccr" => {
                i += 1;
                args.peer_ccr = Some(
                    argv.get(i)
                        .ok_or("--rpki-client-ccr requires a value")?
                        .into(),
                );
            }
            "--out-json" => {
                i += 1;
                args.out_json = Some(argv.get(i).ok_or("--out-json requires a value")?.into());
            }
            "--out-md" => {
                i += 1;
                args.out_md = Some(argv.get(i).ok_or("--out-md requires a value")?.into());
            }
            "--out-dir" => {
                i += 1;
                args.out_dir = Some(argv.get(i).ok_or("--out-dir requires a value")?.into());
            }
            "--trust-anchor" => {
                i += 1;
                args.trust_anchor = argv
                    .get(i)
                    .ok_or("--trust-anchor requires a value")?
                    .clone();
            }
            "--fallback-compare-views" => {
                args.fallback_compare_views = true;
            }
            "-h" | "--help" => return Err(usage().to_string()),
            other => return Err(format!("unknown argument: {other}\n{}", usage())),
        }
        i += 1;
    }
    if args.ours_ccr.is_none() {
        return Err(format!("--ours-ccr is required\n{}", usage()));
    }
    if args.peer_ccr.is_none() {
        return Err(format!("--rpki-client-ccr is required\n{}", usage()));
    }
    if args.out_json.is_none() {
        return Err(format!("--out-json is required\n{}", usage()));
    }
    Ok(args)
}

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

fn run(args: Args) -> Result<(), String> {
    let ours_path = args.ours_ccr.as_ref().unwrap();
    let peer_path = args.peer_ccr.as_ref().unwrap();
    let ours_der = read_file(ours_path)?;
    let peer_der = read_file(peer_path)?;

    let comparison = compare_state_digests(&ours_der, &peer_der).map_err(|e| e.to_string())?;
    let state_digest_match = comparison.matches();
    let mismatched_states = comparison.mismatched_state_names();
    let mut mismatched_components = Vec::new();
    if comparison.ours.version != comparison.peer.version {
        mismatched_components.push("version".to_string());
    }
    if comparison.ours.hash_alg_oid != comparison.peer.hash_alg_oid {
        mismatched_components.push("hashAlgorithm".to_string());
    }
    mismatched_components.extend(mismatched_states.iter().map(|name| (*name).to_string()));
    let run_fallback = args.fallback_compare_views && !state_digest_match;

    let fallback = if run_fallback {
        Some(build_compare_view_fallback(
            &ours_der,
            &peer_der,
            &args.trust_anchor,
            args.out_dir.as_deref(),
        )?)
    } else {
        None
    };

    let all_match = state_digest_match;

    let compare_path = if state_digest_match {
        "ccr_state_digest_match"
    } else if let Some(fallback) = fallback.as_ref() {
        if fallback.vrps.match_ && fallback.vaps.match_ {
            "ccr_state_digest_mismatch_with_compare_views_match"
        } else {
            "ccr_state_digest_mismatch_with_set_diff"
        }
    } else {
        "ccr_state_digest_mismatch"
    };

    let summary = json!({
        "comparePath": compare_path,
        "allMatch": all_match,
        "stateDigestMatch": state_digest_match,
        "mismatchedStates": mismatched_states,
        "mismatchedComponents": mismatched_components,
        "versionMatch": comparison.ours.version == comparison.peer.version,
        "hashAlgorithmMatch": comparison.ours.hash_alg_oid == comparison.peer.hash_alg_oid,
        "ours": {
            "version": comparison.ours.version,
            "hashAlg": comparison.ours.hash_alg_oid,
        },
        "rpkiClient": {
            "version": comparison.peer.version,
            "hashAlg": comparison.peer.hash_alg_oid,
        },
        "states": comparison.states.iter().map(|state| json!({
            "name": state.name,
            "match": state.matches,
            "oursPresent": state.ours_present,
            "rpkiClientPresent": state.peer_present,
            "oursHash": state.ours_hash_hex,
            "rpkiClientHash": state.peer_hash_hex,
        })).collect::<Vec<_>>(),
        "vrps": fallback.as_ref().map(|summary| summary.vrps.to_json()).unwrap_or_else(|| json!({
            "ours": serde_json::Value::Null,
            "rpkiClient": serde_json::Value::Null,
            "match": state_digest_match || !mismatched_states.contains(&"vrps"),
            "onlyInOurs": [],
            "onlyInRpkiClient": [],
        })),
        "vaps": fallback.as_ref().map(|summary| summary.vaps.to_json()).unwrap_or_else(|| json!({
            "ours": serde_json::Value::Null,
            "rpkiClient": serde_json::Value::Null,
            "match": state_digest_match || !mismatched_states.contains(&"vaps"),
            "onlyInOurs": [],
            "onlyInRpkiClient": [],
        })),
    });

    write_json(args.out_json.as_ref().unwrap(), &summary)?;
    if let Some(md_path) = args.out_md.as_ref() {
        write_markdown(md_path, &summary)?;
    }
    println!("{}", args.out_json.as_ref().unwrap().display());
    Ok(())
}

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 write_markdown(path: &Path, summary: &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()))?;
    }
    let lines = vec![
        "# CCR State Compare Summary".to_string(),
        String::new(),
        format!(
            "- `comparePath`: `{}`",
            summary["comparePath"].as_str().unwrap_or("-")
        ),
        format!(
            "- `allMatch`: `{}`",
            summary["allMatch"].as_bool().unwrap_or(false)
        ),
        format!(
            "- `stateDigestMatch`: `{}`",
            summary["stateDigestMatch"].as_bool().unwrap_or(false)
        ),
        format!(
            "- `mismatchedStates`: `{}`",
            summary["mismatchedStates"]
                .as_array()
                .map(|items| {
                    items
                        .iter()
                        .filter_map(|item| item.as_str())
                        .collect::<Vec<_>>()
                        .join(",")
                })
                .unwrap_or_default()
        ),
        format!(
            "- `vrpMatch`: `{}`",
            summary["vrps"]["match"].as_bool().unwrap_or(false)
        ),
        format!(
            "- `vapMatch`: `{}`",
            summary["vaps"]["match"].as_bool().unwrap_or(false)
        ),
    ];
    std::fs::write(path, lines.join("\n") + "\n")
        .map_err(|e| format!("write markdown failed: {}: {e}", path.display()))
}

struct FallbackSummary {
    vrps: SetSummary<VrpCompareRow>,
    vaps: SetSummary<VapCompareRow>,
}

struct SetSummary<T> {
    ours: usize,
    rpki_client: usize,
    match_: bool,
    only_in_ours: Vec<T>,
    only_in_rpki_client: Vec<T>,
}

impl SetSummary<VrpCompareRow> {
    fn to_json(&self) -> serde_json::Value {
        json!({
            "ours": self.ours,
            "rpkiClient": self.rpki_client,
            "match": self.match_,
            "onlyInOurs": self.only_in_ours.iter().map(vrp_row_json).collect::<Vec<_>>(),
            "onlyInRpkiClient": self.only_in_rpki_client.iter().map(vrp_row_json).collect::<Vec<_>>(),
        })
    }
}

impl SetSummary<VapCompareRow> {
    fn to_json(&self) -> serde_json::Value {
        json!({
            "ours": self.ours,
            "rpkiClient": self.rpki_client,
            "match": self.match_,
            "onlyInOurs": self.only_in_ours.iter().map(vap_row_json).collect::<Vec<_>>(),
            "onlyInRpkiClient": self.only_in_rpki_client.iter().map(vap_row_json).collect::<Vec<_>>(),
        })
    }
}

fn build_compare_view_fallback(
    ours_der: &[u8],
    peer_der: &[u8],
    trust_anchor: &str,
    out_dir: Option<&Path>,
) -> Result<FallbackSummary, String> {
    let ours = decode_content_info(ours_der).map_err(|e| e.to_string())?;
    let peer = decode_content_info(peer_der).map_err(|e| e.to_string())?;
    let (ours_vrps, ours_vaps) = decode_views(&ours, trust_anchor)?;
    let (peer_vrps, peer_vaps) = decode_views(&peer, trust_anchor)?;

    if let Some(out_dir) = out_dir {
        write_vrp_csv(&out_dir.join("ours-vrps.csv"), &ours_vrps)?;
        write_vap_csv(&out_dir.join("ours-vaps.csv"), &ours_vaps)?;
        write_vrp_csv(&out_dir.join("rpki-client-vrps.csv"), &peer_vrps)?;
        write_vap_csv(&out_dir.join("rpki-client-vaps.csv"), &peer_vaps)?;
    }

    Ok(FallbackSummary {
        vrps: compare_sets(&ours_vrps, &peer_vrps),
        vaps: compare_sets(&ours_vaps, &peer_vaps),
    })
}

fn decode_views(
    content_info: &CcrContentInfo,
    trust_anchor: &str,
) -> Result<(BTreeSet<VrpCompareRow>, BTreeSet<VapCompareRow>), String> {
    decode_ccr_compare_views(content_info, trust_anchor)
}

fn compare_sets<T: Ord + Clone>(ours: &BTreeSet<T>, rpki_client: &BTreeSet<T>) -> SetSummary<T> {
    SetSummary {
        ours: ours.len(),
        rpki_client: rpki_client.len(),
        match_: ours == rpki_client,
        only_in_ours: ours.difference(rpki_client).take(20).cloned().collect(),
        only_in_rpki_client: rpki_client.difference(ours).take(20).cloned().collect(),
    }
}

fn vrp_row_json(row: &VrpCompareRow) -> serde_json::Value {
    json!([row.asn, row.ip_prefix, row.max_length, row.trust_anchor])
}

fn vap_row_json(row: &VapCompareRow) -> serde_json::Value {
    json!([row.customer_asn, row.providers, row.trust_anchor])
}

#[cfg(test)]
mod tests {
    use super::*;
    use rpki::ccr::{
        CcrContentInfo, CcrDigestAlgorithm, ManifestState, RpkiCanonicalCacheRepresentation,
        build_aspa_payload_state, build_roa_payload_state, encode_content_info,
    };
    use rpki::data_model::roa::{IpPrefix, RoaAfi};
    use rpki::validation::objects::{AspaAttestation, Vrp};

    #[test]
    fn parse_args_accepts_required_flags() {
        let args = parse_args(&[
            "ccr_state_compare".to_string(),
            "--ours-ccr".to_string(),
            "ours.ccr".to_string(),
            "--rpki-client-ccr".to_string(),
            "peer.ccr".to_string(),
            "--out-json".to_string(),
            "summary.json".to_string(),
            "--fallback-compare-views".to_string(),
        ])
        .expect("parse args");
        assert_eq!(args.ours_ccr.as_deref(), Some(Path::new("ours.ccr")));
        assert_eq!(args.peer_ccr.as_deref(), Some(Path::new("peer.ccr")));
        assert_eq!(args.out_json.as_deref(), Some(Path::new("summary.json")));
        assert!(args.fallback_compare_views);
    }

    #[test]
    fn parse_args_rejects_missing_peer() {
        let err = parse_args(&[
            "ccr_state_compare".to_string(),
            "--ours-ccr".to_string(),
            "ours.ccr".to_string(),
            "--out-json".to_string(),
            "summary.json".to_string(),
        ])
        .unwrap_err();
        assert!(err.contains("--rpki-client-ccr is required"), "{err}");
    }

    #[test]
    fn run_reports_digest_match_without_fallback_counts() {
        let temp = tempfile::tempdir().expect("tempdir");
        let ccr = encode_content_info(&sample_content(64496)).expect("encode");
        let ours = temp.path().join("ours.ccr");
        let peer = temp.path().join("peer.ccr");
        let out = temp.path().join("summary.json");
        std::fs::write(&ours, &ccr).expect("write ours");
        std::fs::write(&peer, &ccr).expect("write peer");
        run(Args {
            ours_ccr: Some(ours),
            peer_ccr: Some(peer),
            out_json: Some(out.clone()),
            out_md: None,
            out_dir: Some(temp.path().join("compare")),
            trust_anchor: "apnic".to_string(),
            fallback_compare_views: true,
        })
        .expect("run");
        let summary: serde_json::Value =
            serde_json::from_slice(&std::fs::read(out).expect("read summary")).expect("json");
        assert_eq!(summary["comparePath"], "ccr_state_digest_match");
        assert_eq!(summary["allMatch"], true);
        assert!(summary["vrps"]["ours"].is_null());
        assert!(!temp.path().join("compare/ours-vrps.csv").exists());
    }

    #[test]
    fn run_reports_digest_mismatch_with_vrp_fallback() {
        let temp = tempfile::tempdir().expect("tempdir");
        let ours_ccr = encode_content_info(&sample_content(64496)).expect("encode ours");
        let peer_ccr = encode_content_info(&sample_content(64497)).expect("encode peer");
        let ours = temp.path().join("ours.ccr");
        let peer = temp.path().join("peer.ccr");
        let out = temp.path().join("summary.json");
        std::fs::write(&ours, &ours_ccr).expect("write ours");
        std::fs::write(&peer, &peer_ccr).expect("write peer");
        run(Args {
            ours_ccr: Some(ours),
            peer_ccr: Some(peer),
            out_json: Some(out.clone()),
            out_md: None,
            out_dir: Some(temp.path().join("compare")),
            trust_anchor: "apnic".to_string(),
            fallback_compare_views: true,
        })
        .expect("run");
        let summary: serde_json::Value =
            serde_json::from_slice(&std::fs::read(out).expect("read summary")).expect("json");
        assert_eq!(
            summary["comparePath"],
            "ccr_state_digest_mismatch_with_set_diff"
        );
        assert_eq!(summary["allMatch"], false);
        assert_eq!(summary["vrps"]["ours"], 1);
        assert_eq!(summary["vrps"]["rpkiClient"], 1);
        assert_eq!(summary["vrps"]["match"], false);
        assert!(temp.path().join("compare/ours-vrps.csv").exists());
    }

    #[test]
    fn run_reports_digest_mismatch_but_compare_views_match() {
        let temp = tempfile::tempdir().expect("tempdir");
        let ours_ccr = encode_content_info(&sample_content_with_manifest_hash(64496, 0x11))
            .expect("encode ours");
        let peer_ccr = encode_content_info(&sample_content_with_manifest_hash(64496, 0x22))
            .expect("encode peer");
        let ours = temp.path().join("ours.ccr");
        let peer = temp.path().join("peer.ccr");
        let out = temp.path().join("nested/summary.json");
        let out_md = temp.path().join("nested/summary.md");
        let out_dir = temp.path().join("compare");
        std::fs::write(&ours, &ours_ccr).expect("write ours");
        std::fs::write(&peer, &peer_ccr).expect("write peer");

        run(Args {
            ours_ccr: Some(ours),
            peer_ccr: Some(peer),
            out_json: Some(out.clone()),
            out_md: Some(out_md.clone()),
            out_dir: Some(out_dir.clone()),
            trust_anchor: "apnic".to_string(),
            fallback_compare_views: true,
        })
        .expect("run");

        let summary: serde_json::Value =
            serde_json::from_slice(&std::fs::read(out).expect("read summary")).expect("json");
        assert_eq!(
            summary["comparePath"],
            "ccr_state_digest_mismatch_with_compare_views_match"
        );
        assert_eq!(summary["allMatch"], false);
        assert_eq!(summary["stateDigestMatch"], false);
        assert_eq!(summary["mismatchedStates"], serde_json::json!(["mfts"]));
        assert_eq!(summary["vrps"]["ours"], 1);
        assert_eq!(summary["vrps"]["rpkiClient"], 1);
        assert_eq!(summary["vrps"]["match"], true);
        assert_eq!(summary["vaps"]["match"], true);
        assert!(out_dir.join("ours-vrps.csv").exists());
        assert!(
            std::fs::read_to_string(out_md)
                .expect("read markdown")
                .contains("compare_views_match")
        );
    }

    fn sample_content(asn: u32) -> CcrContentInfo {
        sample_content_with_manifest(asn, None)
    }

    fn sample_content_with_manifest_hash(asn: u32, manifest_hash_fill: u8) -> CcrContentInfo {
        sample_content_with_manifest(
            asn,
            Some(ManifestState {
                mis: Vec::new(),
                most_recent_update: time::OffsetDateTime::UNIX_EPOCH,
                hash: vec![manifest_hash_fill; 32],
            }),
        )
    }

    fn sample_content_with_manifest(asn: u32, mfts: Option<ManifestState>) -> CcrContentInfo {
        let vrps = build_roa_payload_state(&[Vrp {
            asn,
            prefix: IpPrefix {
                afi: RoaAfi::Ipv4,
                prefix_len: 24,
                addr: [192, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
            },
            max_length: 24,
        }])
        .expect("build vrps");
        let vaps = build_aspa_payload_state(&[AspaAttestation {
            customer_as_id: asn,
            provider_as_ids: vec![64497],
        }])
        .expect("build vaps");
        CcrContentInfo::new(RpkiCanonicalCacheRepresentation {
            version: 0,
            hash_alg: CcrDigestAlgorithm::Sha256,
            produced_at: time::OffsetDateTime::UNIX_EPOCH,
            mfts,
            vrps: Some(vrps),
            vaps: Some(vaps),
            tas: None,
            rks: None,
        })
    }
}