Codex

use serde_json::{json, Value};
use std::{
    collections::{HashMap, HashSet, VecDeque},
    io::{BufRead, BufReader, Write},
    path::Path,
    process::{Child, ChildStdin, Command, Stdio},
    sync::{
        atomic::{AtomicBool, AtomicU64, Ordering},
        mpsc, Arc, Mutex,
    },
    thread,
    time::Duration,
};

pub struct Rpc {
    child: Mutex<Child>,
    input: Mutex<ChildStdin>,
    pending: Mutex<HashMap<u64, mpsc::Sender<Value>>>,
    requests: Mutex<HashMap<String, Value>>,
    events: Mutex<VecDeque<Value>>,
    active_turns: Mutex<HashSet<String>>,
    next: AtomicU64,
    sequence: AtomicU64,
    alive: AtomicBool,
    pub mode: String,
}

impl Rpc {
    pub fn start(executable: &Path, proxy: bool) -> Result<Arc<Self>, String> {
        let mut cmd = Command::new(executable);
        cmd.arg("app-server");
        if proxy {
            cmd.arg("proxy");
        } else {
            cmd.arg("--stdio");
        }
        Self::start_command(cmd, proxy)
    }
    fn start_command(mut cmd: Command, proxy: bool) -> Result<Arc<Self>, String> {
        super::process::hidden(&mut cmd)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null());
        let mut child = cmd.spawn().map_err(|e| e.to_string())?;
        let input = child.stdin.take().unwrap();
        let output = child.stdout.take().unwrap();
        let rpc = Arc::new(Self {
            child: Mutex::new(child),
            input: Mutex::new(input),
            pending: Mutex::new(HashMap::new()),
            requests: Mutex::new(HashMap::new()),
            events: Mutex::new(VecDeque::new()),
            active_turns: Mutex::new(HashSet::new()),
            next: AtomicU64::new(1),
            sequence: AtomicU64::new(0),
            alive: AtomicBool::new(true),
            mode: if proxy {
                "desktop-daemon-proxy"
            } else {
                "bridge-owned-app-server"
            }
            .into(),
        });
        let weak = Arc::downgrade(&rpc);
        thread::spawn(move || {
            for line in BufReader::new(output).lines() {
                let Ok(line) = line else {
                    break;
                };
                let Ok(value) = serde_json::from_str::<Value>(&line) else {
                    continue;
                };
                let Some(rpc) = weak.upgrade() else {
                    break;
                };
                if value.get("method").is_none() {
                    if let Some(id) = value["id"].as_u64() {
                        if let Some(tx) = rpc.pending.lock().unwrap().remove(&id) {
                            let _ = tx.send(value);
                        }
                    }
                } else {
                    if let Some(thread_id) = value["params"]["threadId"].as_str() {
                        match value["method"].as_str() {
                            Some("turn/started") => {
                                rpc.active_turns
                                    .lock()
                                    .unwrap()
                                    .insert(thread_id.to_string());
                            }
                            Some("turn/completed") => {
                                rpc.active_turns.lock().unwrap().remove(thread_id);
                            }
                            _ => {}
                        }
                    }
                    if value.get("id").is_some() {
                        rpc.requests
                            .lock()
                            .unwrap()
                            .insert(value["id"].to_string(), value.clone());
                    }
                    let seq = rpc.sequence.fetch_add(1, Ordering::Relaxed) + 1;
                    let mut events = rpc.events.lock().unwrap();
                    events.push_back(json!({"seq":seq,"event":value}));
                    while events.len() > 1000 {
                        events.pop_front();
                    }
                }
            }
            if let Some(rpc) = weak.upgrade() {
                rpc.alive.store(false, Ordering::Relaxed);
                rpc.pending.lock().unwrap().clear();
            }
        });
        let init = rpc.call("initialize", json!({"clientInfo":{"name":"adom_codex_bridge","title":"Adom Codex Bridge","version":env!("CARGO_PKG_VERSION")},"capabilities":{"experimentalApi":true}}), Duration::from_secs(15));
        if let Err(e) = init {
            rpc.stop();
            return Err(e);
        }
        rpc.write(json!({"method":"initialized","params":{}}))?;
        Ok(rpc)
    }
    fn write(&self, value: Value) -> Result<(), String> {
        let mut input = self.input.lock().unwrap();
        writeln!(input, "{value}")
            .and_then(|_| input.flush())
            .map_err(|e| e.to_string())
    }
    pub fn call(&self, method: &str, params: Value, timeout: Duration) -> Result<Value, String> {
        if !self.alive.load(Ordering::Relaxed) {
            return Err("app_server_disconnected".into());
        }
        let id = self.next.fetch_add(1, Ordering::Relaxed);
        let (tx, rx) = mpsc::channel();
        self.pending.lock().unwrap().insert(id, tx);
        if let Err(e) = self.write(json!({"id":id,"method":method,"params":params})) {
            self.pending.lock().unwrap().remove(&id);
            return Err(e);
        }
        let response = rx.recv_timeout(timeout);
        self.pending.lock().unwrap().remove(&id);
        let v = response.map_err(|e| {
            format!("rpc_wait_failed: {e}; inspect codex_events and codex_requests before retrying")
        })?;
        if let Some(err) = v.get("error") {
            return Err(format!("app_server_error: {err}"));
        }
        Ok(v["result"].clone())
    }
    pub fn events(&self, since: u64) -> Value {
        let events = self.events.lock().unwrap();
        json!({"events":events.iter().filter(|e| e["seq"].as_u64().unwrap_or(0)>since).collect::<Vec<_>>(),
            "latestSeq":self.sequence.load(Ordering::Relaxed),"oldestSeq":events.front().map(|v|&v["seq"]),
            "alive":self.alive.load(Ordering::Relaxed),"mode":self.mode})
    }
    pub fn requests(&self) -> Value {
        json!(self.requests.lock().unwrap().values().collect::<Vec<_>>())
    }
    pub fn reply(&self, id: Value, result: Value) -> Result<Value, String> {
        let mut pending = self.requests.lock().unwrap();
        if !pending.contains_key(&id.to_string()) {
            return Err("unknown_request_id".into());
        }
        self.write(json!({"id":id,"result":result}))?;
        pending.remove(&id.to_string());
        Ok(json!({"replied":true}))
    }
    pub fn stop(&self) {
        self.alive.store(false, Ordering::Relaxed);
        let mut child = self.child.lock().unwrap();
        let _ = child.kill();
        let _ = child.wait();
    }
    pub fn alive(&self) -> bool {
        self.alive.load(Ordering::Relaxed)
    }
    pub fn busy(&self) -> bool {
        !self.active_turns.lock().unwrap().is_empty()
            || !self.requests.lock().unwrap().is_empty()
            || !self.pending.lock().unwrap().is_empty()
    }
}

