Documentation
use anyhow::Result;
use async_trait::async_trait;
use tokio::time::{sleep, Duration};

// Import the library
use sein::{AgentEngine, Node, Tool};

// Counter example implementation
#[derive(Debug)]
pub struct CounterContext {
    current: u32,
    max: u32,
}

// Node implementations
pub struct ThinkingNode;

#[async_trait]
impl Node<CounterContext> for ThinkingNode {
    async fn enter(&self, ctx: &mut CounterContext, _tools: &sein::ToolRegistry<CounterContext>) -> Result<()> {
        sleep(Duration::from_millis(500)).await;
        println!("Thinking... current count: {}/{}", ctx.current, ctx.max);
        Ok(())
    }
}

pub struct GreetNode;

#[async_trait]
impl Node<CounterContext> for GreetNode {
    async fn enter(&self, ctx: &mut CounterContext, tools: &sein::ToolRegistry<CounterContext>) -> Result<()> {
        let greet_tool = tools
            .get("greet")
            .ok_or_else(|| anyhow::anyhow!("Tool 'greet' not found in registry"))?;
        greet_tool.execute(ctx).await?;
        Ok(())
    }
}

pub struct ExitNode;

#[async_trait]
impl Node<CounterContext> for ExitNode {
    async fn enter(&self, _ctx: &mut CounterContext, _tools: &sein::ToolRegistry<CounterContext>) -> Result<()> {
        println!("Done! Exiting...");
        Ok(())
    }
}

// Tool implementation
pub struct GreetTool;

#[async_trait]
impl Tool<CounterContext> for GreetTool {
    async fn execute(&self, ctx: &mut CounterContext) -> Result<()> {
        sleep(Duration::from_secs(1)).await;
        ctx.current += 1;
        println!("Hello! This is greeting number {}", ctx.current);
        Ok(())
    }

    fn name(&self) -> &str {
        "greet"
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    // 1) Create engine with context
    let mut engine = AgentEngine::new(CounterContext { current: 0, max: 3 });

    // 2) Add tools
    engine.add_tools(vec![GreetTool]);

    // 3) Add nodes
    let thinking_node = engine.add_node(ThinkingNode);
    let greet_node = engine.add_node(GreetNode);
    let exit_node = engine.add_node(ExitNode);

    // 4) Add transitions
    engine.add_transition(thinking_node, greet_node, |ctx| ctx.current < ctx.max)?;
    engine.add_transition(thinking_node, exit_node, |ctx| ctx.current >= ctx.max)?;
    engine.add_transition(greet_node, thinking_node, |_| true)?;

    // 5) Set start node
    engine.set_start_node(thinking_node)?;

    // 6) Run
    engine.run().await?;

    Ok(())
}