Pup - the AI's own browser

//! pup-cdp: the Chrome DevTools Protocol client. One websocket per browser, flattened sessions
//! multiplexed over it, a reader thread that routes replies by id and events to subscribers, and
//! typed wrappers for the methods pup uses. This is the queue and lifecycle the Node build never
//! had (28 ad-hoc CDPSessions, no pipeline); every call site goes through `send` with a timeout.
//!
//! Sync by design (like the rest of the bridge): the reader thread owns the socket with a short
//! read timeout so senders can take the lock between reads; a command's reply arrives on a
//! channel the sender waits on. No async runtime.

use serde_json::{json, Value};
use std::collections::HashMap;
use std::net::TcpStream;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tungstenite::protocol::WebSocketConfig;

/// pup-bridge#109: when on, page_enable skips `Runtime.enable` (the CDP leak bot managers fingerprint).
static QUIET_RUNTIME: AtomicBool = AtomicBool::new(false);
use tungstenite::stream::MaybeTlsStream;
use tungstenite::{Message, WebSocket};

#[derive(Debug, thiserror::Error)]
pub enum CdpError {
    #[error("connect failed: {0}")]
    Connect(String),
    #[error("browser connection closed")]
    Closed,
    #[error("{method} timed out after {ms}ms")]
    Timeout { method: String, ms: u64 },
    #[error("{method} failed: {code} {message}")]
    Protocol { method: String, code: i64, message: String },
    #[error("bad reply: {0}")]
    Json(String),
}

#[derive(Debug, Clone)]
pub struct Event { pub method: String, pub params: Value, pub session_id: Option<String> }

struct Sub { session: Option<String>, method: Option<String>, tx: Sender<Event> }

/// An alternate transport: every command goes through this instead of a websocket. pup's extension lane
/// (2.0.578) forwards to the extension's own DevTools driver (chrome.debugger), so verbs written against
/// this struct run unchanged on a browser that has no debug port.
pub type Via = Box<dyn Fn(Option<&str>, &str, Value, Duration) -> Result<Value, CdpError> + Send + Sync>;

struct Inner {
    sock: Option<Mutex<WebSocket<MaybeTlsStream<TcpStream>>>>,
    via: Option<Via>,
    next_id: AtomicU64,
    pending: Mutex<HashMap<u64, Sender<Value>>>,
    subs: Mutex<Vec<Sub>>,
    alive: AtomicBool,
    ws_url: String,
    port: u16,
}

#[derive(Clone)]
pub struct Cdp { inner: Arc<Inner> }

impl Cdp {
    /// GET http://127.0.0.1:<port>/json/version then connect the browser websocket.
    pub fn connect_port(port: u16, timeout: Duration) -> Result<Cdp, CdpError> {
        let v: Value = ureq::AgentBuilder::new().timeout(timeout).build()
            .get(&format!("http://127.0.0.1:{port}/json/version")).call()
            .map_err(|e| CdpError::Connect(format!("/json/version on :{port}: {e}")))?
            .into_json().map_err(|e| CdpError::Connect(e.to_string()))?;
        let ws = v.get("webSocketDebuggerUrl").and_then(|s| s.as_str()).ok_or_else(|| CdpError::Connect("no webSocketDebuggerUrl".into()))?.to_string();
        Cdp::connect(&ws, port)
    }

    pub fn connect(ws_url: &str, port: u16) -> Result<Cdp, CdpError> {
        let mut cfg = WebSocketConfig::default();
        cfg.max_message_size = None;   // screenshots come back as one large text frame
        cfg.max_frame_size = None;
        let (sock, _resp) = tungstenite::client::connect_with_config(ws_url, Some(cfg), 3).map_err(|e| CdpError::Connect(e.to_string()))?;
        if let MaybeTlsStream::Plain(tcp) = sock.get_ref() {
            let _ = tcp.set_read_timeout(Some(Duration::from_millis(25)));
            let _ = tcp.set_nodelay(true);
        }
        let inner = Arc::new(Inner {
            sock: Some(Mutex::new(sock)), via: None, next_id: AtomicU64::new(1), pending: Mutex::new(HashMap::new()), subs: Mutex::new(Vec::new()),
            alive: AtomicBool::new(true), ws_url: ws_url.to_string(), port,
        });
        let reader = inner.clone();
        std::thread::Builder::new().name("cdp-reader".into()).spawn(move || reader_loop(reader)).map_err(|e| CdpError::Connect(e.to_string()))?;
        Ok(Cdp { inner })
    }

