use super::super::{HandleConnection, NodeGraphLookups};
use jellyflow_core::core::Graph;
use jellyflow_core::core::{Edge, EdgeId, EdgeKind, EdgeReconnectable};
use jellyflow_core::ops::EdgeEndpoints;

impl NodeGraphLookups {
    pub(super) fn apply_add_edge(&mut self, graph: &Graph, id: EdgeId) -> bool {
        let Some(edge) = graph.edges().get(&id) else {
            return false;
        };
        self.insert_edge_lookup_from_parts(
            graph,
            id,
            edge.kind,
            EdgeEndpoints::from_edge(edge),
            edge.reconnectable,
        )
    }

    pub(super) fn apply_remove_edge(&mut self, id: EdgeId) -> bool {
        self.remove_edge_from_lookups(id);
        true
    }

    pub(super) fn remove_edge_from_lookups(&mut self, edge_id: EdgeId) {
        if let Some(conn) = self.connection_from_edge_lookup(edge_id) {
            self.remove_edge_connection(conn);
        } else {
            self.slow_remove_edge_from_connection_lookup(edge_id);
        }
        self.edge_lookup.remove(&edge_id);
    }

    pub(super) fn remove_edges_from_lookups(&mut self, edges: &[(EdgeId, Edge)]) {
        for (edge_id, _edge) in edges {
            self.remove_edge_from_lookups(*edge_id);
        }
    }

    pub(super) fn apply_set_edge_kind(
        &mut self,
        graph: &Graph,
        id: EdgeId,
        kind: EdgeKind,
    ) -> bool {
        if let Some(e) = self.edge_lookup.get_mut(&id) {
            e.kind = kind;
        } else {
            let Some(edge) = graph.edges().get(&id) else {
                return false;
            };
            return self.recover_edge_lookup_from_graph(graph, id, kind, edge.reconnectable);
        }
        if let Some(conn) = self.connection_from_edge_lookup(id) {
            self.add_edge_connection(conn);
        }
        true
    }

    pub(super) fn apply_set_edge_reconnectable(
        &mut self,
        graph: &Graph,
        id: EdgeId,
        reconnectable: Option<EdgeReconnectable>,
    ) -> bool {
        if let Some(e) = self.edge_lookup.get_mut(&id) {
            e.reconnectable = reconnectable;
            if let Some(conn) = self.connection_from_edge_lookup(id) {
                self.add_edge_connection(conn);
            }
            return true;
        }
        let Some(edge) = graph.edges().get(&id) else {
            return false;
        };
        self.recover_edge_lookup_from_graph(graph, id, edge.kind, reconnectable)
    }

    fn recover_edge_lookup_from_graph(
        &mut self,
        graph: &Graph,
        id: EdgeId,
        kind: EdgeKind,
        reconnectable: Option<EdgeReconnectable>,
    ) -> bool {
        let Some(edge) = graph.edges().get(&id) else {
            return false;
        };
        self.insert_edge_lookup_from_parts(
            graph,
            id,
            kind,
            EdgeEndpoints::from_edge(edge),
            reconnectable,
        )
    }

    pub(super) fn apply_set_edge_endpoints(
        &mut self,
        graph: &Graph,
        id: EdgeId,
        _from: EdgeEndpoints,
        to: EdgeEndpoints,
    ) -> bool {
        if let Some(prev) = self.edge_lookup.get(&id).copied() {
            self.remove_edge_connection(HandleConnection::from_edge_lookup(id, prev));
        } else {
            self.slow_remove_edge_from_connection_lookup(id);
        }

        let Some(kind) = self.edge_kind_for_endpoint_update(graph, id) else {
            return false;
        };

        self.insert_edge_lookup_from_parts(
            graph,
            id,
            kind,
            to,
            self.edge_reconnectable_for_endpoint_update(graph, id),
        )
    }

    fn insert_edge_lookup_from_parts(
        &mut self,
        graph: &Graph,
        id: EdgeId,
        kind: EdgeKind,
        endpoints: EdgeEndpoints,
        reconnectable: Option<EdgeReconnectable>,
    ) -> bool {
        let Some((entry, conn)) =
            Self::edge_lookup_entry_from_graph(graph, id, kind, endpoints, reconnectable)
        else {
            return false;
        };
        self.edge_lookup.insert(id, entry);
        self.add_edge_connection(conn);
        true
    }

    fn edge_kind_for_endpoint_update(&self, graph: &Graph, id: EdgeId) -> Option<EdgeKind> {
        graph
            .edges()
            .get(&id)
            .map(|edge| edge.kind)
            .or_else(|| self.edge_lookup.get(&id).map(|entry| entry.kind))
    }

    fn edge_reconnectable_for_endpoint_update(
        &self,
        graph: &Graph,
        id: EdgeId,
    ) -> Option<EdgeReconnectable> {
        graph
            .edges()
            .get(&id)
            .and_then(|edge| edge.reconnectable)
            .or_else(|| {
                self.edge_lookup
                    .get(&id)
                    .and_then(|entry| entry.reconnectable)
            })
    }
}
