mod loader;

use cch_rs::cch::*;
use petgraph::algo::dijkstra;
use rand::{RngExt, SeedableRng};
use rand_chacha::ChaCha8Rng;
use rustc_hash::FxHashMap;

const GR_PATH: &str = "tests/data/USA-road-t.COL.gr";
const CO_PATH: &str = "tests/data/USA-road-d.COL.co";
const TEST_PROFILE: u8 = 0;
const NUM_TESTS: usize = 50;

#[derive(Clone)]
struct ToyGraph {
    first_out: Vec<u32>,
    head: Vec<u32>,
    weights: Vec<f32>,
    xs: Vec<f32>,
    ys: Vec<f32>,
}

impl CchGraph for ToyGraph {
    fn num_nodes(&self) -> usize { self.first_out.len() - 1 }
    fn num_edges(&self) -> usize { self.head.len() }
    fn first_out(&self) -> &[u32] { &self.first_out }
    fn head(&self) -> &[u32] { &self.head }
    fn x(&self, node_id: u32) -> f32 { self.xs[node_id as usize] }
    fn y(&self, node_id: u32) -> f32 { self.ys[node_id as usize] }
}

fn toy_graph() -> ToyGraph {
    ToyGraph {
        first_out: vec![0, 3, 6, 8, 10],
        head: vec![1, 2, 3, 0, 2, 3, 1, 3, 0, 2],
        weights: vec![1.0, 5.0, 12.0, 1.0, 1.0, 5.0, 1.0, 1.0, 12.0, 1.0],
        xs: vec![0.0; 4],
        ys: vec![0.0; 4],
    }
}

fn toy_petgraph(graph: &ToyGraph, weights: &[f32]) -> petgraph::Graph<(), f32, petgraph::Directed> {
    let mut g = petgraph::Graph::new();
    let nodes: Vec<_> = (0..graph.num_nodes()).map(|_| g.add_node(())).collect();
    for u in 0..graph.num_nodes() {
        for edge_idx in graph.edge_indices(u) {
            g.add_edge(nodes[u], nodes[graph.head[edge_idx] as usize], weights[edge_idx]);
        }
    }
    g
}

fn path_weight<G: CchGraph>(graph: &G, weights: &[f32], path: &[u32]) -> f32 {
    path.windows(2).fold(0.0, |acc, edge| {
        let u = edge[0] as usize;
        let v = edge[1];
        let w = graph
            .edge_indices(u)
            .find_map(|edge_idx| (graph.head()[edge_idx] == v).then_some(weights[edge_idx]))
            .expect("Path must follow original graph edges");
        acc + w
    })
}

#[test]
pub fn verify_cch_against_dijkstra() {
    if !std::path::Path::new(GR_PATH).exists() || !std::path::Path::new(CO_PATH).exists() {
        println!("cargo:warning=Test data not found at {}, {}. Skipping heavy test.", GR_PATH, CO_PATH);
        return;
    }
    let graph = loader::DimacsLoader::load_dimacs(GR_PATH, CO_PATH).expect("Failed to load data");
    let pet_graph = loader::DimacsLoader::to_petgraph(&graph);
    let mut profile_weights = FxHashMap::default();
    profile_weights.insert(TEST_PROFILE, graph.weights.clone());
    let engine = CchEngine::new(&graph, &profile_weights);
    let num_nodes = graph.num_nodes() as u32;
    let mut rng = ChaCha8Rng::seed_from_u64(42);

    for i in 0..NUM_TESTS {
        let start: u32 = rng.random_range(0..num_nodes);
        let end: u32 = rng.random_range(0..num_nodes);
        if start == end {
            continue;
        }
        let cch_res = engine.query_path(TEST_PROFILE, start, end);
        let cch_dist = cch_res.as_ref().map(|r| r.distance).unwrap_or(f32::INFINITY);
        let start_idx = pet_graph.node_indices().nth(start as usize).unwrap();
        let end_idx = pet_graph.node_indices().nth(end as usize).unwrap();
        let dijkstra_map = dijkstra(&pet_graph, start_idx, Some(end_idx), |e| *e.weight());
        let dijkstra_dist = *dijkstra_map.get(&end_idx).unwrap_or(&f32::INFINITY);
        if dijkstra_dist.is_infinite() {
            assert!(cch_dist.is_infinite(), "Dijkstra says unreachable, but CCH found a path!");
        } else {
            let diff = (cch_dist - dijkstra_dist).abs();
            assert!(diff < 0.1, "Test #{} FAILED: {} -> {}. CCH: {}, Dijkstra: {}, Diff: {}", i, start, end, cch_dist, dijkstra_dist, diff);
        }
    }
}