    /// A Cdp whose commands travel through `via` (no socket, no reader). Events are not delivered:
    /// `subscribe` hands back a closed channel so waits return at once instead of timing out.
    pub fn via(via: Via) -> Cdp {
        Cdp { inner: Arc::new(Inner { sock: None, via: Some(via), next_id: AtomicU64::new(1), pending: Mutex::new(HashMap::new()), subs: Mutex::new(Vec::new()), alive: AtomicBool::new(true), ws_url: "via".into(), port: 0 }) }
    }
    pub fn is_via(&self) -> bool { self.inner.via.is_some() }
    pub fn is_alive(&self) -> bool { self.inner.alive.load(Ordering::Relaxed) }
    pub fn ws_url(&self) -> &str { &self.inner.ws_url }
    pub fn port(&self) -> u16 { self.inner.port }

    /// Send one command (browser-level when `session` is None) and wait for its reply.
    pub fn send(&self, session: Option<&str>, method: &str, params: Value, timeout: Duration) -> Result<Value, CdpError> {
        if !self.is_alive() { return Err(CdpError::Closed); }
        if let Some(via) = &self.inner.via { return via(session, method, params, timeout); }
        let id = self.inner.next_id.fetch_add(1, Ordering::Relaxed);
        let (tx, rx) = channel::<Value>();
        self.inner.pending.lock().unwrap().insert(id, tx);
        let mut msg = json!({ "id": id, "method": method, "params": params });
        if let Some(s) = session { msg["sessionId"] = json!(s); }
        {
            let Some(sm) = &self.inner.sock else { return Err(CdpError::Closed) };
            let mut sock = sm.lock().unwrap();
            if let Err(e) = sock.send(Message::Text(msg.to_string().into())) {
                self.inner.pending.lock().unwrap().remove(&id);
                self.inner.alive.store(false, Ordering::Relaxed);
                return Err(CdpError::Connect(format!("send {method}: {e}")));
            }
        }
        match rx.recv_timeout(timeout) {
            Ok(v) => {
                if let Some(e) = v.get("error") {
                    return Err(CdpError::Protocol { method: method.into(), code: e.get("code").and_then(|c| c.as_i64()).unwrap_or(0), message: e.get("message").and_then(|m| m.as_str()).unwrap_or("").into() });
                }
                Ok(v.get("result").cloned().unwrap_or(Value::Null))
            }
            Err(_) => {
                self.inner.pending.lock().unwrap().remove(&id);
                if !self.is_alive() { Err(CdpError::Closed) } else { Err(CdpError::Timeout { method: method.into(), ms: timeout.as_millis() as u64 }) }
            }
        }
    }

    /// Subscribe to events (optionally one session and/or one method). Drop the receiver to
    /// unsubscribe; the reader prunes dead subscribers on the next event.
    pub fn subscribe(&self, session: Option<&str>, method: Option<&str>) -> Receiver<Event> {
        let (tx, rx) = channel::<Event>();
        if self.inner.via.is_some() { drop(tx); return rx; }   // no events on a via transport: recv returns Disconnected at once
        self.inner.subs.lock().unwrap().push(Sub { session: session.map(|s| s.to_string()), method: method.map(|m| m.to_string()), tx });
        rx
    }

    pub fn wait_event(rx: &Receiver<Event>, timeout: Duration) -> Option<Event> { rx.recv_timeout(timeout).ok() }

    // ── typed wrappers (the CDP surface pup actually uses) ──────────────────────────────

    pub fn version(&self) -> Result<Value, CdpError> { self.send(None, "Browser.getVersion", json!({}), Duration::from_secs(5)) }

