//! `atlibs bridge`: the Rust run-module engine driven over stdio (JSON lines), so the JS
//! conformance harness (rust/host/conformance.mjs) can hand it the same ctx and call stubs it
//! hands the JS runModule, and compare. Not a user surface.
//!
//! stdin, first line: { module: { wasm?: path, calls?, anyRecords?, sha256?, owner? }, entry, input,
//!   budget?, inputs?: bool, ledger?: bool, now?: iso, didkey?: { jwk, services }, http?: runModule's http, held?: runModule's held }
//!   (or { rfpPolicyRoutes: { manifest, from, records } }: one { t: done, manifest | error } line, no run)
//! then replies { k, ... } and unsolicited { t: input, data } | { t: inputs-end, error? } | { t: abort }.
//! stdout: { t: run | call | send | close | abort | ledger | fetch, k, ... } and last { t: done, ran | error }.
//! Without module.wasm, ctx.wasm.run goes over the bridge too (the JS harness's own runner).

use std::collections::HashMap;
use std::io::Write;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use serde_json::{Map, Value, json};
use tokio::io::AsyncBufReadExt;
use tokio::sync::{mpsc, oneshot};

use atlibs_host::crypto::DidKey;
use atlibs_host::host::{DidkeyCall, Fetch, Response};
use atlibs_host::run::{self, Call, Ctx, Ledger, Module, NOSEND, Options, Request, Runner, Sub, SubCmd};
use atlibs_host::wasm::{Bound, Wasm};
use atlibs_host::{BoxFut, Signal, js_string};

enum Waiter {
    Once(oneshot::Sender<Value>),
    Stream(mpsc::UnboundedSender<Value>),
}

#[derive(Default)]
struct Io {
    next: AtomicU64,
    waiters: Mutex<HashMap<u64, Waiter>>,
    /// request.emit of pending calls, for { t: emit, k, event }
    emits: Mutex<HashMap<u64, run::Emit>>,
}

impl Io {
    fn write(&self, message: Value) {
        let mut out = std::io::stdout().lock();
        let _ = writeln!(out, "{message}");
        let _ = out.flush();
    }
    fn key(&self) -> u64 {
        self.next.fetch_add(1, Ordering::SeqCst) + 1
    }
    async fn ask(&self, mut message: Value) -> Value {
        let (k, (tx, rx)) = (self.key(), oneshot::channel());
        self.waiters.lock().unwrap().insert(k, Waiter::Once(tx));
        message["k"] = json!(k);
        self.write(message);
        rx.await
            .unwrap_or_else(|_| json!({ "ok": false, "error": "bridge closed" }))
    }
    fn deliver(&self, reply: Value) {
        let Some(k) = reply.get("k").and_then(Value::as_u64) else {
            return;
        };
        let mut waiters = self.waiters.lock().unwrap();
        match waiters.get(&k) {
            Some(Waiter::Once(_)) => {
                if let Some(Waiter::Once(tx)) = waiters.remove(&k) {
                    let _ = tx.send(reply);
                }
            }
            Some(Waiter::Stream(tx)) => {
                let _ = tx.send(reply);
            }
            None => {}
        }
    }
}

fn outcome(reply: &Value) -> Result<Option<Value>, String> {
    if reply.get("ok") == Some(&Value::Bool(true)) {
        Ok(reply.get("output").cloned())
    } else {
        Err(js_string(reply.get("error")))
    }
}

struct BridgeRunner(Arc<Io>);

impl Runner for BridgeRunner {
    fn run(&self, entry: String, input: Value, io_bytes: usize) -> BoxFut<'_, Result<Value, String>> {
        Box::pin(async move {
            outcome(
                &self
                    .0
                    .ask(json!({ "t": "run", "entry": entry, "input": input, "ioBytes": io_bytes }))
                    .await,
            )
            .map(|o| o.unwrap_or(Value::Null))
        })
    }
}

struct BridgeLedger(Arc<Io>);

// The JS ledger's errors come back as { ok: false, error } and are the ledger's errors here too.
fn ledger_reply(reply: Value) -> Result<Value, String> {
    match reply.get("ok") {
        Some(Value::Bool(false)) => Err(js_string(reply.get("error"))),
        _ => Ok(reply),
    }
}

