use crate::error::{GeomError, Result};
use rtx_tensor::Tensor;
use slotmap::{DefaultKey, Key, SlotMap};
use std::collections::HashMap;

/// Unique identifier for nodes in the graph
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct NodeId(DefaultKey);

impl NodeId {
    pub fn from_raw(idx: usize) -> Self {
        // This is unsafe but needed for testing invalid node IDs
        unsafe { std::mem::transmute(idx as u64) }
    }

    pub fn index(&self) -> usize {
        self.0.data().as_ffi() as usize
    }
}

/// Unique identifier for edges in the graph
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct EdgeId(DefaultKey);

impl EdgeId {
    pub fn index(&self) -> usize {
        self.0.data().as_ffi() as usize
    }
}

/// A node in the graph with associated features
#[derive(Debug, Clone)]
pub struct Node {
    features: Tensor,
    metadata: HashMap<String, String>,
}

impl Node {
    pub fn new(features: Tensor) -> Self {
        Self {
            features,
            metadata: HashMap::new(),
        }
    }

    pub fn with_metadata(features: Tensor, metadata: HashMap<String, String>) -> Self {
        Self { features, metadata }
    }

    pub fn features(&self) -> &Tensor {
        &self.features
    }

    pub fn features_mut(&mut self) -> &mut Tensor {
        &mut self.features
    }

    pub fn metadata(&self) -> &HashMap<String, String> {
        &self.metadata
    }

    pub fn set_metadata(&mut self, key: String, value: String) {
        self.metadata.insert(key, value);
    }
}

/// An edge in the graph with weight and optional features
#[derive(Debug, Clone)]
pub struct Edge {
    weight: Tensor,
    features: Option<Tensor>,
    metadata: HashMap<String, String>,
}

impl Edge {
    pub fn new(weight: Tensor) -> Self {
        Self {
            weight,
            features: None,
            metadata: HashMap::new(),
        }
    }

    pub fn with_features(weight: Tensor, features: Tensor) -> Self {
        Self {
            weight,
            features: Some(features),
            metadata: HashMap::new(),
        }
    }

    pub fn weight(&self) -> &Tensor {
        &self.weight
    }

    pub fn weight_mut(&mut self) -> &mut Tensor {
        &mut self.weight
    }

    pub fn features(&self) -> Option<&Tensor> {
        self.features.as_ref()
    }

    pub fn features_mut(&mut self) -> Option<&mut Tensor> {
        self.features.as_mut()
    }

    pub fn metadata(&self) -> &HashMap<String, String> {
        &self.metadata
    }

    pub fn set_metadata(&mut self, key: String, value: String) {
        self.metadata.insert(key, value);
    }
}

/// Graph data structure for storing nodes and edges
#[derive(Debug, Clone)]
pub struct Graph {
    nodes: SlotMap<DefaultKey, Node>,
    edges: SlotMap<DefaultKey, (NodeId, NodeId, Edge)>,
    incoming_adj: HashMap<NodeId, Vec<(NodeId, EdgeId)>>, // incoming edges: to <- [(from, edge_id)]
    outgoing_adj: HashMap<NodeId, Vec<(NodeId, EdgeId)>>, // outgoing edges: from -> [(to, edge_id)]
}

impl Graph {
    pub fn new() -> Self {
        Self {
            nodes: SlotMap::new(),
            edges: SlotMap::new(),
            incoming_adj: HashMap::new(),
            outgoing_adj: HashMap::new(),
        }
    }

    pub fn add_node(&mut self, node: Node) -> NodeId {
        let key = self.nodes.insert(node);

        NodeId(key)
    }

    pub fn add_edge(&mut self, from: NodeId, to: NodeId, edge: Edge) -> Result<EdgeId> {
        // Validate that both nodes exist
        if !self.nodes.contains_key(from.0) {
            return Err(GeomError::InvalidNodeId(from.index()));
        }
        if !self.nodes.contains_key(to.0) {
            return Err(GeomError::InvalidNodeId(to.index()));
        }

        let key = self.edges.insert((from, to, edge));
        let edge_id = EdgeId(key);

        // Update both adjacency lists
        self.incoming_adj
            .entry(to)
            .or_default()
            .push((from, edge_id));
        self.outgoing_adj
            .entry(from)
            .or_default()
            .push((to, edge_id));

        Ok(edge_id)
    }