    pub fn targets(&self) -> Result<Vec<Value>, CdpError> {
        let r = self.send(None, "Target.getTargets", json!({}), Duration::from_secs(10))?;
        Ok(r.get("targetInfos").and_then(|t| t.as_array()).cloned().unwrap_or_default())
    }

    /// Flattened attach; returns the sessionId every page-level command needs.
    pub fn attach(&self, target_id: &str) -> Result<String, CdpError> {
        let r = self.send(None, "Target.attachToTarget", json!({ "targetId": target_id, "flatten": true }), Duration::from_secs(10))?;
        r.get("sessionId").and_then(|s| s.as_str()).map(|s| s.to_string()).ok_or_else(|| CdpError::Json("attachToTarget: no sessionId".into()))
    }

    pub fn detach(&self, session_id: &str) { let _ = self.send(None, "Target.detachFromTarget", json!({ "sessionId": session_id }), Duration::from_secs(5)); }

    /// A NEW OS WINDOW (only Target.createTarget accepts newWindow; Node used exactly this).
    /// `focus`: Some(false) creates a new window WITHOUT Chrome asking Windows for the foreground (2.0.449: no
    /// denied activation, so no attention state on the taskbar button); None leaves Chrome's default (the
    /// window asks, and Windows flags the button when it refuses). Tabs always pass None.
    pub fn create_target(&self, url: &str, new_window: bool, bounds: Option<(i32, i32, i32, i32)>, focus: Option<bool>) -> Result<String, CdpError> {
        let mut p = json!({ "url": url, "newWindow": new_window });
        if let Some(f) = focus { p["focus"] = json!(f); }
        if let Some((l, t, w, h)) = bounds { p["left"] = json!(l); p["top"] = json!(t); p["width"] = json!(w); p["height"] = json!(h); }
        let r = self.send(None, "Target.createTarget", p, Duration::from_secs(20))?;
        r.get("targetId").and_then(|s| s.as_str()).map(|s| s.to_string()).ok_or_else(|| CdpError::Json("createTarget: no targetId".into()))
    }

    pub fn close_target(&self, target_id: &str) -> Result<bool, CdpError> {
        let r = self.send(None, "Target.closeTarget", json!({ "targetId": target_id }), Duration::from_secs(10))?;
        Ok(r.get("success").and_then(|b| b.as_bool()).unwrap_or(true))
    }

    pub fn window_for_target(&self, target_id: &str) -> Result<(i64, Value), CdpError> {
        let r = self.send(None, "Browser.getWindowForTarget", json!({ "targetId": target_id }), Duration::from_secs(10))?;
        Ok((r.get("windowId").and_then(|w| w.as_i64()).unwrap_or(0), r.get("bounds").cloned().unwrap_or(Value::Null)))
    }

    pub fn set_window_bounds(&self, window_id: i64, bounds: Value) -> Result<(), CdpError> {
        self.send(None, "Browser.setWindowBounds", json!({ "windowId": window_id, "bounds": bounds }), Duration::from_secs(10)).map(|_| ())
    }

    /// pup-bridge#109: `Runtime.enable` is THE CDP leak bot managers fingerprint in 2026 — the side effect lands
    /// before any page JavaScript runs, so no in-page stealth can mask it. When this is set, pup keeps the Page
    /// domain (navigation, load events, addScriptToEvaluateOnNewDocument) and skips Runtime. Cost: no console /
    /// exception mirroring and no Runtime.bindingCalled, so the credential-capture and floaty bindings go quiet.
    pub fn quiet_runtime(on: bool) { QUIET_RUNTIME.store(on, std::sync::atomic::Ordering::SeqCst); }
    pub fn quiet_runtime_on() -> bool { QUIET_RUNTIME.load(std::sync::atomic::Ordering::SeqCst) }

