use crate::ir::dtype::DType;
use crate::ir::node::{Node, NodeId};

#[derive(Debug, Clone, PartialEq)]
pub struct TypedDag {
    pub nodes: Vec<Node>,
    pub dtypes: Vec<DType>,
    pub params: Vec<(String, DType)>,
    pub outputs: Vec<NodeId>,
}

#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum TypedDagError {
    #[error("TypedDag.nodes and dtypes length mismatch")]
    NodesDtypesMismatch,
    #[error("output node id {node_id} out of range (nodes len {len})")]
    InvalidOutput { node_id: usize, len: usize },
}

impl TypedDag {
    pub fn try_new(
        nodes: Vec<Node>,
        dtypes: Vec<DType>,
        params: Vec<(String, DType)>,
        outputs: Vec<NodeId>,
    ) -> Result<Self, TypedDagError> {
        if nodes.len() != dtypes.len() {
            return Err(TypedDagError::NodesDtypesMismatch);
        }
        let len = nodes.len();
        for out in &outputs {
            if out.0 >= len {
                return Err(TypedDagError::InvalidOutput {
                    node_id: out.0,
                    len,
                });
            }
        }
        Ok(Self {
            nodes,
            dtypes,
            params,
            outputs,
        })
    }

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

    pub fn is_well_formed(&self) -> bool {
        self.nodes.len() == self.dtypes.len()
            && self.outputs.iter().all(|&id| id.0 < self.nodes.len())
    }
}
