sein 0.1.0

use anyhow::Result;
use async_trait::async_trait;
use petgraph::{graph::NodeIndex, visit::EdgeRef, Graph};
use std::{collections::HashMap, sync::Arc};

// Core types & Trait definitions
pub type NodeId = usize;
pub type Condition<Ctx> = Box<dyn Fn(&Ctx) -> bool + Send + Sync>;
pub type NodeRegistry<Ctx> = HashMap<NodeId, Arc<dyn Node<Ctx>>>;
pub type ToolRegistry<Ctx> = HashMap<String, Arc<dyn Tool<Ctx>>>;

/// Tool: Performs operations or queries on Context
#[async_trait]
pub trait Tool<Ctx>: Send + Sync {
    async fn execute(&self, ctx: &mut Ctx) -> Result<()>;
    fn name(&self) -> &str;
}

/// Node: Executes operations on Context when entered, engine handles subsequent node selection
#[async_trait]
pub trait Node<Ctx>: Send + Sync {
    async fn enter(&self, ctx: &mut Ctx, tools: &ToolRegistry<Ctx>) -> Result<()>;
}

/// AgentEngine: Core state machine implementation
pub struct AgentEngine<Ctx> {
    graph: Graph<NodeId, Condition<Ctx>>,
    nodes: NodeRegistry<Ctx>,
    tools: ToolRegistry<Ctx>,
    current: Option<NodeIndex>,
    context: Ctx,
    next_node_id: NodeId,
}

impl<Ctx> AgentEngine<Ctx> {
    /// Create a new engine with the given context
    pub fn new(context: Ctx) -> Self {
        Self {
            graph: Graph::new(),
            nodes: HashMap::new(),
            tools: HashMap::new(),
            current: None,
            context,
            next_node_id: 0,
        }
    }

    /// Add multiple tools at once
    pub fn add_tools<T>(&mut self, tools: impl IntoIterator<Item = T>)
    where
        T: Tool<Ctx> + 'static,
    {
        for tool in tools {
            let name = tool.name().to_owned();
            self.tools.insert(name, Arc::new(tool));
        }
    }

    /// Add a node and return its NodeId
    pub fn add_node<N: Node<Ctx> + 'static>(&mut self, node: N) -> NodeId {
        let node_id = self.next_node_id;
        self.next_node_id += 1;

        self.graph.add_node(node_id);
        self.nodes.insert(node_id, Arc::new(node));

        node_id
    }

    /// Add a directed edge with transition condition between nodes
    pub fn add_transition<F>(&mut self, from: NodeId, to: NodeId, condition: F) -> Result<()>
    where
        F: Fn(&Ctx) -> bool + Send + Sync + 'static,
    {
        use anyhow::anyhow;

        let from_idx = self
            .graph
            .node_indices()
            .find(|idx| self.graph[*idx] == from)
            .ok_or_else(|| anyhow!("Source node not found: {}", from))?;

        let to_idx = self
            .graph
            .node_indices()
            .find(|idx| self.graph[*idx] == to)
            .ok_or_else(|| anyhow!("Target node not found: {}", to))?;

        self.graph.add_edge(from_idx, to_idx, Box::new(condition));
        Ok(())
    }

    /// Set the starting node
    pub fn set_start_node(&mut self, id: NodeId) -> Result<()> {
        use anyhow::anyhow;

        let start_idx = self
            .graph
            .node_indices()
            .find(|idx| self.graph[*idx] == id)
            .ok_or_else(|| anyhow!("Start node not found: {}", id))?;

        self.current = Some(start_idx);
        Ok(())
    }

    /// Run the state machine until no matching condition is found
    pub async fn run(&mut self) -> Result<()> {
        use anyhow::anyhow;

        let mut current_idx = self
            .current
            .ok_or_else(|| anyhow!("No start node set. Call set_start_node() first"))?;

        loop {
            let node_id = match self.graph.node_weight(current_idx) {
                Some(id) => *id,
                None => break,
            };

            let node = self
                .nodes
                .get(&node_id)
                .ok_or_else(|| anyhow!("Node not found: {}", node_id))?;

            // Enter the node
            node.enter(&mut self.context, &self.tools).await?;

            // Find the next matching edge
            let next_edge = self
                .graph
                .edges(current_idx)
                .find(|edge| (edge.weight())(&self.context));

            match next_edge {
                Some(edge) => {
                    current_idx = edge.target();
                }
                None => {
                    // No matching edge found => exit
                    break;
                }
            }
        }

        // Update current index
        self.current = Some(current_idx);
        Ok(())
    }
}