use mofa_kernel::workflow::{StateGraph, JsonState, START, END};
use mofa_foundation::workflow::{StateGraphImpl, AppendReducer, OverwriteReducer};
use mofa_kernel::workflow::{NodeFunc, Command, RuntimeContext};
use mofa_kernel::agent::error::AgentResult;
use async_trait::async_trait;
use serde_json::json;

// Define a custom node
struct ProcessNode;

#[async_trait]
impl NodeFunc<JsonState> for ProcessNode {
    async fn call(&self, state: &mut JsonState, _ctx: &RuntimeContext) -> AgentResult<Command> {
        // Process state
        let input = state.get_value("input").unwrap_or(json!(null));
        
        // Return command with updates
        Ok(Command::new()
            .update("messages", json!("Processed input"))
            .update("result", json!(format!("Done: {:?}", input)))
            .continue_())
    }
    
    fn name(&self) -> &str {
        "process"
    }
}

#[tokio::main]
async fn main() -> AgentResult<()> {
    // Build graph
    let mut graph = StateGraphImpl::<JsonState>::new("example")
        // Configure reducers
        .add_reducer("messages", Box::new(AppendReducer))
        .add_reducer("result", Box::new(OverwriteReducer))
        // Add nodes
        .add_node("process", Box::new(ProcessNode))
        // Add edges
        .add_edge(START, "process")
        .add_edge("process", END)
        // Compile
        .compile()?;
    
    // Execute
    let mut initial_state = JsonState::new();
    initial_state.apply_update("input", json!("Hello")).await?;
    
    let final_state = graph.invoke(initial_state, None).await?;
    
    println!("Result: {:?}", final_state.get_value("result"));
    
    Ok(())
}

use std::collections::HashMap;

// Node that returns routing decision
struct ClassifierNode;

#[async_trait]
impl NodeFunc<JsonState> for ClassifierNode {
    async fn call(&self, state: &mut JsonState, _ctx: &RuntimeContext) -> AgentResult<Command> {
        let input = state.get_value("query").unwrap_or(json!(null));
        let category = classify(input); // Your classification logic
        
        Ok(Command::new()
            .route(&category)  // Explicit routing
            .update("category", json!(category))
            .continue_())
    }
    
    fn name(&self) -> &str { "classifier" }
}

// Build graph with conditional edges
let mut graph = StateGraphImpl::<JsonState>::new("router")
    .add_node("classify", Box::new(ClassifierNode))
    .add_node("handle_billing", Box::new(BillingNode))
    .add_node("handle_technical", Box::new(TechnicalNode))
    .add_edge(START, "classify")
    .add_conditional_edges(
        "classify",
        HashMap::from([
            ("billing".to_string(), "handle_billing".to_string()),
            ("technical".to_string(), "handle_technical".to_string()),
        ])
    )
    .add_edge("handle_billing", END)
    .add_edge("handle_technical", END)
    .compile()?;

// Build graph with parallel execution
let mut graph = StateGraphImpl::<JsonState>::new("parallel")
    .add_node("fan_out", Box::new(FanOutNode))
    .add_node("task_a", Box::new(TaskA))
    .add_node("task_b", Box::new(TaskB))
    .add_node("task_c", Box::new(TaskC))
    .add_node("merge", Box::new(MergeNode))
    .add_edge(START, "fan_out")
    .add_parallel_edges(
        "fan_out",
        vec!["task_a".to_string(), "task_b".to_string(), "task_c".to_string()]
    )
    .add_edge("task_a", "merge")
    .add_edge("task_b", "merge")
    .add_edge("task_c", "merge")
    .add_edge("merge", END)
    .compile()?;