aura-gui 0.2.2

//! JSON-RPC 2.0 por stdio con un mensaje por línea, común a MCP y ACP: lanza el proceso, reparte las respuestas a
//! quien las espera y entrega al cliente las peticiones y notificaciones que llegan del otro lado.
//! Sin async: un hilo lector por proceso y callbacks.

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

/// Qué se hace con la respuesta a una petición nuestra.
type Callback = Box<dyn FnOnce(Result<Value, AiError>) + Send>;
type Waiting = Arc<Mutex<HashMap<u64, Callback>>>;

/// Lo que el otro extremo nos envía por iniciativa propia.
#[derive(Clone, Debug)]
pub enum Incoming {
    Request { id: Value, method: String, params: Value },
    Notification { method: String, params: Value },
}

/// Qué hizo el manejador con una petición del otro extremo.
pub enum Handled {
    /// Se responderá después con [`RpcPeer::respond`].
    Later,
    /// No se soporta: el lector contesta «método no encontrado».
    Unsupported,
}

pub struct RpcPeer {
    child: Child,
    stdin: Arc<Mutex<ChildStdin>>,
    waiting: Waiting,
    next_id: AtomicU64,
}

fn write_line(stdin: &Mutex<ChildStdin>, message: &Value) -> Result<(), AiError> {
    let mut stdin = stdin.lock().map_err(|_| AiError::Transport("RPC: entrada bloqueada".into()))?;
    writeln!(stdin, "{message}")
        .and_then(|()| stdin.flush())
        .map_err(|error| AiError::Transport(format!("RPC: {error}")))
}

fn error_from(value: &Value) -> AiError {
    AiError::Api(value["message"].as_str().unwrap_or("error del par").to_owned())
}

impl RpcPeer {
    pub fn spawn(
        command: &str,
        args: &[String],
        env: &BTreeMap<String, String>,
        cwd: Option<&Path>,
        mut on_incoming: impl FnMut(Incoming) -> Handled + Send + 'static,
    ) -> Result<Self, AiError> {
        let mut process = Command::new(crate::ai::exec_path::resolve(command));
        process
            .env("PATH", crate::ai::exec_path::child_path())
            .args(args)
            .envs(env)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null());
        if let Some(cwd) = cwd {
            process.current_dir(cwd);
        }
        let mut child =
            process.spawn().map_err(|error| AiError::Transport(format!("no se pudo ejecutar {command}: {error}")))?;
        let (Some(stdin), Some(stdout)) = (child.stdin.take(), child.stdout.take()) else {
            return Err(AiError::Transport("RPC: el proceso no abrió stdio".into()));
        };
        let (stdin, waiting): (_, Waiting) = (Arc::new(Mutex::new(stdin)), Arc::default());
        let (reply_stdin, reply_waiting) = (stdin.clone(), waiting.clone());
        let reader = move || {
            for line in BufReader::new(stdout).lines().map_while(Result::ok) {
                let Ok(message) = serde_json::from_str::<Value>(&line) else { continue };
                match (message["method"].as_str(), message.get("id")) {
                    (Some(method), Some(id)) => {
                        let request = Incoming::Request {
                            id: id.clone(),
                            method: method.to_owned(),
                            params: message["params"].clone(),
                        };
                        if matches!(on_incoming(request), Handled::Unsupported) {
                            let reply = json!({ "jsonrpc": "2.0", "id": id, "error": { "code": -32601, "message": "no soportado" } });
                            if let Err(error) = write_line(&reply_stdin, &reply) {
                                eprintln!("{error}");
                            }
                        }
                    }
                    (Some(method), None) => {
                        on_incoming(Incoming::Notification {
                            method: method.to_owned(),
                            params: message["params"].clone(),
                        });
                    }
                    (None, _) => {
                        let callback = message["id"].as_u64().and_then(|id| reply_waiting.lock().ok()?.remove(&id));
                        if let Some(callback) = callback {
                            callback(match message.get("error") {
                                Some(error) => Err(error_from(error)),
                                None => Ok(message["result"].clone()),
                            });
                        }
                    }
                }
            }
        };
        thread::Builder::new()
            .name("aura-rpc-reader".into())
            .spawn(reader)
            .map_err(|error| AiError::Transport(format!("no se pudo crear el lector: {error}")))?;
        Ok(Self { child, stdin, waiting, next_id: AtomicU64::new(1) })
    }

    pub fn notify(&self, method: &str, params: Value) -> Result<(), AiError> {
        write_line(&self.stdin, &json!({ "jsonrpc": "2.0", "method": method, "params": params }))
    }

    /// Envía una petición y entrega la respuesta a `callback` desde el hilo lector.
    pub fn request_async(
        &self,
        method: &str,
        params: Value,
        callback: impl FnOnce(Result<Value, AiError>) + Send + 'static,
    ) -> Result<(), AiError> {
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        self.waiting
            .lock()
            .map_err(|_| AiError::Transport("RPC: estado bloqueado".into()))?
            .insert(id, Box::new(callback));
        let sent = write_line(&self.stdin, &json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }));
        if sent.is_err()
            && let Ok(mut waiting) = self.waiting.lock()
        {
            waiting.remove(&id);
        }
        sent
    }

    /// Petición bloqueante con tiempo máximo.
    pub fn request(&self, method: &str, params: Value, timeout: Duration) -> Result<Value, AiError> {
        let (sender, receiver) = mpsc::channel();
        self.request_async(method, params, move |result| {
            // Si quien esperaba ya se rindió, la respuesta tardía se descarta.
            sender.send(result).ok();
        })?;
        receiver.recv_timeout(timeout).map_err(|_| AiError::Transport(format!("{method} no respondió")))?
    }

    /// Contesta una petición del otro extremo.
    pub fn respond(&self, id: &Value, result: Result<Value, String>) -> Result<(), AiError> {
        let message = match result {
            Ok(result) => json!({ "jsonrpc": "2.0", "id": id, "result": result }),
            Err(message) => json!({ "jsonrpc": "2.0", "id": id, "error": { "code": -32000, "message": message } }),
        };
        write_line(&self.stdin, &message)
    }
}

impl Drop for RpcPeer {
    fn drop(&mut self) {
        if let Err(error) = self.child.kill()
            && error.kind() != std::io::ErrorKind::InvalidInput
        {
            eprintln!("no se pudo detener el proceso: {error}");
        }
        if let Err(error) = self.child.wait() {
            eprintln!("no se pudo esperar al proceso: {error}");
        }
    }
}