sein 0.1.2

use std::time::Duration;
use anyhow::Result;
use async_trait::async_trait;
use sein::{Engine, EngineControl, NodeTrait, NodeWrapper};
// use std::marker::PhantomData;

// 1. Define your custom context struct - no NodeKey needed!
#[derive(Debug)]
struct MyContext {
    count: u32,
    // Add any other state your nodes need
}

struct ThinkingNode;

#[async_trait]
impl NodeTrait<MyContext> for ThinkingNode {
    async fn enter(&self, ctx: &mut MyContext) -> Result<EngineControl> {
        println!("Entering ThinkingNode (count: {})", ctx.count);
        tokio::time::sleep(Duration::from_secs(1)).await; // Simulate some async work
        
        ctx.count += 1;

        match ctx.count {
            1 => {
                println!("ThinkingNode: Count is 1, continuing.");
                Ok(EngineControl::Continue)
            }
            2 => {
                println!("ThinkingNode: Count is 2, adding GreetToolNode.");
                // Properly wrap the GreetToolNode to implement AnyNode
                let wrapper = NodeWrapper::new(GreetToolNode);
                Ok(EngineControl::AddNode("greet".to_string(), Box::new(wrapper)))
            }
            3 => {
                println!("ThinkingNode: Count is 3, jumping to greet node.");
                // Jump to node by name
                Ok(EngineControl::JumpToNode("greet".to_string()))
            }
            4 => {
                println!("ThinkingNode: Count is 4, ending.");
                Ok(EngineControl::End)
            }
            _ => {
                println!("ThinkingNode: Count is {}, ending.", ctx.count);
                Ok(EngineControl::End)
            }
        }
    }
}

struct GreetToolNode;

#[async_trait]
impl NodeTrait<MyContext> for GreetToolNode {
    async fn enter(&self, ctx: &mut MyContext) -> Result<EngineControl> {
        println!("Entering GreetToolNode (count: {})", ctx.count);
        println!("GreetToolNode: Incremented count to {}, continuing.", ctx.count);        
        Ok(EngineControl::JumpToNode("thinking".to_string()))
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    // Create your custom context instance
    let ctx = MyContext { count: 0 };

    // Create the engine
    let mut engine = Engine::new(ctx);

    // Add nodes with names
    engine.add_node("thinking", ThinkingNode)?;
    println!("Added 'thinking' node");
    
    // Set the start node by name
    engine.set_start_node("thinking")?;
    println!("Set start node to 'thinking'");

    println!("Starting engine run...");
    engine.run().await?;

    println!("Engine run finished.");
    println!("Final context state: {:?}", engine.ctx);

    // Verify the count reached 4 as expected
    assert_eq!(engine.ctx.count, 4, "Count should be 4 after execution");
    println!("Assertion passed: Final count is 4.");

    Ok(())
}