mixvideo-v2/cargos/uni-comfyui-sdk/examples/advanced_usage.rs

use std::collections::HashMap;
use std::time::Duration;
use tokio::time::sleep;
use uni_comfyui_sdk::{
    UniComfyUIClient, ComfyUIError, Result, ServerHeartbeatRequest, ServerRegistrationRequest,
    ServerUnregisterRequest,
};

#[tokio::main]
async fn main() -> Result<()> {
    let client = UniComfyUIClient::new("http://192.168.0.148:18000")?;

    // Demonstrate comprehensive server management
    println!("=== Server Management Demo ===");
    
    // Register a new server
    let server_name = "demo-server";
    let registration = ServerRegistrationRequest {
        name: server_name.to_string(),
        http_url: "http://192.168.0.148:8188".to_string(),
        ws_url: "ws://192.168.0.148:8188".to_string(),
    };

    match client.register_server(registration).await {
        Ok(response) => println!("✓ Server registered: {:?}", response),
        Err(ComfyUIError::Api { status: 409, .. }) => {
            println!("ℹ Server already exists (409 conflict)")
        }
        Err(e) => println!("✗ Failed to register server: {}", e),
    }

    // Send heartbeat
    let heartbeat = ServerHeartbeatRequest {
        name: server_name.to_string(),
    };

    match client.update_heartbeat(heartbeat).await {
        Ok(response) => println!("✓ Heartbeat sent: {:?}", response),
        Err(e) => println!("✗ Failed to send heartbeat: {}", e),
    }

    // List all servers
    match client.list_all_servers().await {
        Ok(servers) => {
            println!("✓ Found {} servers:", servers.len());
            for server in servers {
                println!("  - {}: {} ({})", server.name, server.http_url, server.status);
            }
        }
        Err(e) => println!("✗ Failed to list servers: {}", e),
    }

    // Demonstrate workflow management
    println!("\n=== Workflow Management Demo ===");

    // Get all workflows
    match client.get_all_workflows().await {
        Ok(workflows) => println!("✓ Found {} workflows", workflows.len()),
        Err(e) => println!("✗ Failed to get workflows: {}", e),
    }

    // Publish a sample workflow
    let mut sample_workflow = HashMap::new();
    sample_workflow.insert("name".to_string(), serde_json::json!("demo_workflow"));
    sample_workflow.insert("version".to_string(), serde_json::json!("1.0.0"));
    sample_workflow.insert(
        "description".to_string(),
        serde_json::json!("A demo workflow for testing"),
    );

    match client.publish_workflow(sample_workflow).await {
        Ok(_) => println!("✓ Workflow published successfully"),
        Err(e) => println!("✗ Failed to publish workflow: {}", e),
    }

    // Demonstrate workflow execution
    println!("\n=== Workflow Execution Demo ===");

    let mut execution_data = HashMap::new();
    execution_data.insert("input_text".to_string(), serde_json::json!("Hello, ComfyUI!"));
    execution_data.insert("steps".to_string(), serde_json::json!(20));
    execution_data.insert("seed".to_string(), serde_json::json!(42));

    match client
        .run_workflow("demo_workflow", Some("1.0.0"), execution_data)
        .await
    {
        Ok(response) => {
            println!("✓ Workflow execution started");
            
            // Try to extract run ID from response for status checking
            if let Some(run_id) = response.get("run_id").and_then(|v| v.as_str()) {
                println!("  Run ID: {}", run_id);
                
                // Poll for status (demo - in real usage you'd want better polling logic)
                for i in 1..=3 {
                    sleep(Duration::from_secs(2)).await;
                    match client.get_run_status(run_id).await {
                        Ok(status) => {
                            println!("  Status check {}: {:?}", i, status);
                        }
                        Err(e) => println!("  Status check {} failed: {}", i, e),
                    }
                }
            }
        }
        Err(e) => println!("✗ Failed to start workflow: {}", e),
    }

    // Demonstrate monitoring
    println!("\n=== Monitoring Demo ===");

    // System health
    match client.get_system_health().await {
        Ok(health) => println!("✓ System health: {:?}", health),
        Err(e) => println!("✗ Failed to get system health: {}", e),
    }

    // System stats
    match client.get_system_stats().await {
        Ok(stats) => println!("✓ System stats retrieved"),
        Err(e) => println!("✗ Failed to get system stats: {}", e),
    }

    // Recent tasks
    match client.get_recent_tasks(Some(5)).await {
        Ok(tasks) => println!("✓ Found {} recent tasks", tasks.len()),
        Err(e) => println!("✗ Failed to get recent tasks: {}", e),
    }

    // Force health check
    match client.force_health_check().await {
        Ok(response) => println!("✓ Health check forced: {:?}", response),
        Err(e) => println!("✗ Failed to force health check: {}", e),
    }

    // Demonstrate RunX endpoints
    println!("\n=== RunX Demo ===");

    // Get schema for model_with_multi_dress
    match client.model_with_multi_dress_json_schema().await {
        Ok(schema) => println!("✓ Got model_with_multi_dress schema"),
        Err(e) => println!("✗ Failed to get schema: {}", e),
    }

    // Try to run model_with_multi_dress workflow
    let mut dress_data = HashMap::new();
    dress_data.insert("model_image".to_string(), serde_json::json!("base64_image_data"));
    dress_data.insert("dress_images".to_string(), serde_json::json!(["dress1.jpg", "dress2.jpg"]));

    match client
        .model_with_multi_dress(None, dress_data)
        .await
    {
        Ok(response) => println!("✓ Model with multi dress executed: {:?}", response),
        Err(e) => println!("✗ Failed to execute model with multi dress: {}", e),
    }

    // Cleanup: Unregister the demo server
    println!("\n=== Cleanup ===");
    let unregister = ServerUnregisterRequest {
        name: server_name.to_string(),
    };

    match client.unregister_server(unregister).await {
        Ok(response) => println!("✓ Server unregistered: {:?}", response),
        Err(e) => println!("✗ Failed to unregister server: {}", e),
    }

    println!("\n=== Demo Complete ===");
    Ok(())
}