//! A* Pathfinding Algorithm for CFG with 3D Spatial Heuristics

use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap, HashSet};

#[derive(Debug, Clone)]
pub struct CfgGraphNode {
    pub id: u64,
    pub x: f32,
    pub y: f32,
    pub z: f32,
    pub successors: Vec<u64>,
}

#[derive(Debug, Clone)]
pub struct CfgPath {
    pub node_ids: Vec<u64>,
    pub total_cost: f32,
    pub heuristic_cost: f32,
    pub actual_cost: f32,
}

#[derive(Debug, Clone, Copy, PartialEq)]
struct AStarNode {
    node_id: u64,
    f_score: f32,
    g_score: f32,
}

impl Eq for AStarNode {}

impl Ord for AStarNode {
    fn cmp(&self, other: &Self) -> Ordering {
        other
            .f_score
            .partial_cmp(&self.f_score)
            .unwrap_or(Ordering::Equal)
    }
}

impl PartialOrd for AStarNode {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[inline]
fn euclidean_distance(x1: f32, y1: f32, z1: f32, x2: f32, y2: f32, z2: f32) -> f32 {
    let dx = x2 - x1;
    let dy = y2 - y1;
    let dz = z2 - z1;
    (dx * dx + dy * dy + dz * dz).sqrt()
}

pub fn astar_find_path(nodes: &[CfgGraphNode], start_id: u64, goal_id: u64) -> Option<CfgPath> {
    eprintln!(
        "DEBUG astar_find_path: start_id={}, goal_id={}, nodes.len()={}",
        start_id,
        goal_id,
        nodes.len()
    );
    let node_map: HashMap<u64, &CfgGraphNode> = nodes.iter().map(|n| (n.id, n)).collect();
    let start_node = node_map.get(&start_id)?;

    // Special case: goal_id = u64::MAX means "find path to farthest node"
    let actual_goal_id = if goal_id == u64::MAX {
        // Find the farthest node from start
        let mut farthest_id = start_id;
        let mut max_dist = 0.0f32;
        for (&id, node) in &node_map {
            if id == start_id {
                continue;
            }
            let dist = euclidean_distance(
                start_node.x,
                start_node.y,
                start_node.z,
                node.x,
                node.y,
                node.z,
            );
            if dist > max_dist {
                max_dist = dist;
                farthest_id = id;
            }
        }
        eprintln!(
            "DEBUG: start_id={}, goal_id=MAX, farthest_id={}, max_dist={}",
            start_id, farthest_id, max_dist
        );
        farthest_id
    } else {
        goal_id
    };

    let goal_node = node_map.get(&actual_goal_id)?;

    let mut open_set = BinaryHeap::new();
    let mut in_open_set: HashSet<u64> = HashSet::new();
    let mut closed_set: HashSet<u64> = HashSet::new();
    let mut g_score: HashMap<u64, f32> = HashMap::new();
    let mut came_from: HashMap<u64, u64> = HashMap::new();

    let initial_h = euclidean_distance(
        start_node.x,
        start_node.y,
        start_node.z,
        goal_node.x,
        goal_node.y,
        goal_node.z,
    );
    g_score.insert(start_id, 0.0);
    open_set.push(AStarNode {
        node_id: start_id,
        f_score: initial_h,
        g_score: 0.0,
    });
    in_open_set.insert(start_id);

    while let Some(current) = open_set.pop() {
        let current_id = current.node_id;
        if closed_set.contains(&current_id) {
            continue;
        }
        if current_id == actual_goal_id {
            let mut path = vec![actual_goal_id];
            let mut cur = actual_goal_id;
            while let Some(&prev) = came_from.get(&cur) {
                path.push(prev);
                cur = prev;
            }
            path.reverse();
            return Some(CfgPath {
                node_ids: path,
                total_cost: current.f_score,
                heuristic_cost: initial_h,
                actual_cost: current.g_score,
            });
        }
        in_open_set.remove(&current_id);
        closed_set.insert(current_id);
        let current_node = node_map.get(&current_id)?;
        eprintln!(
            "DEBUG: Processing node_id={}, successors={:?}",
            current_id, current_node.successors
        );
        for &succ_id in &current_node.successors {
            if closed_set.contains(&succ_id) {
                continue;
            }
            let succ_node = match node_map.get(&succ_id) {
                Some(n) => n,
                None => continue,
            };
            let tentative_g = current.g_score + 1.0;
            let existing_g = g_score.get(&succ_id).copied().unwrap_or(f32::INFINITY);
            if tentative_g < existing_g {
                came_from.insert(succ_id, current_id);
                g_score.insert(succ_id, tentative_g);
                let h = euclidean_distance(
                    succ_node.x,
                    succ_node.y,
                    succ_node.z,
                    goal_node.x,
                    goal_node.y,
                    goal_node.z,
                );
                if !in_open_set.contains(&succ_id) {
                    open_set.push(AStarNode {
                        node_id: succ_id,
                        f_score: tentative_g + h,
                        g_score: tentative_g,
                    });
                    in_open_set.insert(succ_id);
                }
            }
        }
    }
    None
}

pub fn astar_find_k_paths(
    nodes: &[CfgGraphNode],
    start_id: u64,
    goal_id: u64,
    k: usize,
) -> Vec<CfgPath> {
    let mut paths: Vec<CfgPath> = Vec::new();
    if let Some(first) = astar_find_path(nodes, start_id, goal_id) {
        paths.push(first);
    } else {
        return paths;
    }
    let mut candidates: Vec<CfgPath> = Vec::new();
    for i in 0..(k - 1).min(paths.len()) {
        let base = &paths[i];
        for j in 0..(base.node_ids.len() - 1) {
            let spur = base.node_ids[j];
            let modified: Vec<CfgGraphNode> = nodes
                .iter()
                .map(|n| {
                    let mut m = n.clone();
                    if j > 0 && n.id == base.node_ids[j - 1] {
                        m.successors.retain(|&s| s != spur);
                    }
                    m
                })
                .collect();
            if let Some(spur_path) = astar_find_path(&modified, spur, goal_id) {
                let mut combined: Vec<u64> = base.node_ids[..=j].to_vec();
                combined.extend_from_slice(&spur_path.node_ids[1..]);
                let cp = CfgPath {
                    node_ids: combined,
                    total_cost: spur_path.total_cost + j as f32,
                    heuristic_cost: spur_path.heuristic_cost,
                    actual_cost: spur_path.actual_cost + j as f32,
                };
                if !candidates.iter().any(|p| p.node_ids == cp.node_ids) {
                    candidates.push(cp);
                }
            }
        }
        candidates.sort_by(|a, b| {
            a.total_cost
                .partial_cmp(&b.total_cost)
                .unwrap_or(Ordering::Equal)
        });
        if let Some(best) = candidates.pop() {
            paths.push(best);
        }
    }
    paths
}

#[derive(Debug, Clone)]
pub struct PathComplexity {
    pub path_length: usize,
    pub max_loop_depth: u32,
    pub branch_count: usize,
    pub total_spatial_distance: f32,
    pub heuristic_efficiency: f32,
}

pub fn analyze_path_complexity(path: &CfgPath, nodes: &[CfgGraphNode]) -> PathComplexity {
    let node_map: HashMap<u64, &CfgGraphNode> = nodes.iter().map(|n| (n.id, n)).collect();
    let mut max_ld = 0.0f32;
    let mut bc = 0;
    let mut tsd = 0.0f32;
    for i in 0..path.node_ids.len() {
        if let Some(&n) = node_map.get(&path.node_ids[i]) {
            max_ld = max_ld.max(n.y);
            if n.z > 0.0 {
                bc += 1;
            }
            if i > 0 {
                if let Some(&p) = node_map.get(&path.node_ids[i - 1]) {
                    tsd += euclidean_distance(p.x, p.y, p.z, n.x, n.y, n.z);
                }
            }
        }
    }
    PathComplexity {
        path_length: path.node_ids.len(),
        max_loop_depth: max_ld as u32,
        branch_count: bc,
        total_spatial_distance: tsd,
        heuristic_efficiency: if path.heuristic_cost > 0.0 {
            path.actual_cost / path.heuristic_cost
        } else {
            1.0
        },
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_astar_simple() {
        let n = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![2],
            },
            CfgGraphNode {
                id: 2,
                x: 2.0,
                y: 0.0,
                z: 0.0,
                successors: vec![3],
            },
            CfgGraphNode {
                id: 3,
                x: 3.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];
        assert_eq!(
            astar_find_path(&n, 0, 3).unwrap().node_ids,
            vec![0, 1, 2, 3]
        );
    }
    #[test]
    fn test_astar_branch() {
        let n = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1, 2],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 1.0,
                successors: vec![3],
            },
            CfgGraphNode {
                id: 2,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![3],
            },
            CfgGraphNode {
                id: 3,
                x: 2.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];
        assert!(astar_find_path(&n, 0, 3).is_some());
    }
    #[test]
    fn test_astar_nopath() {
        let n = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];
        assert!(astar_find_path(&n, 0, 2).is_none());
    }
    #[test]
    fn test_astar_kpaths() {
        let n = vec![
            CfgGraphNode {
                id: 0,
                x: 0.0,
                y: 0.0,
                z: 0.0,
                successors: vec![1, 2, 3],
            },
            CfgGraphNode {
                id: 1,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![4],
            },
            CfgGraphNode {
                id: 2,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![4],
            },
            CfgGraphNode {
                id: 3,
                x: 1.0,
                y: 0.0,
                z: 0.0,
                successors: vec![4],
            },
            CfgGraphNode {
                id: 4,
                x: 2.0,
                y: 0.0,
                z: 0.0,
                successors: vec![],
            },
        ];
        let paths = astar_find_k_paths(&n, 0, 4, 3);
        assert!(paths.len() >= 2 && paths.len() <= 3);
    }
}
