//! The broker: message routing, mailboxes, subscriptions, the version store,
//! and the requests to `broker` (spawn, upgrade, rollback, stop, start,
//! subscribe, unsubscribe, scan, exit). `supervisor.rs` runs each WASM actor.
//!
//! The broker accepts only `.wasm` files. It does not know any source language.
//! The world folder is the source of truth: `scan` boots every
//! `actor/<id>/actor.wasm` and treats a changed file as an upgrade.

use std::collections::{BTreeMap, HashMap};
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

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

use crate::candidate;
use crate::runtime::{InstanceConfig, Mount, Runtime};
use crate::supervisor::{Control, Reply, Supervisor};
use crate::world::{self, ActorMeta, NATIVE, VersionMeta, World};

pub const MAILBOX_SIZE: usize = 1000;
pub const SHUTDOWN_DEADLINE: Duration = Duration::from_secs(5);

#[derive(Clone, Debug)]
pub struct Envelope {
    pub id: String,
    pub from: String,
    pub to: String,
    pub reply_to: Option<String>,
    pub body: Value,
}

impl Envelope {
    /// The NDJSON line that the recipient reads on stdin.
    pub fn to_line(&self) -> String {
        let mut v = json!({"id": self.id, "from": self.from, "body": self.body});
        if let Some(r) = &self.reply_to {
            v["reply_to"] = json!(r);
        }
        v.to_string()
    }
}

pub fn shutdown_line() -> String {
    json!({"id": "broker/shutdown", "from": "broker", "body": {"type": "shutdown"}}).to_string()
}

pub struct Broker {
    pub world: World,
    pub runtime: Arc<Runtime>,
    mailboxes: Mutex<HashMap<String, mpsc::Sender<Envelope>>>,
    controls: Mutex<HashMap<String, mpsc::Sender<Control>>>,
    /// `actor.json` of every actor, native and WASM.
    metas: Mutex<BTreeMap<String, ActorMeta>>,
    op_locks: Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>,
    seq: AtomicU64,
    /// Set by the `exit` request: the host process should stop.
    exit: tokio::sync::Notify,
}

impl Broker {
    pub fn new(world: World, runtime: Arc<Runtime>) -> Result<Arc<Self>> {
        let metas = world.actor_ids()?.into_iter().map(|id| (id.clone(), world.read_meta(&id))).collect();
        Ok(Arc::new(Self {
            world,
            runtime,
            mailboxes: Mutex::new(HashMap::new()),
            controls: Mutex::new(HashMap::new()),
            metas: Mutex::new(metas),
            op_locks: Mutex::new(HashMap::new()),
            seq: AtomicU64::new(1),
            exit: tokio::sync::Notify::new(),
        }))
    }

    /// Wait until an actor sends `{"type": "exit"}` to `broker`. The `sausage`
    /// binary then stops all actors and exits, so that a wrapper such as
    /// `nix run .#loop` can start a fresh (updated) host.
    pub async fn exit_requested(&self) {
        self.exit.notified().await;
    }

    pub fn meta(&self, id: &str) -> Option<ActorMeta> {
        self.metas.lock().unwrap().get(id).cloned()
    }

    /// Change the `actor.json` of `id` (creating it if needed) and save it.
    fn update_meta(&self, id: &str, f: impl FnOnce(&mut ActorMeta)) -> Result<()> {
        let mut metas = self.metas.lock().unwrap();
        let meta = metas.entry(id.to_string()).or_insert_with(|| ActorMeta { running: true, ..Default::default() });
        f(meta);
        self.world.write_meta(id, meta)
    }

    /// Create the mailbox of an actor. Messages wait here while the actor is
    /// down, during a swap, or while it restarts.
    pub fn open_mailbox(&self, id: &str) -> mpsc::Receiver<Envelope> {
        let (tx, rx) = mpsc::channel(MAILBOX_SIZE);
        self.mailboxes.lock().unwrap().insert(id.to_string(), tx);
        rx
    }

