use ordered_float::OrderedFloat;
use std::{
	collections::{BinaryHeap, HashMap},
	hash::Hash,
	ops::Add,
};

/// Module for grid node implementation
/// Uses lifetime-bound references to the grid for obstacle checking
mod node {
	use crate::Pos;
	use crate::astar::OrderedFloat;

	use crate::Direction;
	use crate::Tile;
	use std::hash::Hash;

	/// Implementation of a simple grid node for demonstration
	#[derive(Clone, Copy)]
	pub struct NodeStruct<'a> {
		pub pos: Pos,
		grid: &'a [Tile],
	}
	impl<'a> NodeStruct<'a> {
		pub fn new(pos: Pos, grid: &'a [Tile]) -> Self {
			Self { pos, grid }
		}

		/// Step the node in a given direction, returning None if out of bounds or obstacle
		fn step(self, dir: Direction) -> Option<Self> {
			let pos = self.pos.step(dir)?;
			if self.grid[pos.idx()].is_wall() {
				return None;
			}

			Some(Self {
				pos,
				grid: self.grid,
			})
		}
	}

	impl PartialEq for NodeStruct<'_> {
		fn eq(&self, other: &Self) -> bool {
			self.pos == other.pos
		}
	}
	impl Eq for NodeStruct<'_> {}

	impl PartialOrd for NodeStruct<'_> {
		fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
			Some(self.cmp(other))
		}
	}
	impl Ord for NodeStruct<'_> {
		fn cmp(&self, other: &Self) -> std::cmp::Ordering {
			self.pos.cmp(&other.pos)
		}
	}

	impl Hash for NodeStruct<'_> {
		fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
			self.pos.hash(state);
		}
	}

	impl crate::astar::Node<OrderedFloat<f64>> for NodeStruct<'_> {
		/// Get neighbors of this node with their movement costs
		fn neighbors(&self) -> impl IntoIterator<Item = (Self, OrderedFloat<f64>)> {
			use Direction as D;
			[
				self.step(D::North),
				self.step(D::South),
				self.step(D::East),
				self.step(D::West),
			]
			.into_iter()
			.flatten()
			.map(|node| (node, OrderedFloat(1.0)))
		}
	}
}
use node::NodeStruct;

/// Trait representing a node in the graph
pub trait Node<C>
where
	Self: Sized,
{
	/// Get neighbors of this node with their movement costs
	/// Returns an iterator of (neighbor_node, cost) pairs
	fn neighbors(&self) -> impl IntoIterator<Item = (Self, C)>;
}

mod state {
	use super::Node;
	use std::cmp::Ordering;

	/// State wrapper for priority queue ordering
	pub struct State<N, C>
	where
		N: Node<C>,
	{
		pub cost: C,
		pub estimated_cost: C,
		pub node: N,
	}

	impl<N, C> Clone for State<N, C>
	where
		N: Node<C> + Clone,
		C: Clone,
	{
		fn clone(&self) -> Self {
			let cost = self.cost.clone();
			let estimated_cost = self.estimated_cost.clone();
			let node = self.node.clone();
			Self {
				cost,
				estimated_cost,
				node,
			}
		}
	}

	impl<N, C> Copy for State<N, C>
	where
		N: Node<C> + Copy,
		C: Copy,
	{
	}

	impl<N, C> PartialEq for State<N, C>
	where
		N: Node<C>,
		C: PartialEq,
	{
		fn eq(&self, other: &Self) -> bool {
			self.estimated_cost.eq(&other.estimated_cost) && self.cost.eq(&other.cost)
		}
	}

	impl<N, C> Eq for State<N, C>
	where
		N: Node<C>,
		C: Eq,
	{
	}

