//! The supervisor of one WASM actor: it runs the instance, feeds it the
//! mailbox, routes its output, restarts it after a crash, and swaps versions.
//!
//! The broker talks to a supervisor only through `Control` requests.

use std::sync::Arc;
use std::time::{Duration, Instant};

use anyhow::{Result, bail};
use serde_json::{Value, json};
use tokio::sync::{mpsc, oneshot};
use wasmtime::component::Component;

use crate::broker::{Broker, Envelope, SHUTDOWN_DEADLINE, shutdown_line};
use crate::runtime::{Exit, Instance};

/// A new version that crashes this often within `CRASH_WINDOW` after an
/// upgrade is rolled back automatically.
const CRASH_LIMIT: usize = 3;
const CRASH_WINDOW: Duration = Duration::from_secs(30);

pub type Reply = oneshot::Sender<Result<Value>>;

pub enum Control {
    /// Swap to `version` (already in `versions/`, compiled as `component`).
    Upgrade { version: u64, component: Component, requester: String, reply: Reply },
    Rollback { version: u64, reply: Reply },
    Stop { reply: Reply },
    Start { reply: Reply },
    Shutdown { reply: Reply },
}

/// The time after an upgrade in which crashes count toward a rollback.
struct Probation {
    previous: u64,
    since: Instant,
    crashes: usize,
    requester: String,
}

pub struct Supervisor {
    broker: Arc<Broker>,
    id: String,
    version: u64,
    running: bool,
    mailbox: mpsc::Receiver<Envelope>,
    control: mpsc::Receiver<Control>,
    component: Option<Component>,
    instance: Option<Instance>,
    started_at: Instant,
    restarts: u32,
    restart_at: Option<Instant>,
    probation: Option<Probation>,
    /// A control request that arrived during a blocked delivery.
    next_control: Option<Control>,
}

async fn next_line(instance: &mut Option<Instance>) -> Option<String> {
    match instance {
        Some(i) => i.stdout.recv().await,
        None => std::future::pending().await,
    }
}

async fn sleep_until(t: Option<Instant>) {
    match t {
        Some(t) => tokio::time::sleep_until(t.into()).await,
        None => std::future::pending().await,
    }
}

impl Supervisor {
    pub fn new(
        broker: Arc<Broker>,
        id: String,
        version: u64,
        running: bool,
        mailbox: mpsc::Receiver<Envelope>,
        control: mpsc::Receiver<Control>,
    ) -> Self {
        Self {
            broker,
            id,
            version,
            running,
            mailbox,
            control,
            component: None,
            instance: None,
            started_at: Instant::now(),
            restarts: 0,
            restart_at: None,
            probation: None,
            next_control: None,
        }
    }

    pub async fn run(mut self) {
        loop {
            if let Some(ctl) = self.next_control.take() {
                if self.handle(ctl).await {
                    break;
                }
                continue;
            }
            if self.running && self.instance.is_none() && self.restart_at.is_none_or(|t| Instant::now() >= t) {
                self.restart_at = None;
                let _ = self.start().await;
            }
            let has_instance = self.instance.is_some();
            tokio::select! {
                ctl = self.control.recv() => {
                    let Some(ctl) = ctl else { break };
                    if self.handle(ctl).await {
                        break;
                    }
                }
                msg = self.mailbox.recv(), if has_instance => {
                    let Some(msg) = msg else { break };
                    self.deliver(msg).await;
                }
                line = next_line(&mut self.instance) => match line {
                    Some(line) => self.broker.route_line(&self.id, &line),
                    None => {
                        let exit = self.instance.as_mut().unwrap().wait().await;
                        self.instance = None;
                        self.on_exit(exit).await;
                    }
                },
                _ = sleep_until(self.restart_at), if !has_instance && self.running => {}
            }
        }
    }

    /// Write a message to the instance's stdin, and route its output
    /// meanwhile. If a control request arrives before the actor reads the
    /// message, the instance is killed (a partly written line would corrupt
    /// the next message) and the request runs.
    async fn deliver(&mut self, msg: Envelope) {
        let Some(inst) = self.instance.as_mut() else { return };
        let (broker, id, line) = (self.broker.clone(), self.id.clone(), msg.to_line());
        let stuck = tokio::select! {
            written = inst.send_line(&line, |l| broker.route_line(&id, &l)) => match written {
                Ok(()) => return,
                // The instance is ending. The main loop sees its closed stdout.
                Err(_) => false,
            },
            Some(ctl) = self.control.recv() => {
                self.next_control = Some(ctl);
                true
            }
        };
        if msg.from != "broker" {
            let body = json!({"type": "undeliverable", "reason": "the actor stopped", "to": self.id});
            self.broker.send("broker", Some(&msg.from), None, Some(&msg.id), body);
        }
        if stuck && let Some(inst) = &self.instance {
            self.broker.log(&self.id, "killed: it did not read stdin when the broker needed it");
            inst.kill();
        }
    }