impl Ledger for BridgeLedger {
    fn get(&self, uri: &str) -> BoxFut<'_, Result<Option<Value>, String>> {
        let uri = uri.to_string();
        Box::pin(async move {
            let reply = ledger_reply(self.0.ask(json!({ "t": "ledger", "op": "get", "uri": uri })).await)?;
            Ok(reply.get("value").filter(|v| !v.is_null()).cloned())
        })
    }
    fn put(&self, uri: String, value: Value) -> BoxFut<'_, Result<(), String>> {
        Box::pin(async move { ledger_reply(self.0.ask(json!({ "t": "ledger", "op": "put", "uri": uri, "value": value })).await).map(drop) })
    }
    fn delete(&self, uri: String) -> BoxFut<'_, Result<(), String>> {
        Box::pin(async move { ledger_reply(self.0.ask(json!({ "t": "ledger", "op": "delete", "uri": uri })).await).map(drop) })
    }
}

/// The JS side's ctx.http.fetch; with a scheme ("webrtc:"), its fetchers[scheme].
struct BridgeFetch(Arc<Io>, Option<String>);

impl Fetch for BridgeFetch {
    fn fetch(
        &self,
        method: &str,
        url: &str,
        headers: Vec<(String, String)>,
        body: Option<Vec<u8>>,
    ) -> BoxFut<'_, Result<Response, String>> {
        let headers: Map<String, Value> = headers.into_iter().map(|(k, v)| (k, json!(v))).collect();
        let message = json!({ "t": "fetch", "scheme": self.1, "method": method, "url": url, "headers": headers, "body": body.map(|b| String::from_utf8_lossy(&b).to_string()) });
        Box::pin(async move {
            let reply = self.0.ask(message).await;
            Ok(Response {
                status: reply.get("status").and_then(Value::as_u64).unwrap_or(200) as u16,
                headers: reply.get("headers").and_then(Value::as_object).map(|h| h.iter().map(|(k, v)| (k.clone(), js_string(Some(v)))).collect()).unwrap_or_default(),
                body: js_string(reply.get("body")).into_bytes(),
            })
        })
    }
}

/// The JS side's call; session: its `session` option's (runModule's, which mints http serviceAuth) instead.
struct BridgeCall(Arc<Io>, bool);

impl BridgeCall {
    // The JS stub sees signal.aborted (with this reason) while the call is pending, and when the
    // Rust side drops the call unfinished (a subscription stopped while opening), as dropping a real
    // call's future cancels it.
    fn watch_abort(&self, k: u64, signal: Signal, done: oneshot::Receiver<()>) {
        let io = self.0.clone();
        tokio::spawn(async move {
            tokio::select! {
                biased;
                reason = signal.aborted() => io.write(json!({ "t": "abort", "k": k, "reason": reason })),
                done = done => if done.is_err() {
                    io.write(json!({ "t": "abort", "k": k, "reason": signal.reason().unwrap_or_else(|| "cancelled".into()) }));
                },
            }
        });
    }
}