	impl<N, C> PartialOrd for State<N, C>
	where
		N: Node<C>,
		C: Ord,
	{
		fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
			Some(self.cmp(other))
		}
	}

	impl<N, C> Ord for State<N, C>
	where
		N: Node<C>,
		C: Ord,
	{
		fn cmp(&self, other: &Self) -> Ordering {
			other
				.estimated_cost
				.cmp(&self.estimated_cost)
				.then_with(|| other.cost.cmp(&self.cost))
		}
	}
}
use state::State;

use crate::{Pos, Tile};

pub fn find_path(grid: &[Tile], start: Pos, goal: Pos) -> Option<(Vec<Pos>, f64)> {
	// Create nodes for astar
	let start_node = NodeStruct::new(start, grid);
	let goal_node = NodeStruct::new(goal, grid);

	// Use closure for heuristic with lifetime-bound reference
	let heuristic = |node: &NodeStruct| {
		let abs_pos = node.pos.abs_diff(goal);
		OrderedFloat((abs_pos.x() + abs_pos.y()) as f64)
	};

	// Call A* algorithm
	astar(&start_node, &goal_node, heuristic).map(|(path, cost)| {
		(
			path.into_iter().map(|node| node.pos).collect(),
			cost.into_inner(),
		)
	})
}

/// A* pathfinding implementation
///
/// # Type Parameters
/// * `N` - Node type implementing the Node trait
/// * `H` - Heuristic function type (closure)
///
/// # Arguments
/// * `start` - Starting node (owned)
/// * `goal` - Goal node (borrowed reference)
/// * `heuristic` - Heuristic function estimating cost from any node to goal
///
/// # Returns
/// * `Some((path, total_cost))` if path found
/// * `None` if no path exists
fn astar<N, C, H>(start: &N, goal: &N, heuristic: H) -> Option<(Vec<N>, C)>
where
	N: Node<C> + Eq + Hash + Clone,
	H: Fn(&N) -> C,
	C: Default + Add<Output = C> + Ord + Copy,
{
	// list of nodes we need to look at next (BinaryHeap for priority queue behavior)
	let mut open_set = BinaryHeap::new();

	// The best known source node to get to a
	// destination node with minimum weight cost (HashMap for O(1) access)
	let mut came_from: HashMap<N, N> = HashMap::new();

	// The current known cost to get to a given node from the start node
	// (HashMap for O(1) access)
	let mut known_cost_map: HashMap<N, C> = HashMap::new();

	// Initialize with start node (add to known cost map and nodes to explore)
	known_cost_map.insert(start.clone(), C::default());
	open_set.push(State {
		cost: C::default(),
		estimated_cost: C::default(),
		node: start.clone(),
	});

	// Explore the nodes until we run out of nodes
	while let Some(State { cost, node, .. }) = open_set.pop() {
		// Goal reached! Reconstruct the path
		if &node == goal {
			return Some((reconstruct_path(&came_from, node), cost));
		}

		// Explore neighbors
		for (neighbor, move_cost) in node.neighbors() {
			let tentative_cost = cost + move_cost;

			// If the path we are about to take is worse then the best known so far
			// for out neighbor, then we should not take it!
			if let Some(known_cost) = known_cost_map.get(&neighbor)
				&& *known_cost <= tentative_cost
			{
				continue;
			}

			// We have found a better path, we should update, so
			// our structures to reflect this
			came_from.insert(neighbor.clone(), node.clone());
			known_cost_map.insert(neighbor.clone(), tentative_cost);
			open_set.push(State {
				cost: tentative_cost,
				estimated_cost: tentative_cost + heuristic(&neighbor),
				node: neighbor,
			});
		}
	}

	None // No path found
}

/// Reconstruct path from goal to start using came_from map
fn reconstruct_path<N, C>(came_from: &HashMap<N, N>, mut current: N) -> Vec<N>
where
	N: Node<C> + Eq + Hash + Clone,
{
	let mut path = vec![current.clone()];

	// Backtrack from goal to start
	while let Some(prev) = came_from.get(&current) {
		current = prev.clone();
		path.push(current.clone());
	}

	path.reverse();
	path
}