    /// Send a message. `to: None` publishes to the subscribers of `from`.
    /// With `own_id`, the message ID is `"<from>/<own_id>"`. Returns the ID.
    pub fn send(&self, from: &str, to: Option<&str>, own_id: Option<&str>, reply_to: Option<&str>, body: Value) -> String {
        let id = match own_id {
            Some(own) => format!("{from}/{own}"),
            None => format!("m-{}", self.seq.fetch_add(1, Ordering::Relaxed)),
        };
        let recipients: Vec<String> = match to {
            Some(to) => vec![to.to_string()],
            None => self
                .metas
                .lock()
                .unwrap()
                .iter()
                .filter(|(_, m)| m.subscriptions.contains(from))
                .map(|(id, _)| id.clone())
                .collect(),
        };
        for to in recipients {
            self.deliver(Envelope {
                id: id.clone(),
                from: from.to_string(),
                to,
                reply_to: reply_to.map(str::to_string),
                body: body.clone(),
            });
        }
        id
    }

    /// Route one stdout line of a WASM actor.
    pub fn route_line(&self, from: &str, line: &str) {
        let line = line.trim();
        if line.is_empty() {
            return;
        }
        let Ok(Value::Object(obj)) = serde_json::from_str::<Value>(line) else {
            self.log(from, &format!("stdout line is not a JSON object, ignored: {}", truncate(line, 500)));
            return;
        };
        let Some(body) = obj.get("body").cloned() else {
            self.log(from, &format!("message without \"body\", ignored: {}", truncate(line, 500)));
            return;
        };
        let own_id = obj.get("id").and_then(Value::as_str);
        let to = obj.get("to").and_then(Value::as_str);
        let reply_to = obj.get("reply_to").and_then(Value::as_str);
        self.send(from, to, own_id, reply_to, body);
    }

    fn deliver(&self, env: Envelope) {
        let tx = self.mailboxes.lock().unwrap().get(&env.to).cloned();
        let error = match tx {
            None => format!("no actor with ID {}", env.to),
            Some(tx) => match tx.try_send(env.clone()) {
                Ok(()) => return,
                Err(mpsc::error::TrySendError::Full(_)) => format!("mailbox of {} is full", env.to),
                Err(mpsc::error::TrySendError::Closed(_)) => format!("actor {} is gone", env.to),
            },
        };
        if env.from != "broker" {
            let body = json!({"type": "undeliverable", "reason": error, "to": env.to});
            self.send("broker", Some(&env.from), None, Some(&env.id), body);
        }
    }

    /// Append a host note to the stderr log of an actor.
    pub fn log(&self, id: &str, text: &str) {
        use std::io::Write;
        let path = self.world.actor_dir(id).join("stderr.log");
        if let Ok(mut f) = std::fs::OpenOptions::new().create(true).append(true).open(path) {
            let _ = writeln!(f, "[sausage {}] {text}", world::now());
        }
    }

    fn op_lock(&self, id: &str) -> Arc<tokio::sync::Mutex<()>> {
        self.op_locks.lock().unwrap().entry(id.to_string()).or_default().clone()
    }

    /// What an instance of `id` gets (see PROTOCOL.md). A candidate runs in
    /// `scratch`, which holds a copy of the actor's `data/` and `public/`,
    /// and has no network.
    pub fn instance_config(&self, id: &str, version: u64, scratch: Option<&Path>) -> InstanceConfig {
        let own = scratch.map_or_else(|| self.world.actor_dir(id), Path::to_path_buf);
        let mut env = vec![("SAUSAGE_ACTOR".into(), id.to_string()), ("SAUSAGE_VERSION".into(), version.to_string())];
        if scratch.is_some() {
            env.push(("SAUSAGE_CANDIDATE".into(), "1".into()));
        }
        InstanceConfig {
            mounts: vec![
                Mount { guest: "/actor".into(), host: self.world.actors_root(), writable: false },
                Mount { guest: "/public".into(), host: own.join("public"), writable: true },
                Mount { guest: "/data".into(), host: own.join("data"), writable: true },
            ],
            network: scratch.is_none(),
            stderr_log: own.join("stderr.log"),
            env,
        }
    }

    // ---------------------------------------------------------------------
    // Lifecycle
    // ---------------------------------------------------------------------