    /// Start an instance of the current version. If that fails, the actor
    /// stays stopped until the next start, upgrade, or rollback.
    async fn start(&mut self) -> Result<()> {
        let started = self.try_start();
        match &started {
            Ok(()) => self.started_at = Instant::now(),
            Err(e) => {
                self.broker.log(&self.id, &format!("cannot start version {}: {e:#}", self.version));
                self.running = false;
            }
        }
        started
    }

    fn try_start(&mut self) -> Result<()> {
        if self.component.is_none() {
            self.component = Some(self.broker.load_version(&self.id, self.version)?);
        }
        let cfg = self.broker.instance_config(&self.id, self.version, None);
        self.instance = Some(self.broker.runtime.start(self.component.as_ref().unwrap(), cfg)?);
        Ok(())
    }

    /// Stop the instance gracefully. Its last messages still get routed.
    async fn stop(&mut self) {
        let Some(inst) = self.instance.take() else { return };
        let (broker, id) = (self.broker.clone(), self.id.clone());
        let exit = inst.stop(&shutdown_line(), SHUTDOWN_DEADLINE, |line| broker.route_line(&id, &line)).await;
        self.broker.log(&self.id, &format!("stopped: {exit}"));
    }

    async fn on_exit(&mut self, exit: Exit) {
        self.broker.log(&self.id, &format!("instance ended unexpectedly: {exit}"));
        if let Some(p) = &mut self.probation
            && p.since.elapsed() < CRASH_WINDOW
        {
            p.crashes += 1;
            if p.crashes >= CRASH_LIMIT {
                let (from, to, requester) = (self.version, p.previous, p.requester.clone());
                self.broker.log(&self.id, &format!("version {from} crashed {CRASH_LIMIT} times after the upgrade; rolling back to {to}"));
                let body = match self.rollback(to).await {
                    Ok(_) => json!({"type": "rolled_back", "actor": self.id, "from": from, "to": to,
                                    "reason": "crash loop after upgrade", "last_exit": exit.to_string()}),
                    Err(e) => json!({"type": "error", "actor": self.id, "reason": format!("automatic rollback failed: {e:#}")}),
                };
                self.broker.send("broker", Some(&requester), None, None, body);
                return;
            }
        }
        // Restart after 0.5 s, doubling up to 30 s; a minute of uptime resets it.
        if self.started_at.elapsed() > Duration::from_secs(60) {
            self.restarts = 0;
        }
        let delay = Duration::from_millis(500 << self.restarts.min(6)).min(Duration::from_secs(30));
        self.restarts += 1;
        self.restart_at = Some(Instant::now() + delay);
    }

    /// Make `version` current. It is compiled first (unless `component` is
    /// given): if that fails, nothing changes. `data/` and `public/` stay as
    /// they are. A stopped actor stays stopped.
    async fn install(&mut self, version: u64, component: Option<Component>) -> Result<()> {
        let component = match component {
            Some(c) => c,
            None => self.broker.load_version(&self.id, version)?,
        };
        self.stop().await;
        self.broker.install(&self.id, version, self.running)?;
        self.version = version;
        self.component = Some(component);
        self.restarts = 0;
        self.restart_at = None;
        Ok(())
    }

    async fn upgrade(&mut self, version: u64, component: Component, requester: String) -> Result<Value> {
        let previous = self.version;
        self.install(version, Some(component)).await?;
        if !self.running {
            // Installed, not started: main runs and mail is read at the next `start`.
            return Ok(json!({"type": "upgraded", "actor": self.id, "version": version, "previous": previous, "running": false}));
        }
        self.probation = Some(Probation { previous, since: Instant::now(), crashes: 0, requester });
        if let Err(e) = self.start().await {
            self.rollback(previous).await?;
            bail!("version {version} did not start ({e:#}); rolled back to {previous}");
        }
        Ok(json!({"type": "upgraded", "actor": self.id, "version": version, "previous": previous, "running": true}))
    }

    async fn rollback(&mut self, version: u64) -> Result<Value> {
        let from = self.version;
        self.install(version, None).await?;
        self.probation = None;
        if self.running {
            self.start().await?;
        }
        Ok(json!({"type": "rolled_back", "actor": self.id, "from": from, "to": version, "running": self.running}))
    }

    /// Handle a control request. Returns `true` when the supervisor ends.
    async fn handle(&mut self, ctl: Control) -> bool {
        match ctl {
            Control::Upgrade { version, component, requester, reply } => {
                let _ = reply.send(self.upgrade(version, component, requester).await);
            }
            Control::Rollback { version, reply } => {
                let _ = reply.send(self.rollback(version).await);
            }
            Control::Stop { reply } => {
                self.running = false;
                self.stop().await;
                let r = self.broker.set_running(&self.id, false);
                let _ = reply.send(r.map(|_| json!({"type": "stopped", "actor": self.id})));
            }
            Control::Start { reply } => {
                self.running = true;
                self.restart_at = None;
                let r = self.broker.set_running(&self.id, true);
                let _ = reply.send(r.map(|_| json!({"type": "started", "actor": self.id})));
            }
            Control::Shutdown { reply } => {
                self.stop().await;
                let _ = reply.send(Ok(json!({"type": "shut_down"})));
                return true;
            }
        }
        false
    }
}