#[test]
pub fn diagnostic_cch_flow() {
    if !std::path::Path::new(GR_PATH).exists() || !std::path::Path::new(CO_PATH).exists() {
        println!("cargo:warning=Test data not found at {}, {}. Skipping heavy test.", GR_PATH, CO_PATH);
        return;
    }
    let graph = loader::DimacsLoader::load_dimacs(GR_PATH, CO_PATH).expect("Failed to load data");
    let mut profile_weights = FxHashMap::default();
    profile_weights.insert(TEST_PROFILE, graph.weights.clone());
    let engine = CchEngine::new(&graph, &profile_weights);
    let data = engine.data.load();
    let cch = &data.cch;
    let profile_data = data.profiles.get(&TEST_PROFILE).unwrap().read();
    println!("Checking Rank/Order consistency...");
    for i in 0..graph.num_nodes() {
        let r = cch.ranks[i];
        if cch.order[r as usize] != i as u32 {
            panic!("CONSISTENCY ERROR: Node {} maps to rank {}, but rank {} maps back to node {}", i, r, r, cch.order[r as usize]);
        }
    }
    println!("Checking Topology completeness...");
    let mut missing_edges = 0;
    for u in 0..graph.num_nodes() {
        for v in graph.neighbors(u as u32) {
            let (r_u, r_v) = (cch.ranks[u], cch.ranks[v as usize]);
            let (low, high) = if r_u < r_v { (r_u, r_v) } else { (r_v, r_u) };
            if unsafe { cch.topology.find_edge_id(low, high).is_none() } {
                missing_edges += 1;
            }
        }
    }
    if missing_edges > 0 {
        panic!("TOPOLOGY ERROR: {} edges from the original graph are missing in CCH!", missing_edges);
    }
    let start_node = 97624;
    let end_node = 297078;
    let pet_graph = loader::DimacsLoader::to_petgraph(&graph);
    let start_idx = pet_graph.node_indices().nth(start_node).unwrap();
    let end_idx = pet_graph.node_indices().nth(end_node).unwrap();
    let (dijkstra_dist, path_nodes) = dijkstra_with_path(&pet_graph, start_idx, end_idx);
    println!("Dijkstra Distance: {}", dijkstra_dist);
    println!("Verifying path edges in CCH...");
    for window in path_nodes.windows(2) {
        let u = window[0];
        let v = window[1];
        let r_u = cch.ranks[u as usize];
        let r_v = cch.ranks[v as usize];
        let (low, high) = if r_u < r_v { (r_u, r_v) } else { (r_v, r_u) };
        match unsafe { cch.topology.find_edge_id(low, high) } {
            Some(edge_id) => {
                let weight = if r_u < r_v {
                    profile_data.weights.up[edge_id as usize]
                } else {
                    profile_data.weights.down[edge_id as usize]
                };
                if weight > dijkstra_dist {
                    println!("  Edge {}->{} (Rank {}->{}) has INF or WRONG weight in CCH!", u, v, r_u, r_v);
                }
            },
            None => println!("  Edge {}->{} (Rank {}->{}) MISSING in Topology!", u, v, r_u, r_v),
        }
    }
    let cch_res = engine.query_path(TEST_PROFILE, start_node as u32, end_node as u32).unwrap();
    println!("CCH Distance: {}", cch_res.distance);
}

fn dijkstra_with_path(g: &petgraph::Graph<(), f32, petgraph::Directed>, start: petgraph::graph::NodeIndex, end: petgraph::graph::NodeIndex) -> (f32, Vec<u32>) {
    let res = petgraph::algo::dijkstra(g, start, Some(end), |e| *e.weight());
    let dist = *res.get(&end).unwrap();
    (dist, vec![])
}

#[test]
fn partial_weight_update_matches_full_customization() {
    let graph = toy_graph();
    let mut profile_weights = FxHashMap::default();
    profile_weights.insert(TEST_PROFILE, graph.weights.clone());

    let partial_engine = CchEngine::new(&graph, &profile_weights);
    let full_engine = CchEngine::new(&graph, &profile_weights);

    let mut new_weights = graph.weights.clone();
    new_weights[4] = 10.0;

    partial_engine.update_weights_partial(TEST_PROFILE, &[(4, 10.0)]);
    full_engine.update_weights(TEST_PROFILE, &new_weights);

    let pet_graph = toy_petgraph(&graph, &new_weights);
    for from in 0..graph.num_nodes() as u32 {
        for to in 0..graph.num_nodes() as u32 {
            if from == to {
                continue;
            }

            let partial = partial_engine.query_path(TEST_PROFILE, from, to);
            let full = full_engine.query_path(TEST_PROFILE, from, to);

            let partial_dist = partial.as_ref().map(|r| r.distance).unwrap_or(f32::INFINITY);
            let full_dist = full.as_ref().map(|r| r.distance).unwrap_or(f32::INFINITY);
            assert!((partial_dist - full_dist).abs() < 1e-4, "Distance mismatch for {from}->{to}: partial={partial_dist}, full={full_dist}");

            let start = pet_graph.node_indices().nth(from as usize).unwrap();
            let end = pet_graph.node_indices().nth(to as usize).unwrap();
            let dijkstra_dist = *dijkstra(&pet_graph, start, Some(end), |e| *e.weight()).get(&end).unwrap_or(&f32::INFINITY);
            assert!(
                (partial_dist - dijkstra_dist).abs() < 1e-4,
                "Dijkstra mismatch for {from}->{to}: partial={partial_dist}, dijkstra={dijkstra_dist}"
            );

            if let Some(result) = partial {
                let path_dist = path_weight(&graph, &new_weights, &result.path);
                assert!(
                    (path_dist - result.distance).abs() < 1e-4,
                    "Path unpack mismatch for {from}->{to}: path={path_dist}, query={}",
                    result.distance
                );
            }
        }
    }
}
