use std::{future::Future, time::Duration};

use noprocess::{HandleID, Process, ProcessHandler, ProcessManager, ProcessResult};
use tokio_util::sync::CancellationToken;

#[tokio::main]
async fn main() -> Result<(), noprocess::Error> {
    tracing_subscriber::fmt::init();

    let manager = ProcessManager::new();

    // Start with 2 workers
    println!("--- Starting 2 workers ---");
    let mut workers = vec![];
    for i in 0..2 {
        let id = spawn_worker(&manager, i).await;
        workers.push(id);
    }

    tokio::time::sleep(Duration::from_secs(2)).await;

    // Scale up to 4 workers
    println!("\n--- Scaling up to 4 workers ---");
    for i in 2..4 {
        let id = spawn_worker(&manager, i).await;
        workers.push(id);
    }

    tokio::time::sleep(Duration::from_secs(2)).await;

    // Scale down to 2 workers
    println!("\n--- Scaling down to 2 workers ---");
    for id in workers.drain(2..) {
        manager.stop_process(&id).await?;
    }

    tokio::time::sleep(Duration::from_secs(2)).await;

    // Stop all remaining workers
    println!("\n--- Stopping all workers ---");
    for id in workers {
        manager.stop_process(&id).await?;
    }

    println!("All workers stopped");
    Ok(())
}

async fn spawn_worker(manager: &ProcessManager, id: usize) -> HandleID {
    let process = Process::builder(Worker { id })
        .handle_id(format!("worker-{id}"))
        .build();

    let handle_id = manager.add_process(process).await;
    manager.start_process(&handle_id).await.unwrap();
    handle_id
}

struct Worker {
    id: usize,
}

impl ProcessHandler for Worker {
    fn call(&self, cancel: CancellationToken) -> impl Future<Output = ProcessResult> + Send {
        let id = self.id;
        async move {
            println!("[worker-{id}] started");
            let mut interval = tokio::time::interval(Duration::from_millis(500));
            let mut jobs = 0;

            loop {
                tokio::select! {
                    _ = cancel.cancelled() => {
                        println!("[worker-{id}] stopping after {jobs} jobs");
                        break;
                    }
                    _ = interval.tick() => {
                        // Simulate doing work
                        jobs += 1;
                        println!("[worker-{id}] completed job #{jobs}");
                    }
                }
            }
            Ok(())
        }
    }
}