    fn start_supervisor(self: &Arc<Self>, id: &str, version: u64, running: bool) {
        let mailbox = self.open_mailbox(id);
        let (ctl_tx, ctl_rx) = mpsc::channel(16);
        self.controls.lock().unwrap().insert(id.to_string(), ctl_tx);
        tokio::spawn(Supervisor::new(self.clone(), id.to_string(), version, running, mailbox, ctl_rx).run());
    }

    fn is_supervised(&self, id: &str) -> bool {
        self.controls.lock().unwrap().contains_key(id)
    }

    async fn control(&self, id: &str, make: impl FnOnce(Reply) -> Control) -> Result<Value> {
        let tx = self.controls.lock().unwrap().get(id).cloned().ok_or_else(|| anyhow!("no WASM actor with ID {id}"))?;
        let (reply, rx) = oneshot::channel();
        tx.send(make(reply)).await.map_err(|_| anyhow!("supervisor of {id} is gone"))?;
        rx.await.map_err(|_| anyhow!("supervisor of {id} did not answer"))?
    }

    /// Stop every actor gracefully, in parallel (host shutdown). The
    /// `running` flags stay, so the next `sausage run` starts them again.
    pub async fn shutdown_all(self: &Arc<Self>) {
        let ids: Vec<String> = self.controls.lock().unwrap().keys().cloned().collect();
        let mut tasks = tokio::task::JoinSet::new();
        for id in ids {
            let this = self.clone();
            tasks.spawn(async move { this.control(&id, |reply| Control::Shutdown { reply }).await });
        }
        tasks.join_all().await;
    }

    /// Compile `versions/<version>/actor.wasm`.
    pub fn load_version(&self, id: &str, version: u64) -> Result<Component> {
        self.load(id, &self.world.version_dir(id, version).join("actor.wasm"))
    }

    fn load(&self, id: &str, wasm: &Path) -> Result<Component> {
        let started = Instant::now();
        let component = self.runtime.load(wasm)?;
        self.log(id, &format!("compiled {} in {} ms", wasm.display(), started.elapsed().as_millis()));
        Ok(component)
    }

    /// Test a new version where it lies. Only when the tests pass, copy it
    /// to `versions/<n>/`: that folder holds only versions that run(ran).
    /// Returns the component and the test report.
    async fn try_version(&self, id: &str, meta: &VersionMeta, wasm: &Path, source: Option<&Path>) -> Result<(Component, Value)> {
        let component = self.load(id, wasm)?;
        let tests = candidate::test(self, id, meta.version, &component, source).await;
        if let Err(e) = &tests {
            self.log(id, &format!("version {} from {} failed its tests: {e:#}", meta.version, meta.wasm));
        }
        let tests = tests?;
        let dir = self.world.version_dir(id, meta.version);
        std::fs::create_dir_all(&dir)?;
        std::fs::copy(wasm, dir.join("actor.wasm")).with_context(|| format!("copy {}", wasm.display()))?;
        world::replace_source(source.unwrap_or(&dir.join("no-source")), &dir.join("source"))?;
        world::write_atomic(&dir.join("meta.json"), serde_json::to_string_pretty(meta)?.as_bytes())?;
        Ok((component, tests))
    }

    /// Make `version` the current `actor.wasm`, `source/`, and `actor.json`.
    pub fn install(&self, id: &str, version: u64) -> Result<()> {
        self.world.install_files(id, version)?;
        self.update_meta(id, |m| {
            m.version = Some(version);
            m.running = true;
        })
    }

    pub fn set_running(&self, id: &str, running: bool) -> Result<()> {
        self.update_meta(id, |m| m.running = running)
    }

    /// Create a new actor from a `.wasm` file: test it, install it, start it.
    pub async fn create_actor(self: &Arc<Self>, id: &str, wasm: &Path, source: Option<&Path>, meta: VersionMeta) -> Result<Value> {
        if !world::valid_id(id) {
            bail!("invalid actor ID {id:?}: use a-z, 0-9, '-' and '_'");
        }
        let lock = self.op_lock(id);
        let _guard = lock.lock().await;
        if NATIVE.contains(&id) || self.is_supervised(id) || self.world.wasm_path(id).exists() {
            bail!("an actor with ID {id} already exists");
        }
        let meta = VersionMeta { version: self.world.next_version(id), ..meta };
        self.create_version(id, &meta, wasm, source).await
    }