    pub fn page_enable(&self, session: &str) -> Result<(), CdpError> {
        self.send(Some(session), "Page.enable", json!({}), Duration::from_secs(10)).map(|_| ())?;
        if Self::quiet_runtime_on() { return Ok(()); }
        // 2.0.458: focus emulation is NO LONGER set here. It used to be set on every attached page
        // ("puppeteer sets this on every page"), but a page that reports focus also reports VISIBLE, so
        // Chrome never throttled that tab's timers and Memory Saver never discarded it. pup_verbs::focus_emulate
        // now puts it on just the tabs pup needs it on.
        self.send(Some(session), "Runtime.enable", json!({}), Duration::from_secs(10)).map(|_| ())
    }

    /// Emulation.setFocusEmulationEnabled: make this page report focus even while its window sits behind the
    /// user's windows, so CDP input lands and focus/caret events fire. Costly to leave on (see page_enable).
    pub fn set_focus_emulation(&self, session: &str, enabled: bool) {
        let _ = self.send(Some(session), "Emulation.setFocusEmulationEnabled", json!({ "enabled": enabled }), Duration::from_secs(5));
    }

    /// Navigate and wait for the load event (or the timeout, which is NOT an error: the page is
    /// still moving; callers report `loaded:false`).
    pub fn navigate(&self, session: &str, url: &str, wait: Duration) -> Result<(String, bool), CdpError> {
        let rx = self.subscribe(Some(session), Some("Page.loadEventFired"));
        let r = self.send(Some(session), "Page.navigate", json!({ "url": url }), Duration::from_secs(30))?;
        if let Some(e) = r.get("errorText").and_then(|e| e.as_str()) { if !e.is_empty() { return Err(CdpError::Protocol { method: "Page.navigate".into(), code: 0, message: e.into() }); } }
        let frame = r.get("frameId").and_then(|f| f.as_str()).unwrap_or("").to_string();
        let loaded = Cdp::wait_event(&rx, wait).is_some();
        Ok((frame, loaded))
    }

    /// Runtime.evaluate with the devtools-console semantics the Node build used: expression
    /// first (returnByValue, awaitPromise), the error surfaced as `__error` like pup_eval did.
    pub fn evaluate(&self, session: &str, expression: &str, timeout: Duration) -> Result<Value, CdpError> {
        // Expression form first (awaitPromise resolves an async IIFE); a SyntaxError means the caller
        // wrote STATEMENTS, so retry in replMode (devtools-console semantics), like pup_eval did.
        let mut r = self.send(Some(session), "Runtime.evaluate", json!({ "expression": expression, "returnByValue": true, "awaitPromise": true, "userGesture": true }), timeout)?;
        let is_syntax = r.get("exceptionDetails").map(|ex| ex.get("exception").and_then(|e| e.get("description")).and_then(|d| d.as_str()).unwrap_or("").starts_with("SyntaxError") || ex.get("text").and_then(|t| t.as_str()).unwrap_or("").contains("SyntaxError")).unwrap_or(false);
        if is_syntax { r = self.send(Some(session), "Runtime.evaluate", json!({ "expression": expression, "returnByValue": true, "awaitPromise": true, "userGesture": true, "replMode": true }), timeout)?; }
        if let Some(ex) = r.get("exceptionDetails") {
            let desc = ex.get("exception").and_then(|e| e.get("description")).and_then(|d| d.as_str()).or_else(|| ex.get("text").and_then(|t| t.as_str())).unwrap_or("evaluation threw");
            return Ok(json!({ "__error": desc }));
        }
        Ok(r.get("result").and_then(|x| x.get("value")).cloned().unwrap_or(Value::Null))
    }

