xagent-pi 0.2.10

//! Pi process transport: spawn `pi --mode rpc`, JSON-line RPC over stdin/stdout,
//! with id-correlated `response` matching for awaitable commands.

use anyhow::{anyhow, Result};
use serde_json::Value;
use std::collections::HashMap;
use std::path::Path;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};
use tokio::process::{Child, Command, ChildStdin, ChildStdout};
use tokio::sync::{mpsc, oneshot, Mutex};

pub struct PiTransport {
    stdin_tx: mpsc::UnboundedSender<String>,
    pub events_rx: Mutex<mpsc::UnboundedReceiver<Value>>,
    pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>>,
    next_id: Arc<AtomicU64>,
    _child: Child,
}

impl PiTransport {
    pub fn spawn(pi_path: &str, cwd: &Path, resume_session: Option<&str>) -> Result<Self> {
        let mut cmd = Command::new(pi_path);
        cmd.args(["--mode", "rpc"]);
        if let Some(id) = resume_session {
            cmd.args(["--session", id]);
        }
        cmd.current_dir(cwd)
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::piped());

        let mut child = cmd.spawn()?;
        let stdin: ChildStdin = child
            .stdin
            .take()
            .ok_or_else(|| anyhow!("no pi stdin"))?;
        let stdout: ChildStdout = child
            .stdout
            .take()
            .ok_or_else(|| anyhow!("no pi stdout"))?;
        let stderr = child.stderr.take();

        let pending: Arc<Mutex<HashMap<u64, oneshot::Sender<Value>>>> =
            Arc::new(Mutex::new(HashMap::new()));
        let next_id = Arc::new(AtomicU64::new(0));

        // stdin writer
        let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel::<String>();
        tokio::spawn(async move {
            let mut w = BufWriter::new(stdin);
            while let Some(line) = stdin_rx.recv().await {
                if w.write_all(line.as_bytes()).await.is_err() {
                    break;
                }
                if w.write_all(b"\n").await.is_err() {
                    break;
                }
                let _ = w.flush().await;
            }
        });

        // stdout reader: route `response` events with id to pending; others to events channel.
        let (events_tx, events_rx) = mpsc::unbounded_channel::<Value>();
        let pending_clone = Arc::clone(&pending);
        tokio::spawn(async move {
            let mut reader = BufReader::new(stdout).lines();
            loop {
                match reader.next_line().await {
                    Ok(Some(line)) => {
                        let line = line.trim();
                        if line.is_empty() {
                            continue;
                        }
                        let v: Value = match serde_json::from_str(line) {
                            Ok(v) => v,
                            Err(e) => {
                                log::debug!("pi: skipping non-json line: {e}: {}", short(line));
                                continue;
                            }
                        };
                        if v.get("type").and_then(|t| t.as_str()) == Some("response") {
                            if let Some(id_str) = v.get("id").and_then(|i| i.as_str()) {
                                if let Ok(id) = id_str.parse::<u64>() {
                                    let mut g = pending_clone.lock().await;
                                    if let Some(tx) = g.remove(&id) {
                                        let _ = tx.send(v);
                                        continue;
                                    }
                                }
                            }
                            // response without a known id → still emit (startup get_state etc.)
                        }
                        let _ = events_tx.send(v);
                    }
                    Ok(None) => {
                        log::info!("pi stdout ended");
                        break;
                    }
                    Err(e) => {
                        log::warn!("pi stdout read error: {e}");
                        break;
                    }
                }
            }
        });

        // stderr → debug log
        if let Some(stderr) = stderr {
            tokio::spawn(async move {
                let mut reader = BufReader::new(stderr).lines();
                while let Ok(Some(line)) = reader.next_line().await {
                    log::debug!("[pi:stderr] {}", line.trim_end());
                }
            });
        }

        Ok(Self {
            stdin_tx,
            events_rx: Mutex::new(events_rx),
            pending,
            next_id,
            _child: child,
        })
    }

    /// Fire-and-forget command (no id).
    pub fn send(&self, cmd: Value) -> Result<()> {
        self.stdin_tx
            .send(serde_json::to_string(&cmd)?)
            .map_err(|_| anyhow!("pi stdin closed"))
    }

    /// Awaitable command: injects `id`, returns the full `response` event Value.
    pub async fn send_and_wait(&self, mut cmd: Value, timeout: Duration) -> Result<Value> {
        let id = self.next_id.fetch_add(1, Ordering::SeqCst) + 1;
        if let Some(obj) = cmd.as_object_mut() {
            obj.insert("id".to_string(), Value::String(id.to_string()));
        }
        let (tx, rx) = oneshot::channel();
        self.pending.lock().await.insert(id, tx);
        self.send(cmd)?;
        match tokio::time::timeout(timeout, rx).await {
            Ok(Ok(v)) => Ok(v),
            _ => {
                self.pending.lock().await.remove(&id);
                Err(anyhow!("pi rpc timeout (id={id})"))
            }
        }
    }

}

fn short(s: &str) -> String {
    if s.len() > 200 {
        format!("{}…", &s[..200])
    } else {
        s.to_string()
    }
}