honey-biscuit-workshop/workshop-baker/src/notify/handler.rs

use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use chrono::Utc;
use tokio::sync::Notify;

use crate::finalize::NotificationTemplateDef;
use crate::notify::template::{
    build_batch_template_context, render_template, resolve_template,
    select_template_str,
};
use crate::notify::types::{
    MethodStatus, NotificationEvent,
    NotificationEventSender, NotificationPluginMap, NotificationRenderer,
    NotificationRulesMap,
};
use crate::notify::util::{
    find_rule, get_fallible, is_temporary_error, notify_delay, wait_time,
};

/// Shared state for a single dispatch bucket.
struct BucketEntry {
    events: Arc<Mutex<Vec<NotificationEvent>>>,
    result: Arc<Mutex<Option<(MethodStatus, Vec<NotificationEvent>)>>>,
    done: Arc<AtomicBool>,
    notify: Arc<Notify>,
    method: String,
    rule_id: usize,
    _handle: tokio::task::JoinHandle<()>,
}

/// Concurrent notification handler with bucket-based dispatch.
///
/// Events are routed to matching buckets. Buckets cool down and dispatch
/// plugins concurrently. The caller awaits bucket completion, then handles
/// retry/escalation decisions.
pub struct NotificationHandler {
    buckets: Mutex<HashMap<(String, usize), BucketEntry>>,
    pub config: Arc<NotificationRulesMap>,
    plugins: Arc<NotificationPluginMap>,
    templates: Arc<HashMap<String, NotificationTemplateDef>>,
    ctx: Arc<workshop_engine::ExecutionContext>,
    retry_tx: NotificationEventSender,
    period: u32,
    watermark: u32,
}