    /// Runtime.evaluate inside one execution context (a frame's isolated world), same semantics as `evaluate`.
    pub fn evaluate_in_context(&self, session: &str, context_id: u64, expression: &str, timeout: Duration) -> Result<Value, CdpError> {
        let mut r = self.send(Some(session), "Runtime.evaluate", json!({ "expression": expression, "contextId": context_id, "returnByValue": true, "awaitPromise": true, "userGesture": true }), timeout)?;
        let is_syntax = r.get("exceptionDetails").map(|ex| ex.get("exception").and_then(|e| e.get("description")).and_then(|d| d.as_str()).unwrap_or("").starts_with("SyntaxError")).unwrap_or(false);
        if is_syntax { r = self.send(Some(session), "Runtime.evaluate", json!({ "expression": expression, "contextId": context_id, "returnByValue": true, "awaitPromise": true, "userGesture": true, "replMode": true }), timeout)?; }
        if let Some(ex) = r.get("exceptionDetails") {
            let desc = ex.get("exception").and_then(|e| e.get("description")).and_then(|d| d.as_str()).or_else(|| ex.get("text").and_then(|t| t.as_str())).unwrap_or("evaluation threw");
            return Ok(json!({ "__error": desc }));
        }
        Ok(r.get("result").and_then(|x| x.get("value")).cloned().unwrap_or(Value::Null))
    }
    /// Every frame of the page, depth-first: [{id, name, url, parentId}] (Page.getFrameTree).
    pub fn frames(&self, session: &str) -> Result<Vec<Value>, CdpError> {
        let r = self.send(Some(session), "Page.getFrameTree", json!({}), Duration::from_secs(10))?;
        let mut out = Vec::new();
        fn walk(n: &Value, out: &mut Vec<Value>) { let f = &n["frame"]; out.push(json!({ "id": f["id"], "name": f["name"], "url": f["url"], "parentId": f["parentId"] })); if let Some(c) = n["childFrames"].as_array() { for k in c { walk(k, out); } } }
        walk(&r["frameTree"], &mut out);
        Ok(out)
    }
    /// An execution context for a frame (an isolated world with universal access), for evaluate_in_context.
    pub fn frame_context(&self, session: &str, frame_id: &str) -> Result<u64, CdpError> {
        let r = self.send(Some(session), "Page.createIsolatedWorld", json!({ "frameId": frame_id, "worldName": "pup-eval", "grantUniveralAccess": true }), Duration::from_secs(10))?;
        r.get("executionContextId").and_then(|v| v.as_u64()).ok_or_else(|| CdpError::Json("createIsolatedWorld: no executionContextId".into()))
    }
    /// Target.getTargetInfo with a caller-chosen budget (the list verbs use a short one so a jammed page cannot stall them).
    pub fn target_info_within(&self, session: &str, timeout: Duration) -> Result<Value, CdpError> {
        let r = self.send(Some(session), "Target.getTargetInfo", json!({}), timeout)?;
        Ok(r.get("targetInfo").cloned().unwrap_or(Value::Null))
    }

    /// PNG bytes (base64-decoded) of the viewport or the full page.
    pub fn screenshot_png(&self, session: &str, full_page: bool, timeout: Duration) -> Result<Vec<u8>, CdpError> {
        let mut p = json!({ "format": "png", "captureBeyondViewport": full_page });
        if full_page {
            if let Ok(m) = self.send(Some(session), "Page.getLayoutMetrics", json!({}), Duration::from_secs(10)) {
                let cs = m.get("cssContentSize").or_else(|| m.get("contentSize")).cloned().unwrap_or(Value::Null);
                let (w, h) = (cs.get("width").and_then(|x| x.as_f64()).unwrap_or(0.0), cs.get("height").and_then(|x| x.as_f64()).unwrap_or(0.0));
                if w > 0.0 && h > 0.0 { p["clip"] = json!({ "x": 0, "y": 0, "width": w, "height": h, "scale": 1 }); }
            }
        }
        let r = self.send(Some(session), "Page.captureScreenshot", p, timeout)?;
        let b64 = r.get("data").and_then(|d| d.as_str()).ok_or_else(|| CdpError::Json("captureScreenshot: no data".into()))?;
        base64_decode(b64).map_err(|e| CdpError::Json(format!("screenshot base64: {e}")))
    }

    pub fn target_info(&self, session: &str) -> Result<Value, CdpError> {
        let r = self.send(Some(session), "Target.getTargetInfo", json!({}), Duration::from_secs(5))?;
        Ok(r.get("targetInfo").cloned().unwrap_or(Value::Null))
    }

