xagent-pi 0.2.3

//! High-level hub client: create session, connect Socket.IO /cli, keep-alive,
//! versioned update-metadata, send agent messages, RPC ack plumbing.

use crate::config::Config;
use crate::hub::rest;
use crate::hub::sio::{self, Inbound, SioClient};
use anyhow::{bail, Result};
use serde_json::{json, Value};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use tokio::sync::Mutex;

const PROTOCOL_VERSION: &str = "1";
const APP_VERSION: &str = env!("CARGO_PKG_VERSION");

pub struct HubClient {
    pub sio: Arc<SioClient>,
    pub session_id: String,
    pub machine_id: String,
    pub session_name: String,
    metadata_version: Arc<AtomicU64>,
    thinking: Arc<AtomicBool>,
    events_rx: Mutex<tokio::sync::mpsc::UnboundedReceiver<Inbound>>,
}

fn now_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_millis() as u64)
        .unwrap_or(0)
}

fn hostname() -> String {
    std::env::var("HOSTNAME")
        .or_else(|_| std::env::var("HOST"))
        .unwrap_or_else(|_| "pi-rs".to_string())
}

fn ws_url_from(hub_url: &str) -> String {
    let base = hub_url.trim_end_matches('/');
    let ws_base = base
        .strip_prefix("http://")
        .map(|s| format!("ws://{s}"))
        .or_else(|| base.strip_prefix("https://").map(|s| format!("wss://{s}")))
        .unwrap_or_else(|| format!("ws://{base}"));
    format!("{ws_base}/socket.io/?EIO=4&transport=websocket")
}

impl HubClient {
    pub async fn start(cfg: &Config) -> Result<Self> {
        // 1. REST create session.
        let created = rest::create_session(
            &cfg.hub_url,
            &cfg.token,
            &cfg.machine_id,
            cfg.session_id.as_deref(),
            &cfg.session_name,
            &cfg.cwd.to_string_lossy(),
            &hostname(),
            concat!("xagent-pi ", env!("CARGO_PKG_VERSION")),
        )
        .await?;
        log::info!(
            "created/loaded session id={} metadataVersion={}",
            created.id,
            created.metadata_version
        );

        let metadata_version = Arc::new(AtomicU64::new(created.metadata_version));
        let thinking = Arc::new(AtomicBool::new(false));

        // 2. Connect Socket.IO /cli.
        let ws_url = ws_url_from(&cfg.hub_url);
        let auth = json!({
            "token": cfg.token,
            "clientType": "session-scoped",
            "sessionId": created.id,
        });
        let (sio, events_rx, connect_rx) = sio::connect(&ws_url, auth).await?;
        let sio = Arc::new(sio);

        // 3. Await CONNECT confirmation.
        match tokio::time::timeout(Duration::from_secs(15), connect_rx).await {
            Ok(Ok(Ok(()))) => log::info!("connected to /cli namespace"),
            Ok(Ok(Err(e))) => bail!("socket.io connect rejected: {e}"),
            Ok(Err(_)) => bail!("socket.io connect dropped"),
            Err(_) => bail!("socket.io connect timeout"),
        }

        // 4. Immediate keep-alive + 2s heartbeat loop.
        let sio_for_keep = Arc::clone(&sio);
        let sid_for_keep = created.id.clone();
        let thinking_for_keep = Arc::clone(&thinking);
        let _ = sio_for_keep
            .emit(
                "session-alive",
                json!({ "sid": sid_for_keep, "time": now_ms(), "thinking": false, "mode": "remote" }),
            )
            .await;
        tokio::spawn(async move {
            let mut tick = tokio::time::interval(Duration::from_secs(2));
            tick.tick().await; // skip immediate (already sent)
            loop {
                tick.tick().await;
                let payload = json!({
                    "sid": sid_for_keep,
                    "time": now_ms(),
                    "thinking": thinking_for_keep.load(Ordering::Relaxed),
                    "mode": "remote",
                });
                if let Err(e) = sio_for_keep.emit("session-alive", payload).await {
                    log::warn!("keepalive emit failed: {e}");
                    break;
                }
            }
        });

        // 5. Register RPC methods we handle.
        for method in ["abort", "set-session-config", "listPiModels"] {
            let _ = sio
                .emit(
                    "rpc-register",
                    json!({ "method": format!("{}:{}", created.id, method) }),
                )
                .await;
        }

        Ok(Self {
            sio,
            session_id: created.id,
            machine_id: cfg.machine_id.clone(),
            session_name: cfg.session_name.clone(),
            metadata_version,
            thinking,
            events_rx: Mutex::new(events_rx),
        })
    }

    pub fn set_thinking(&self, on: bool) {
        self.thinking.store(on, Ordering::Relaxed);
    }

    pub async fn send_session_ready(&self) -> Result<()> {
        self.sio
            .emit(
                "session-ready",
                json!({ "sid": self.session_id, "time": now_ms() }),
            )
            .await
    }

    /// Send an agent message. `data` is the CodexMessage content
    /// (e.g. {type:"message", message:"..."}).
    pub async fn send_agent_message(&self, data: Value) -> Result<()> {
        let envelope = json!({
            "sid": self.session_id,
            "message": {
                "role": "agent",
                "content": { "type": "codex", "data": data },
                "meta": { "sentFrom": "cli" }
            }
        });
        self.sio.emit("message", envelope).await
    }

    /// Versioned update-metadata with mismatch retry.
    pub async fn update_metadata(&self, metadata: Value) -> Result<()> {
        for _ in 0..10 {
            let expected = self.metadata_version.load(Ordering::Relaxed);
            let payload = json!({
                "sid": self.session_id,
                "expectedVersion": expected,
                "metadata": metadata,
            });
            let ack = self
                .sio
                .emit_with_ack("update-metadata", payload, Duration::from_secs(15))
                .await?;
            let result = ack.get("result").and_then(|v| v.as_str()).unwrap_or("");
            match result {
                "success" => {
                    if let Some(v) = ack.get("version").and_then(|v| v.as_u64()) {
                        self.metadata_version.store(v, Ordering::Relaxed);
                    }
                    return Ok(());
                }
                "version-mismatch" => {
                    if let Some(v) = ack.get("version").and_then(|v| v.as_u64()) {
                        self.metadata_version.store(v, Ordering::Relaxed);
                    }
                    tokio::time::sleep(Duration::from_millis(250)).await;
                    continue;
                }
                other => bail!("update-metadata error: {other} ({ack})"),
            }
        }
        bail!("update-metadata exhausted retries")
    }

    pub async fn send_rpc_ack(&self, ack_id: u64, value: Value) -> Result<()> {
        self.sio.send_ack(ack_id, value).await
    }

    pub fn protocol_version() -> &'static str {
        PROTOCOL_VERSION
    }

    pub fn app_version() -> &'static str {
        APP_VERSION
    }

    pub async fn next_event(&self) -> Option<Inbound> {
        self.events_rx.lock().await.recv().await
    }
}