    async fn create_version(self: &Arc<Self>, id: &str, meta: &VersionMeta, wasm: &Path, source: Option<&Path>) -> Result<Value> {
        let (_, tests) = self.try_version(id, meta, wasm, source).await?;
        self.install(id, meta.version)?;
        self.start_supervisor(id, meta.version, true);
        Ok(json!({"type": "spawned", "actor": id, "version": meta.version, "tests": tests}))
    }

    pub async fn upgrade_actor(self: &Arc<Self>, id: &str, wasm: &Path, source: Option<&Path>, meta: VersionMeta) -> Result<Value> {
        let lock = self.op_lock(id);
        let _guard = lock.lock().await;
        if !self.is_supervised(id) {
            bail!("no WASM actor with ID {id}");
        }
        let version = self.world.next_version(id);
        self.upgrade_version(id, &VersionMeta { version, ..meta }, wasm, source).await
    }

    async fn upgrade_version(self: &Arc<Self>, id: &str, meta: &VersionMeta, wasm: &Path, source: Option<&Path>) -> Result<Value> {
        let (component, tests) = self.try_version(id, meta, wasm, source).await?;
        let (version, requester) = (meta.version, meta.requested_by.clone());
        let mut result = self.control(id, |reply| Control::Upgrade { version, component, requester, reply }).await?;
        result["tests"] = tests;
        Ok(result)
    }

    /// Boot the actors in `actor/`. Called at start and on `scan`.
    ///
    /// - An `actor.wasm` equal to the installed version starts as is.
    /// - A new or changed `actor.wasm` is tested where it lies. A new actor
    ///   is installed, a changed one is an upgrade. If its tests fail,
    ///   nothing changes, and the file stays for the next scan.
    pub async fn scan(self: &Arc<Self>) -> Vec<Value> {
        let mut report = Vec::new();
        for id in self.world.actor_ids().unwrap_or_default() {
            if NATIVE.contains(&id.as_str()) || !self.world.wasm_path(&id).is_file() {
                continue;
            }
            let lock = self.op_lock(&id);
            let _guard = lock.lock().await;
            if let Some(r) = self.scan_actor(&id).await {
                report.push(r);
            }
        }
        report
    }

    async fn scan_actor(self: &Arc<Self>, id: &str) -> Option<Value> {
        let meta = self.world.read_meta(id);
        let installed = meta.version.filter(|v| self.world.version_dir(id, *v).join("actor.wasm").is_file());
        self.metas.lock().unwrap().insert(id.to_string(), meta.clone());
        if let Some(v) = installed
            && !self.is_supervised(id)
        {
            self.start_supervisor(id, v, meta.running);
        }
        let wasm = self.world.wasm_path(id);
        if installed.is_some_and(|v| same_file(&wasm, &self.world.version_dir(id, v).join("actor.wasm"))) {
            return None;
        }
        let source = self.world.source_dir(id);
        let vmeta = VersionMeta {
            version: self.world.next_version(id),
            wasm: format!("/actor/{id}/actor.wasm"),
            source: Some(format!("/actor/{id}/source")),
            built_by: None,
            requested_by: "scan".into(),
            created_at: world::now(),
        };
        let source = source.is_dir().then_some(source.as_path());
        let result = match installed {
            None => self.create_version(id, &vmeta, &wasm, source).await,
            Some(_) => self.upgrade_version(id, &vmeta, &wasm, source).await,
        };
        Some(result.unwrap_or_else(|e| json!({"type": "error", "actor": id, "reason": format!("{e:#}")})))
    }

    /// Handle one message to `broker`.
    async fn handle(self: Arc<Self>, env: Envelope) {
        let body = self.clone().command(&env).await.unwrap_or_else(|e| {
            json!({"type": "error", "request": env.body.get("type"), "reason": format!("{e:#}")})
        });
        self.send("broker", Some(&env.from), None, Some(&env.id), body);
    }