    pub fn close(&self) {
        self.inner.alive.store(false, Ordering::Relaxed);
        if let Some(sm) = &self.inner.sock { if let Ok(mut s) = sm.lock() { let _ = s.close(None); } }
    }
}

fn reader_loop(inner: Arc<Inner>) {
    let mut idle = 0u32;
    while inner.alive.load(Ordering::Relaxed) {
        let msg = {
            let Some(sm) = &inner.sock else { break };
            let mut sock = match sm.lock() { Ok(s) => s, Err(_) => break };
            match sock.read() {
                Ok(m) => Some(m),
                Err(tungstenite::Error::Io(e)) if matches!(e.kind(), std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut) => None,
                Err(_) => { inner.alive.store(false, Ordering::Relaxed); break; }
            }
        };
        match msg {
            Some(Message::Text(t)) => {
                idle = 0;
                if let Ok(v) = serde_json::from_str::<Value>(&t) {
                    if let Some(id) = v.get("id").and_then(|i| i.as_u64()) {
                        if let Some(tx) = inner.pending.lock().unwrap().remove(&id) { let _ = tx.send(v); }
                    } else if let Some(method) = v.get("method").and_then(|m| m.as_str()) {
                        let ev = Event { method: method.to_string(), params: v.get("params").cloned().unwrap_or(Value::Null), session_id: v.get("sessionId").and_then(|s| s.as_str()).map(|s| s.to_string()) };
                        let mut subs = inner.subs.lock().unwrap();
                        subs.retain(|s| {
                            let want = s.session.as_deref().map_or(true, |x| Some(x) == ev.session_id.as_deref()) && s.method.as_deref().map_or(true, |m| m == ev.method);
                            if !want { return true; }
                            s.tx.send(ev.clone()).is_ok()
                        });
                    }
                }
            }
            Some(Message::Close(_)) => { inner.alive.store(false, Ordering::Relaxed); break; }
            Some(_) => {}
            None => { idle = idle.saturating_add(1); if idle > 4 { std::thread::sleep(Duration::from_millis(3)); } }
        }
    }
    // wake every waiter so nobody blocks on a dead connection
    inner.pending.lock().unwrap().clear();
}

/// Minimal base64 (standard alphabet) decoder, so the crate needs no extra dependency.
pub fn base64_decode(s: &str) -> Result<Vec<u8>, String> {
    fn val(c: u8) -> Result<u32, String> {
        match c { b'A'..=b'Z' => Ok((c - b'A') as u32), b'a'..=b'z' => Ok((c - b'a' + 26) as u32), b'0'..=b'9' => Ok((c - b'0' + 52) as u32), b'+' => Ok(62), b'/' => Ok(63), _ => Err(format!("bad base64 byte {c}")) }
    }
    let bytes: Vec<u8> = s.bytes().filter(|b| !b.is_ascii_whitespace()).collect();
    let mut out = Vec::with_capacity(bytes.len() * 3 / 4);
    let mut chunk = [0u32; 4];
    let mut n = 0;
    let mut pad = 0;
    for &b in &bytes {
        if b == b'=' { pad += 1; chunk[n] = 0; n += 1; } else { chunk[n] = val(b)?; n += 1; }
        if n == 4 {
            let v = (chunk[0] << 18) | (chunk[1] << 12) | (chunk[2] << 6) | chunk[3];
            out.push((v >> 16) as u8);
            if pad < 2 { out.push((v >> 8) as u8); }
            if pad < 1 { out.push(v as u8); }
            n = 0;
        }
    }
    Ok(out)
}

#[allow(dead_code)]
fn _unused(_: Instant) {}

#[cfg(test)]
mod tests {
    use super::*;
    #[test] fn b64() { assert_eq!(base64_decode("aGVsbG8=").unwrap(), b"hello"); assert_eq!(base64_decode("aGk=").unwrap(), b"hi"); assert_eq!(base64_decode("aGVsbG8gd29ybGQ=").unwrap(), b"hello world"); }
}