folk-plugin-process 0.2.0

use std::sync::Arc;

use anyhow::Result;
use async_trait::async_trait;
use bytes::Bytes;
use folk_api::{PluginContext, RpcMethodDef, ServerPlugin};
use tracing::info;

use crate::config::ProcessConfig;
use crate::supervisor::{ProcessStatus, ProcessSupervisor};

pub struct ProcessPlugin {
    config: ProcessConfig,
}

impl ProcessPlugin {
    pub fn new(config: ProcessConfig) -> Self {
        Self { config }
    }
}

#[async_trait]
impl ServerPlugin for ProcessPlugin {
    fn name(&self) -> &'static str {
        "process"
    }

    async fn run(&self, ctx: PluginContext) -> Result<()> {
        let supervisors: Vec<Arc<ProcessSupervisor>> = self
            .config
            .processes
            .iter()
            .map(|def| Arc::new(ProcessSupervisor::new(def.clone())))
            .collect();

        if let Some(reg) = &ctx.rpc_registrar {
            let sups = supervisors.clone();
            reg.register_raw(
                "process.list".into(),
                Arc::new(move |_: Bytes| {
                    let sups = sups.clone();
                    Box::pin(async move {
                        let mut statuses = Vec::new();
                        for s in &sups {
                            let st = s.status().await;
                            statuses.push(format!("{}: {:?}", s.name(), st));
                        }
                        Ok(Bytes::from(rmp_serde::to_vec(&statuses)?))
                    })
                }),
            )
            .await;
        }

        if let Some(health) = &ctx.health_registry {
            let sups = supervisors.clone();
            health
                .register(
                    "process".into(),
                    Arc::new(move || {
                        let sups = sups.clone();
                        Box::pin(async move {
                            let statuses =
                                futures_util::future::join_all(sups.iter().map(|s| s.status()))
                                    .await;
                            let all_ok = statuses
                                .iter()
                                .all(|st| !matches!(st, ProcessStatus::Failed { .. }));
                            if all_ok {
                                folk_api::HealthStatus::ok()
                            } else {
                                folk_api::HealthStatus::degraded("one or more processes failed")
                            }
                        })
                    }),
                )
                .await;
        }

        let handles: Vec<tokio::task::JoinHandle<()>> = supervisors
            .into_iter()
            .map(|s| {
                let sd = ctx.shutdown.clone();
                tokio::spawn(async move {
                    s.run(sd).await;
                })
            })
            .collect();

        info!("process plugin running");

        let mut sd = ctx.shutdown.clone();
        sd.changed().await.ok();
        for h in handles {
            let _ = h.await;
        }
        Ok(())
    }

    fn rpc_methods(&self) -> Vec<RpcMethodDef> {
        vec![
            RpcMethodDef::new("process.list", "list supervised process statuses"),
            RpcMethodDef::new("process.restart", "restart a named process"),
        ]
    }
}