    pub fn node(&self, id: NodeId) -> Option<&Node> {
        self.nodes.get(id.0)
    }

    pub fn node_mut(&mut self, id: NodeId) -> Option<&mut Node> {
        self.nodes.get_mut(id.0)
    }

    pub fn edge(&self, id: EdgeId) -> Option<&Edge> {
        self.edges.get(id.0).map(|(_, _, edge)| edge)
    }

    pub fn edge_mut(&mut self, id: EdgeId) -> Option<&mut Edge> {
        self.edges.get_mut(id.0).map(|(_, _, edge)| edge)
    }

    pub fn edge_endpoints(&self, id: EdgeId) -> Option<(NodeId, NodeId)> {
        self.edges.get(id.0).map(|(from, to, _)| (*from, *to))
    }

    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }

    pub fn edge_count(&self) -> usize {
        self.edges.len()
    }

    pub fn nodes(&self) -> Vec<NodeId> {
        self.nodes.keys().map(NodeId).collect()
    }

    pub fn edges(&self) -> Vec<EdgeId> {
        self.edges.keys().map(EdgeId).collect()
    }

    /// Get outgoing neighbors (nodes that this node has edges TO)
    pub fn neighbors(&self, node: NodeId) -> Vec<NodeId> {
        self.outgoing_adj
            .get(&node)
            .map(|neighbors| {
                neighbors
                    .iter()
                    .map(|(neighbor_id, _)| *neighbor_id)
                    .collect()
            })
            .unwrap_or_default()
    }

    /// Get incoming neighbors (nodes that have edges TO this node) - used for message passing
    pub fn incoming_neighbors(&self, node: NodeId) -> Vec<NodeId> {
        self.incoming_adj
            .get(&node)
            .map(|neighbors| {
                neighbors
                    .iter()
                    .map(|(neighbor_id, _)| *neighbor_id)
                    .collect()
            })
            .unwrap_or_default()
    }

    pub fn neighbor_edges(&self, node: NodeId) -> Vec<EdgeId> {
        self.outgoing_adj
            .get(&node)
            .map(|neighbors| neighbors.iter().map(|(_, edge_id)| *edge_id).collect())
            .unwrap_or_default()
    }

    pub fn incoming_neighbor_edges(&self, node: NodeId) -> Vec<EdgeId> {
        self.incoming_adj
            .get(&node)
            .map(|neighbors| neighbors.iter().map(|(_, edge_id)| *edge_id).collect())
            .unwrap_or_default()
    }

    /// Generate adjacency matrix as a tensor
    pub fn adjacency_matrix(&self) -> Tensor {
        let n = self.node_count();
        let mut matrix_data = vec![0.0f32; n * n];

        // Create mapping from NodeId to matrix index
        let node_to_index: HashMap<NodeId, usize> = self
            .nodes()
            .into_iter()
            .enumerate()
            .map(|(idx, node_id)| (node_id, idx))
            .collect();

        // Fill the adjacency matrix
        for (from_id, neighbors) in &self.outgoing_adj {
            if let Some(&from_idx) = node_to_index.get(from_id) {
                for &(to_id, edge_id) in neighbors {
                    if let Some(&to_idx) = node_to_index.get(&to_id)
                        && let Some(edge) = self.edge(edge_id)
                    {
                        let weight = edge.weight().to_cpu().unwrap()[0];
                        matrix_data[from_idx * n + to_idx] = weight;
                    }
                }
            }
        }

        Tensor::from_data(matrix_data, [n, n], &rtx_tensor::Device::default()).unwrap()
    }

    /// Get node features as a batch tensor
    pub fn node_features(&self) -> Option<Tensor> {
        if self.nodes.is_empty() {
            return None;
        }

        let first_node = self.nodes.values().next().unwrap();
        let feature_dim = first_node.features().shape().dims()[0];
        let num_nodes = self.node_count();

        let mut features_data = Vec::with_capacity(num_nodes * feature_dim);

        // Collect features in consistent order
        for node_id in self.nodes() {
            let node = self.node(node_id).unwrap();
            let node_features = node.features().to_cpu().unwrap();
            features_data.extend(node_features);
        }

        Some(
            Tensor::from_data(
                features_data,
                [num_nodes, feature_dim],
                &rtx_tensor::Device::default(),
            )
            .unwrap(),
        )
    }
}

impl Default for Graph {
    fn default() -> Self {
        Self::new()
    }
}