impl Drop for Rpc {
    fn drop(&mut self) {
        self.stop();
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn mock_app_server() {
        if std::env::var_os("CODEX_BRIDGE_TEST_SERVER").is_none() {
            return;
        }
        for line in std::io::stdin().lock().lines() {
            let v: Value = serde_json::from_str(&line.unwrap()).unwrap();
            let output = match v["method"].as_str() {
                Some("initialize") => json!({"id":v["id"],"result":{}}),
                Some("test/echo") => json!({"id":v["id"],"result":v["params"]}),
                Some("test/error") => {
                    json!({"id":v["id"],"error":{"code":-1,"message":"expected failure"}})
                }
                Some("test/request") => {
                    println!(
                        "\n{}",
                        json!({"id":"approval-1","method":"item/commandExecution/requestApproval","params":{}})
                    );
                    json!({"id":v["id"],"result":{}})
                }
                None if v["id"] == "approval-1" => {
                    json!({"method":"test/replied","params":v["result"]})
                }
                _ => continue,
            };
            println!("\n{output}");
            std::io::stdout().flush().unwrap();
        }
    }
    #[test]
    fn native_rpc_correlates_concurrent_calls_errors_and_approval_replies() {
        let mut cmd = Command::new(std::env::current_exe().unwrap());
        cmd.args(["--exact", "rpc::tests::mock_app_server", "--nocapture"])
            .env("CODEX_BRIDGE_TEST_SERVER", "1");
        let rpc = Rpc::start_command(cmd, false).unwrap();
        let calls: Vec<_> = (0..8)
            .map(|i| {
                let rpc = rpc.clone();
                thread::spawn(move || {
                    assert_eq!(
                        rpc.call("test/echo", json!({"index":i}), Duration::from_secs(2))
                            .unwrap()["index"],
                        i
                    )
                })
            })
            .collect();
        for call in calls {
            call.join().unwrap();
        }
        assert!(rpc
            .call("test/error", json!({}), Duration::from_secs(2))
            .unwrap_err()
            .contains("expected failure"));
        rpc.call("test/request", json!({}), Duration::from_secs(2))
            .unwrap();
        assert!(rpc.busy());
        assert_eq!(rpc.requests().as_array().unwrap().len(), 1);
        assert!(rpc.reply(json!("wrong"), json!({})).is_err());
        rpc.reply(json!("approval-1"), json!({"decision":"decline"}))
            .unwrap();
        assert!(!rpc.busy());
        let deadline = std::time::Instant::now() + Duration::from_secs(2);
        while rpc.events(0)["latestSeq"].as_u64().unwrap() < 2 {
            assert!(std::time::Instant::now() < deadline);
            thread::sleep(Duration::from_millis(10));
        }
        assert_eq!(rpc.events(1)["events"].as_array().unwrap().len(), 1);
        rpc.stop();
        assert!(!rpc.alive());
        assert!(rpc
            .call("test/echo", json!({}), Duration::from_secs(1))
            .is_err());
    }
}