    async fn command(self: Arc<Self>, env: &Envelope) -> Result<Value> {
        let b = &env.body;
        let ty = b["type"].as_str().unwrap_or("");
        let actor = || b["actor"].as_str().map(str::to_string).ok_or_else(|| anyhow!("missing \"actor\""));
        let version_meta = || -> Result<(std::path::PathBuf, Option<std::path::PathBuf>, VersionMeta)> {
            let wasm_s = b["wasm"].as_str().ok_or_else(|| anyhow!("missing \"wasm\""))?;
            let wasm = self.world.resolve(&env.from, wasm_s)?;
            if !wasm.is_file() {
                bail!("{wasm_s} is not a file");
            }
            let source_s = b["source"].as_str();
            let source = source_s.map(|s| self.world.resolve(&env.from, s)).transpose()?;
            if let (Some(s), Some(p)) = (source_s, &source)
                && !p.is_dir()
            {
                bail!("{s} is not a folder");
            }
            let meta = VersionMeta {
                version: 0,
                wasm: wasm_s.to_string(),
                source: source_s.map(str::to_string),
                built_by: b["built_by"].as_str().map(str::to_string),
                requested_by: env.from.clone(),
                created_at: world::now(),
            };
            Ok((wasm, source, meta))
        };
        match ty {
            "spawn" => {
                let (wasm, source, meta) = version_meta()?;
                self.create_actor(&actor()?, &wasm, source.as_deref(), meta).await
            }
            "upgrade" => {
                let (wasm, source, meta) = version_meta()?;
                self.upgrade_actor(&actor()?, &wasm, source.as_deref(), meta).await
            }
            "rollback" => {
                let id = actor()?;
                let version = b["version"].as_u64().ok_or_else(|| anyhow!("missing \"version\""))?;
                let lock = self.op_lock(&id);
                let _guard = lock.lock().await;
                self.control(&id, |reply| Control::Rollback { version, reply }).await
            }
            "stop" => self.control(&actor()?, |reply| Control::Stop { reply }).await,
            "start" => self.control(&actor()?, |reply| Control::Start { reply }).await,
            "subscribe" | "unsubscribe" => {
                let publisher = actor()?;
                let on = ty == "subscribe";
                self.update_meta(&env.from, |m| {
                    if on {
                        m.subscriptions.insert(publisher.clone());
                    } else {
                        m.subscriptions.remove(&publisher);
                    }
                })?;
                Ok(json!({"type": format!("{ty}d"), "actor": publisher}))
            }
            "scan" => Ok(json!({"type": "scanned", "changes": self.scan().await})),
            "exit" => {
                self.log("broker", &format!("exit requested by {}", env.from));
                // notify_one keeps the permit if nobody waits yet.
                self.exit.notify_one();
                Ok(json!({"type": "exiting"}))
            }
            other => bail!("unknown broker request {other:?}. Known: spawn, upgrade, rollback, stop, start, subscribe, unsubscribe, scan, exit"),
        }
    }

    /// The native `broker` actor. Requests run concurrently. Aborting this
    /// task aborts the requests that still run. Replies (messages with
    /// `reply_to`) get no answer, so that two actors never answer each other
    /// forever.
    pub async fn run_actor(self: Arc<Self>, mut mailbox: mpsc::Receiver<Envelope>) {
        let mut requests = tokio::task::JoinSet::new();
        loop {
            tokio::select! {
                env = mailbox.recv() => match env {
                    Some(env) if env.reply_to.is_none() => { requests.spawn(self.clone().handle(env)); }
                    Some(_) => {}
                    None => break,
                },
                Some(_) = requests.join_next() => {}
            }
        }
    }
}

fn same_file(a: &Path, b: &Path) -> bool {
    matches!((std::fs::read(a), std::fs::read(b)), (Ok(x), Ok(y)) if x == y)
}

pub fn truncate(s: &str, max: usize) -> String {
    if s.len() <= max {
        return s.to_string();
    }
    let mut end = max;
    while !s.is_char_boundary(end) {
        end -= 1;
    }
    format!("{}… ({} bytes)", &s[..end], s.len())
}