impl NotificationHandler {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        config: Arc<NotificationRulesMap>,
        plugins: Arc<NotificationPluginMap>,
        templates: Arc<HashMap<String, NotificationTemplateDef>>,
        ctx: Arc<workshop_engine::ExecutionContext>,
        retry_tx: NotificationEventSender,
        period: u32,
        watermark: u32,
    ) -> Self {
        NotificationHandler {
            buckets: Mutex::new(HashMap::new()),
            config,
            plugins,
            templates,
            ctx,
            retry_tx,
            period,
            watermark,
        }
    }

    /// Process a single event: route to buckets, await completion,
    /// then handle retry/escalation for each bucket result.
    pub async fn process_event(self: Arc<Self>, event: NotificationEvent) {
        let entries = self.route_event(&event);
        if entries.is_empty() {
            return;
        }

        // Await all buckets — yields CPU
        for (_, notify) in &entries {
            notify.notified().await;
        }

        // Collect results and remove entries
        let mut results: Vec<(String, usize, Option<(MethodStatus, Vec<NotificationEvent>)>)> =
            Vec::new();
        {
            let mut buckets = self.buckets.lock().unwrap();
            for (key, _) in &entries {
                if let Some(entry) = buckets.remove(key) {
                    let r = entry.result.lock().unwrap().take();
                    results.push((entry.method, entry.rule_id, r));
                }
            }
        }

        // Handle retry/escalation per method
        for (method, _rule_id, result) in results {
            let Some((status, mut events)) = result else { continue };

            match &status {
                MethodStatus::Success => {
                    log::debug!("[NOTIFY] {}: success ({} events)", method, events.len());
                }
                MethodStatus::TemporaryError(_) => {
                    log::warn!("[NOTIFY] {}: temporary error, retrying", method);
                    for event in events.iter_mut() {
                        event.target.rules = Some(vec![method.clone()]);
                        event.target.count += 1;
                        event.not_before = Utc::now() + notify_delay(event);
                        let _ = self.retry_tx.send(event.clone()).await;
                    }
                }
                MethodStatus::PermanentError(_) | MethodStatus::Exhausted(_) => {
                    let fallible = get_fallible(&self.config, &method);
                    if fallible {
                        log::error!("[NOTIFY] {}: discard (fallible)", method);
                    } else {
                        log::warn!("[NOTIFY] {}: escalating priority", method);
                        for event in events.iter_mut() {
                            event.target.priority += 1;
                            event.target.rules = None;
                            event.attempted = event.target.count;
                            event.target.count = 0;
                            event.effective_priority = event.priority;
                            let _ = self.retry_tx.send(event.clone()).await;
                        }
                    }
                }
            }
        }
    }

    /// Route a single event to all matching buckets.
    ///
    /// Returns Notify handles for awaiting bucket completion.
    fn route_event(
        &self,
        event: &NotificationEvent,
    ) -> Vec<((String, usize), Arc<Notify>)> {
        let mut matches: Vec<(String, usize)> = Vec::new();

        for ruleset in self.config.values() {
            for (method, rules) in ruleset.iter() {
                if let Some(ref method_rules) = event.target.rules {
                    if !method_rules.contains(method) {
                        continue;
                    }
                }
                for (rule_id, rule) in rules.iter().enumerate() {
                    if !rule.enabled {
                        continue;
                    }
                    if rule.accept(&event.outcome, &event.stage, &event.substage) {
                        matches.push((method.clone(), rule_id));
                    }
                }
            }
        }

        let mut buckets = self.buckets.lock().unwrap();
        let mut result: Vec<((String, usize), Arc<Notify>)> = Vec::new();

        for key in matches {
            if let Some(entry) = buckets.get(&key) {
                if entry.done.load(Ordering::Acquire) {
                    buckets.remove(&key);
                }
            }

            if let Some(entry) = buckets.get(&key) {
                entry.events.lock().unwrap().push(event.clone());
                result.push((key, entry.notify.clone()));
            } else {
                let (entry, notify) = Self::spawn_bucket(
                    &key.0,
                    key.1,
                    &self.config,
                    &self.plugins,
                    &self.templates,
                    &self.ctx,
                    self.period,
                    self.watermark,
                );
                entry.events.lock().unwrap().push(event.clone());
                result.push((key.clone(), notify));
                buckets.insert(key, entry);
            }
        }

        result
    }

    /// Spawn a bucket task: cooldown → take → dispatch → store → notify.
    fn spawn_bucket(
        method: &str,
        rule_id: usize,
        config: &Arc<NotificationRulesMap>,
        plugins: &Arc<NotificationPluginMap>,
        templates: &Arc<HashMap<String, NotificationTemplateDef>>,
        ctx: &Arc<workshop_engine::ExecutionContext>,
        period: u32,
        watermark: u32,
    ) -> (BucketEntry, Arc<Notify>) {
        let events: Arc<Mutex<Vec<NotificationEvent>>> = Arc::new(Mutex::new(Vec::new()));
        let result: Arc<Mutex<Option<(MethodStatus, Vec<NotificationEvent>)>>> =
            Arc::new(Mutex::new(None));
        let notify = Arc::new(Notify::new());
        let done = Arc::new(AtomicBool::new(false));

        let ec = events.clone();
        let rc = result.clone();
        let nc = notify.clone();
        let dc = done.clone();
        let m = method.to_string();
        let c = config.clone();
        let p = plugins.clone();
        let t = templates.clone();
        let x = ctx.clone();

        let handle = tokio::spawn(async move {
            // Phase 1: Cooldown
            let count = ec.lock().unwrap().len() as u32;
            let wait_secs = wait_time(period, watermark, count);
            tokio::time::sleep(Duration::from_secs_f32(wait_secs)).await;

            // Phase 2: Take events
            let batch = std::mem::take(&mut *ec.lock().unwrap());
            if batch.is_empty() {
                dc.store(true, Ordering::Release);
                nc.notify_waiters();
                return;
            }

            // Phase 3: Dispatch
            let rule = match find_rule(&c, &m, rule_id) {
                Ok(r) => r,
                Err(_) => {
                    *rc.lock().unwrap() = Some((MethodStatus::PermanentError("rule not found".into()), batch));
                    dc.store(true, Ordering::Release);
                    nc.notify_waiters();
                    return;
                }
            };

            let plugin = match p.get(&m) {
                Some(p) => p,
                None => {
                    *rc.lock().unwrap() = Some((MethodStatus::PermanentError("plugin not found".into()), batch));
                    dc.store(true, Ordering::Release);
                    nc.notify_waiters();
                    return;
                }
            };

            let renderer = plugin.metadata.renderer.clone().unwrap_or_default();
            let argument = build_plugin_argument(&renderer, rule, &m, &t, &x, &batch);
            let status = match argument {
                Some(arg) => {
                    match tokio::time::timeout(rule.timeout, plugin.call(arg, &x, &None)).await {
                        Ok(Ok(())) => MethodStatus::Success,
                        Ok(Err(e)) if is_temporary_error(&e) => {
                            MethodStatus::TemporaryError(format!("{}", e))
                        }
                        Ok(Err(e)) => MethodStatus::PermanentError(format!("{}", e)),
                        Err(_) => MethodStatus::TemporaryError("timeout".into()),
                    }
                }
                None => MethodStatus::PermanentError("argument build failed".into()),
            };

            *rc.lock().unwrap() = Some((status, batch));
            dc.store(true, Ordering::Release);
            nc.notify_waiters();
        });

        let entry = BucketEntry {
            events,
            result,
            done,
            notify: notify.clone(),
            method: method.to_string(),
            rule_id,
            _handle: handle,
        };

        (entry, notify)
    }
}

/// Build the YAML argument for a plugin call based on renderer mode.
fn build_plugin_argument(
    renderer: &NotificationRenderer,
    rule: &crate::notify::types::NotificationRule,
    method: &str,
    templates: &HashMap<String, NotificationTemplateDef>,
    ctx: &workshop_engine::ExecutionContext,
    events: &[NotificationEvent],
) -> Option<serde_yaml::Value> {
    let first = events.first()?;

    match renderer {
        NotificationRenderer::Server => {
            let template_def = resolve_template(
                &rule.template.schema,
                method,
                templates,
                ctx.standalone,
                None,
            )
            .ok()?;

            let template_str = select_template_str(&template_def, &first.outcome)
                .or_else(|| rule.template.on_success.as_deref())
                .or_else(|| rule.template.on_failure.as_deref())
                .or_else(|| {
                    rule.template
                        .on_finish_of
                        .as_ref()
                        .map(|o| o.content.as_str())
                })?;

            let data = build_batch_template_context(ctx, events);
            let rendered = render_template(template_str, &data).ok()?;

            let mut argument = rule.config.clone();
            if let serde_yaml::Value::Mapping(ref mut map) = argument {
                map.insert(
                    serde_yaml::Value::String("_rendered_content".into()),
                    serde_yaml::Value::String(rendered),
                );
            }
            log::info!("SSR: rendered content for method '{}'", method);
            Some(argument)
        }
        NotificationRenderer::Plugin => {
            unimplemented!("Plugin renderer is not yet implemented")
        }
        NotificationRenderer::Interactive => {
            unimplemented!("Interactive renderer is not yet implemented")
        }
    }
}