//! Grace periods, service restarts, and retirement of replaced versions (PROTOCOL.md section 3.7).
use crate::Host;
use serde_json::json;
use std::{collections::BTreeMap, sync::atomic::Ordering, time::Duration};
use tokio::sync::watch;

pub(crate) struct OldProcess {
    pid: String,
    version: String,
    ended: watch::Receiver<bool>,
}
pub(crate) struct Upgrade {
    plugin: String,
    previous: String,
    current: String,
    grace: Duration,
    processes: Vec<OldProcess>,
}
impl Host {
    /// Snapshots old admissions while activation holds the admission write guard.
    pub(crate) fn snapshot_upgrade(
        &self,
        plugin: &str,
        previous: String,
        current: &str,
        grace: f64,
    ) -> Upgrade {
        let mut state = self.0.state.lock().unwrap();
        let processes = state
            .processes
            .iter_mut()
            .filter(|(_, process)| {
                process.info.plugin == plugin
                    && process.info.version != current
                    && process.info.method != "rpc.discover"
            })
            .map(|(pid, process)| {
                // The service supervisor owns the restart, including an exit on
                // the superseded notice before the grace timer is installed.
                process.superseded = !process.stopped;
                OldProcess {
                    pid: pid.clone(),
                    version: process.info.version.clone(),
                    ended: process.ended.subscribe(),
                }
            })
            .collect();
        // A cleanly returned service is dormant until an upgrade. A pending
        // admission will choose the current version itself. Disabled services
        // (explicit stop or crash-limit give-up) must never be re-armed here.
        for service in &mut state.services {
            if service.enabled
                && !service.pending
                && service.owner.is_none()
                && service.version.as_deref() != Some(current)
                && service
                    .method
                    .split_once('.')
                    .is_some_and(|(name, _)| name == plugin)
            {
                self.queue_service(service);
            }
        }
        Upgrade {
            plugin: plugin.into(),
            previous,
            current: current.into(),
            grace: Duration::from_secs_f64(grace),
            processes,
        }
    }
    /// Lets old calls finish, stops them after grace, and starts each lifecycle hook once.
    pub(crate) fn begin_upgrade(&self, upgrade: Upgrade) {
        // Queue the notice before returning from activation, without waiting on guest code.
        for process in &upgrade.processes {
            self.send(
                &process.pid,
                json!({"jsonrpc":"2.0","method":"rpc.superseded","params":{"current":upgrade.current}}),
            );
        }
        let host = self.clone();
        tokio::spawn(async move {
            if host.0.shutdown.load(Ordering::SeqCst) {
                return;
            }
            host.0.journal.write(
                "core",
                &upgrade.plugin,
                "core",
                &format!("upgrade grace {} seconds", upgrade.grace.as_secs_f64()),
            );
            let mut versions: BTreeMap<String, Vec<watch::Receiver<bool>>> = BTreeMap::new();
            versions.entry(upgrade.previous).or_default();
            for process in upgrade.processes {
                versions
                    .entry(process.version)
                    .or_default()
                    .push(process.ended.clone());
                let host = host.clone();
                let grace = upgrade.grace;
                tokio::spawn(async move {
                    let mut ended = process.ended;
                    tokio::select! {
                        _ = ended.wait_for(|ended|*ended) => (),
                        _ = tokio::time::sleep(grace) => { host.begin_retire(&process.pid); }
                    }
                });
            }
            for (version, processes) in versions {
                let host = host.clone();
                let plugin = upgrade.plugin.clone();
                let retirement = (plugin.clone(), version.clone());
                if !host
                    .0
                    .retirements
                    .lock()
                    .unwrap()
                    .insert(retirement.clone())
                {
                    continue;
                }
                tokio::spawn(async move {
                    for mut ended in processes {
                        let _ = ended.wait_for(|ended| *ended).await;
                    }
                    loop {
                        if host.0.shutdown.load(Ordering::SeqCst) {
                            host.0.retirements.lock().unwrap().remove(&retirement);
                            return;
                        }
                        let live: Vec<_> = host
                            .0
                            .state
                            .lock()
                            .unwrap()
                            .processes
                            .values()
                            .filter(|process| {
                                process.info.plugin == plugin
                                    && process.info.version == version
                                    && process.info.method != "rpc.discover"
                            })
                            .map(|process| process.ended.subscribe())
                            .collect();
                        if live.is_empty() {
                            break;
                        }
                        for mut ended in live {
                            let _ = ended.wait_for(|ended| *ended).await;
                        }
                    }
                    if host
                        .0
                        .graph
                        .lock()
                        .unwrap()
                        .plugins
                        .get(&plugin)
                        .is_some_and(|current| current.hash == version)
                    {
                        host.0.retirements.lock().unwrap().remove(&retirement);
                        return;
                    }
                    host.0.journal.write(
                        "core",
                        &plugin,
                        "core",
                        &format!("retired start {version}"),
                    );
                    let result = host
                        .route(
                            "core".into(),
                            Some(json!(uuid::Uuid::now_v7().to_string())),
                            format!("{plugin}.rpc.retired"),
                            json!({"previous":version}),
                            true,
                        )
                        .await;
                    host.0.journal.write(
                        "core",
                        &plugin,
                        "core",
                        &format!("retired result {result:?}"),
                    );
                    host.0.retirements.lock().unwrap().remove(&retirement);
                });
            }
        });
    }
}