impl Call for BridgeCall {
    fn call(&self, nsid: String, request: Request) -> BoxFut<'_, Result<Option<Value>, String>> {
        Box::pin(async move {
            let (k, (tx, rx)) = (self.0.key(), oneshot::channel());
            self.0.waiters.lock().unwrap().insert(k, Waiter::Once(tx));
            let (done_tx, done_rx) = oneshot::channel();
            self.watch_abort(k, request.signal.clone(), done_rx);
            if let Some(emit) = request.emit.clone() {
                self.0.emits.lock().unwrap().insert(k, emit);
            }
            self.0
                .write(json!({ "t": "call", "k": k, "nsid": nsid, "request": request.to_json(), "session": self.1 }));
            let reply = rx
                .await
                .unwrap_or_else(|_| json!({ "ok": false, "error": "bridge closed" }));
            self.0.emits.lock().unwrap().remove(&k);
            let _ = done_tx.send(());
            outcome(&reply)
        })
    }

    fn subscribe(&self, nsid: String, request: Request) -> BoxFut<'_, Result<Sub, String>> {
        Box::pin(async move {
            let k = self.0.key();
            let (tx, mut replies) = mpsc::unbounded_channel();
            self.0.waiters.lock().unwrap().insert(k, Waiter::Stream(tx));
            let (done_tx, done_rx) = oneshot::channel();
            self.watch_abort(k, request.signal.clone(), done_rx);
            self.0
                .write(json!({ "t": "call", "k": k, "nsid": nsid, "request": request.to_json(), "session": self.1 }));
            let first = replies
                .recv()
                .await
                .unwrap_or_else(|| json!({ "ok": false, "error": "bridge closed" }));
            if first.get("open") != Some(&Value::Bool(true)) {
                self.0.waiters.lock().unwrap().remove(&k);
                return Err(js_string(first.get("error")));
            }
            let (msg_tx, messages) = mpsc::unbounded_channel();
            let (commands, mut cmd_rx) = mpsc::unbounded_channel::<SubCmd>();
            let io = self.0.clone();
            tokio::spawn(async move {
                while let Some(reply) = replies.recv().await {
                    if let Some(data) = reply.get("message") {
                        let _ = msg_tx.send(Ok(data.clone()));
                    } else if reply.get("end").is_some() {
                        if let Some(e) = reply.get("error") {
                            let _ = msg_tx.send(Err(js_string(Some(e))));
                        }
                        break;
                    }
                }
                io.waiters.lock().unwrap().remove(&k);
            });
            let io = self.0.clone();
            tokio::spawn(async move {
                while let Some(command) = cmd_rx.recv().await {
                    match command {
                        SubCmd::Send(data, ack) => {
                            let reply =
                                io.ask(json!({ "t": "send", "sub": k, "data": data })).await;
                            let _ = ack.send(if reply.get("nosend").is_some() {
                                Err(NOSEND.into())
                            } else {
                                outcome(&reply).map(drop)
                            });
                        }
                        SubCmd::Close => break,
                    }
                }
                io.write(json!({ "t": "close", "sub": k }));
                let _ = done_tx.send(());
            });
            Ok(Sub { messages, commands })
        })
    }
}

