use crate::rules::ConnectPlan;
use jellyflow_core::core::{EdgeKind, Graph, Port, PortDirection, PortId};
use jellyflow_core::interaction::NodeGraphConnectionMode;
use jellyflow_core::ops::EdgeEndpoints;

use super::rejections::{
    reject_incompatible_port_kinds, reject_missing_port, reject_opposite_directions_required,
    reject_self_connection,
};

pub(in crate::rules::connection) struct ConnectionEndpoints<'a> {
    pub from_id: PortId,
    pub to_id: PortId,
    pub from: &'a Port,
    pub to: &'a Port,
    pub edge_kind: EdgeKind,
}

impl ConnectionEndpoints<'_> {
    pub(in crate::rules::connection) fn edge_endpoints(&self) -> EdgeEndpoints {
        EdgeEndpoints::new(self.from_id, self.to_id)
    }

    pub(in crate::rules::connection) fn is_out_to_in(&self) -> bool {
        self.from.dir == PortDirection::Out && self.to.dir == PortDirection::In
    }
}

pub(in crate::rules::connection) fn resolve_connection_endpoints(
    graph: &Graph,
    a: PortId,
    b: PortId,
    mode: NodeGraphConnectionMode,
) -> Result<ConnectionEndpoints<'_>, ConnectPlan> {
    if a == b {
        return Err(reject_self_connection());
    }

    let Some(port_a) = graph.ports().get(&a) else {
        return Err(reject_missing_port(a));
    };
    let Some(port_b) = graph.ports().get(&b) else {
        return Err(reject_missing_port(b));
    };

    let (from_id, to_id) = match mode {
        NodeGraphConnectionMode::Strict => match (port_a.dir, port_b.dir) {
            (PortDirection::Out, PortDirection::In) => (a, b),
            (PortDirection::In, PortDirection::Out) => (b, a),
            _ => {
                return Err(reject_opposite_directions_required());
            }
        },
        NodeGraphConnectionMode::Loose => match port_a.dir {
            PortDirection::Out => (a, b),
            PortDirection::In => (b, a),
        },
    };

    resolve_ordered_connection_endpoints(graph, from_id, to_id)
}

pub(in crate::rules::connection) fn resolve_ordered_connection_endpoints(
    graph: &Graph,
    from_id: PortId,
    to_id: PortId,
) -> Result<ConnectionEndpoints<'_>, ConnectPlan> {
    if from_id == to_id {
        return Err(reject_self_connection());
    }

    let (from, to) = connection_ports(graph, from_id, to_id)?;

    if from.kind != to.kind {
        return Err(reject_incompatible_port_kinds(from.kind, to.kind));
    }

    let edge_kind = from.kind.edge_kind();

    Ok(ConnectionEndpoints {
        from_id,
        to_id,
        from,
        to,
        edge_kind,
    })
}

pub(in crate::rules::connection) fn connection_ports(
    graph: &Graph,
    from_id: PortId,
    to_id: PortId,
) -> Result<(&Port, &Port), ConnectPlan> {
    let Some(from) = graph.ports().get(&from_id) else {
        return Err(reject_missing_port(from_id));
    };
    let Some(to) = graph.ports().get(&to_id) else {
        return Err(reject_missing_port(to_id));
    };
    Ok((from, to))
}