pub async fn main() -> Result<(), String> {
    let io = Arc::new(Io::default());
    let mut lines = tokio::io::BufReader::new(tokio::io::stdin()).lines();
    let start: Value = serde_json::from_str(
        &lines
            .next_line()
            .await
            .map_err(|e| e.to_string())?
            .ok_or("bridge: no start line")?,
    )
    .map_err(|e| e.to_string())?;
    // { rfpPolicyRoutes: { manifest, from, records } }: compose::rfp_policy_routes alone, no run
    if let Some(ask) = start.get("rfpPolicyRoutes") {
        let records = ask["records"].as_array().cloned().unwrap_or_default();
        io.write(match atlibs_host::compose::rfp_policy_routes(&ask["manifest"], &js_string(ask.get("from")), &records) {
            Ok(manifest) => json!({ "t": "done", "manifest": manifest }),
            Err(error) => json!({ "t": "done", "error": error }),
        });
        return Ok(());
    }
    let (input_tx, input_rx) = mpsc::unbounded_channel();
    let signal = Signal::new();
    {
        let (io, signal) = (io.clone(), signal.clone());
        tokio::spawn(async move {
            let mut input_tx = Some(input_tx);
            while let Ok(Some(line)) = lines.next_line().await {
                let Ok(message) = serde_json::from_str::<Value>(&line) else {
                    continue;
                };
                match message.get("t").and_then(Value::as_str) {
                    Some("input") => drop(input_tx.as_ref().map(|tx| {
                        tx.send(Ok(message.get("data").cloned().unwrap_or(Value::Null)))
                    })),
                    Some("inputs-end") => {
                        if let (Some(tx), Some(e)) = (&input_tx, message.get("error")) {
                            let _ = tx.send(Err(js_string(Some(e))));
                        }
                        input_tx = None;
                    }
                    Some("abort") => signal.abort("aborted"),
                    Some("emit") => {
                        let emit = message
                            .get("k")
                            .and_then(Value::as_u64)
                            .and_then(|k| io.emits.lock().unwrap().get(&k).cloned());
                        if let Some(emit) = emit {
                            emit.emit(message.get("event").cloned().unwrap_or(Value::Null));
                        }
                    }
                    _ => io.deliver(message),
                }
            }
        });
    }
    let spec = &start["module"];
    // a refused module is the run's error, as the JS runModule throws it
    let mut module = match Module::from_manifest(vec![], spec) {
        Ok(module) => module,
        Err(error) => {
            io.write(json!({ "t": "done", "error": error }));
            return Ok(());
        }
    };
    module.sha256 = spec.get("sha256").and_then(Value::as_str).map(String::from);
    // the JS side's moduleOwner (it has the bytes when ctx.wasm.run is bridged)
    module.owner = spec.get("owner").and_then(Value::as_str).map(String::from);
    let runner: Arc<dyn Runner> = match spec.get("wasm").and_then(Value::as_str) {
        Some(path) => {
            module.bytes = Arc::new(std::fs::read(path).map_err(|e| format!("{path}: {e}"))?);
            Arc::new(Bound {
                wasm: Wasm::new(atlibs_host::wasm::WALL_TIME_MS),
                bytes: module.bytes.clone(),
                sha256: module.sha256.clone(),
            })
        }
        None => Arc::new(BridgeRunner(io.clone())),
    };
    let ledger: Option<Arc<dyn Ledger>> = if start.get("ledger") == Some(&Value::Bool(true)) {
        Some(Arc::new(BridgeLedger(io.clone())))
    } else {
        None
    };
    let ctx = Ctx {
        runner,
        ledger,
        now: start.get("now").and_then(Value::as_str).map(String::from),
    };
    let mut call: Arc<dyn Call> = Arc::new(BridgeCall(io.clone(), false));
    if let Some(didkey) = start.get("didkey") {
        let services = didkey
            .get("services")
            .and_then(Value::as_array)
            .map(|s| {
                s.iter()
                    .filter_map(|x| x.as_str().map(String::from))
                    .collect()
            })
            .unwrap_or_default();
        call = Arc::new(DidkeyCall {
            key: DidKey::new(didkey.get("jwk"))?,
            services,
            fetch: Arc::new(BridgeFetch(io.clone(), None)),
            inner: call,
            store: None, // the suite's own did:key, no state dir
        });
    }
    let entry = js_string(start.get("entry"));
    let input = start
        .get("input")
        .and_then(Value::as_object)
        .cloned()
        .unwrap_or_default();
    let options = Options {
        signal: Some(signal),
        inputs: (start.get("inputs") == Some(&Value::Bool(true))).then_some(input_rx),
        budget: start.get("budget").cloned(),
        on_event: (start.get("onEvent") == Some(&Value::Bool(true))).then(|| {
            let io = io.clone();
            Arc::new(move |event: &Value, out: &Value| {
                io.write(json!({ "t": "event", "event": event, "out": out }))
            }) as run::OnEvent
        }),
        load: None,
        provides: None,
        http: match start.get("http") {
            Some(allow) => {
                let mut http = run::Http::from_json(Arc::new(BridgeFetch(io.clone(), None)), allow)?;
                for scheme in start["fetchers"].as_array().into_iter().flatten().filter_map(Value::as_str) {
                    http.schemes.insert(scheme.into(), Arc::new(BridgeFetch(io.clone(), Some(scheme.into()))));
                }
                Some(Arc::new(http))
            }
            None => None,
        },
        held: start.get("held").and_then(Value::as_object).cloned(),
        session: (start.get("session") == Some(&Value::Bool(true))).then(|| Arc::new(BridgeCall(io.clone(), true)) as Arc<dyn Call>),
    };
    let ran = run::run_module(&ctx, &module, &entry, &input, call, options).await;
    // Let the run end's aborts and closes reach the stubs before the process exits.
    tokio::time::sleep(std::time::Duration::from_millis(5)).await;
    io.write(json!({ "t": "done", "ran": ran }));
    Ok(())
}