Pup - the AI's own browser

//! Annotate and snip: the in-page overlay (src/extension-annotate + the embedded overlay script),
//! the capture watcher, tight-crop screenshot, caption bar (drawn by the page's own canvas), the
//! iTXt sidecar, the clipboard (ab desktop_clipboard_set, adom-bridge#145) and the on-page
//! confirmation toast. Ports startAnnotation / processAnnotateCapture / addCaptionBar / pngAddITXt.

use crate::{touch, Ctx};
use pup_core::{err, ok, Caller};
use serde_json::{json, Value};
use std::sync::Arc;
use std::time::{Duration, Instant};

pub const OVERLAY_JS: &str = include_str!("../../../src/extension/annotate-overlay.js");   // ONE copy: shipped in the extension (file-injected on the ext lane), evaluated on CDP
pub const FLOATY_JS: &str = include_str!("../assets/floaty.js");
const T: Duration = Duration::from_secs(10);

pub fn floaty_script(ctx: &Ctx) -> String {
    let px = ctx.settings.get("floatyHitPx", "24");
    FLOATY_JS.replace("${BRIDGE_VERSION}", pup_core::version()).replace("__HITPX__", &px).replace("__FLOATY_ON__", if ctx.settings.get("tbFloatyPencil", "off") == "on" { "1" } else { "" }).replace("__HOTKEY_ON__", if ctx.settings.get("tbSnipHotkey", "on") != "off" { "1" } else { "" })
}

/// pup_annotate {tool: pen|snip, mode: auto|multi}: arm the overlay and watch for the capture.
pub fn start(ctx: &Ctx, caller: &Caller, args: &Value) -> Value {
    // 2.0.702: say why in the log. A toolbar click that resolved no window used to leave no trace at all.
    let s = match crate::resolve(ctx, caller, args) { Ok(s) => s, Err(e) => { pup_core::log(format!("[annotate] could not start (by {}): {} (tab {:?}, url {:?})", caller.thread, e["error"].as_str().unwrap_or("?"), args.get("_tabId"), args.get("_tabUrl").and_then(|u| u.as_str()).map(|u| u.chars().take(80).collect::<String>()))); return e } };
    let Some(t) = s.active().cloned() else { return err("no_tab", "the window has no active tab", None) };
    let lane = match crate::transport::lane_for(ctx, &s) { Ok(l) => l, Err(e) => return e };
    let tool = args["tool"].as_str().unwrap_or("pen").to_string();
    let mode = args["mode"].as_str().filter(|m| *m == "multi" || *m == "auto").map(|m| m.to_string());
    // 2.0.569: both lanes. The overlay is one file in the extension folder: file-injected on the extension
    // lane (a page CSP cannot refuse that, unlike eval; Google Docs and the wiki used to fail here), and
    // evaluated over CDP on the CDP lane.
    if let Err(e) = crate::transport::inject_script(&lane, &t, "annotate-overlay.js", OVERLAY_JS, "window.__pupStartAnnotate") { pup_core::log(format!("[annotate] \"{}\" could not inject the overlay: {}", s.id, e["error"].as_str().unwrap_or("?"))); return pup_core::merge(e, json!({ "errorCode": "annotate_inject_failed" })); }
    let mode_arg = mode.as_ref().map(|m| json!(m)).unwrap_or(Value::Null);
    if let Err(e) = crate::transport::call_window_fn(&lane, &t, "__pupStartAnnotate", &[json!(tool), mode_arg]) { pup_core::log(format!("[annotate] \"{}\" the overlay failed to start in the page: {}", s.id, e["error"].as_str().unwrap_or("?"))); return pup_core::merge(e, json!({ "errorCode": "annotate_start_failed" })); }
    pup_core::log(format!("[annotate] \"{}\" overlay started (by {}, tool={tool}, mode={}, lane={})", s.id, caller.thread, mode.as_deref().unwrap_or("default"), if crate::transport::is_ext(&s) { "ext" } else { "cdp" }));
    if let Some(a) = ctx.arc() { watch(a, s.id.clone(), caller.clone()); }
    touch(ctx, &s.id, None);
    // Node v2.0.293: a HUMAN caller (jump-list task, toolbar pen) clicked Annotate, so the window
    // comes to their screen now, captioned; an AI caller arms the overlay in the BACKGROUND and must
    // go through pup_raise_os_window + foregroundReason.
    let user = crate::USER_CALLERS.contains(&caller.thread.as_str());
    let raised = if user {
        match &lane {
            crate::transport::Lane::Cdp(b) => { let _ = b.cdp.send(Some(&t.cdp_session), "Page.bringToFront", json!({}), T); }
            crate::transport::Lane::Ext => { let _ = crate::transport::activate(&lane, &t); }
        }
        crate::window::raise_for_user(ctx, &s, caller, &format!("Annotate: draw on the page, then hit Done ({})", if tool == "snip" { "drag a region" } else { "pen" }))
    } else { false };
    ok(json!({ "success": true, "output": json!({ "ok": true, "started": true, "raised": raised, "sessionId": s.id, "tool": tool, "mode": mode, "armed": true }).to_string(), "sessionId": s.id, "window": s.id, "aiThread": s.owner,
        "_hint": if user { "Overlay is live and the window is on the user's screen: pen/arrow/box/text, Done captures (the drawing stays visible so the owning AI thread's next screenshot shows it), the X cancels." } else { "Overlay armed in the BACKGROUND. To let the user draw, pup_raise_os_window with a foregroundReason; Done captures to the clipboard and the owning thread's last_annotation." } }))
}

/// Poll window.__pupAnnotateResult (500 ms, up to 15 min) and process each capture.
pub fn watch(ctx: Arc<Ctx>, sid: String, caller: Caller) {
    std::thread::Builder::new().name("annotate-watch".into()).spawn(move || {
        let t0 = Instant::now();
        while t0.elapsed() < Duration::from_secs(15 * 60) {
            std::thread::sleep(Duration::from_millis(500));
            let Some(s) = ctx.sessions.get(&sid) else { return };
            let Some(t) = s.active().cloned() else { return };
            let Ok(lane) = crate::transport::lane_for(&ctx, &s) else { return };
            let Ok((res, dpr)) = crate::transport::read_window(&lane, &t, "__pupAnnotateResult", true) else { return };
            if res.is_null() { continue; }
            // 2.0.703: Esc / the X leave {done:false}. That is a CANCEL: capture nothing, leave the clipboard alone,
            // stop watching. Every result used to be processed as a capture, so cancelling copied a screenshot of the
            // page over whatever the user had on the clipboard (John's laptop, 2026-10-07 12:09).
            if res["done"] == json!(false) { pup_core::log(format!("[annotate] \"{sid}\" cancelled; nothing captured")); return; }
            if let Err(e) = process(&ctx, &sid, &lane, &t, dpr, &res, &caller) { pup_core::log(format!("[annotate] \"{sid}\" capture failed: {e}")); }
        }
    }).ok();
}

fn process(ctx: &Ctx, sid: &str, lane: &crate::transport::Lane, t: &pup_session::Tab, dpr: f64, res: &Value, caller: &Caller) -> Result<(), String> {
    let s = ctx.sessions.get(sid).ok_or("session gone")?;
    let is_snip = res["snip"].is_object();
    let clip = if is_snip { res["snip"].clone() } else { res["bbox"].clone() };
    pup_core::log(format!("[annotate] \"{sid}\" capture: screenshotting ({})", if is_snip { "snip" } else { "tight-crop" }));
    let clip_css = if clip.is_object() { Some((clip["x"].as_f64().unwrap_or(0.0), clip["y"].as_f64().unwrap_or(0.0), clip["w"].as_f64().unwrap_or(1.0), clip["h"].as_f64().unwrap_or(1.0))) } else { None };
    let shot = crate::transport::capture_viewport_png(lane, t, clip_css, dpr).map_err(|e| e["error"].as_str().unwrap_or("capture failed").to_string())?;
    let owner = s.owner.clone().unwrap_or_else(|| s.id.clone());
    let url = s.active().map(|t| t.url.clone()).unwrap_or_default();
    let short_url = url.chars().take(80).collect::<String>();
    let when = pup_core::iso_now();
    // caption bar + preview thumbnail, composed in the binary (2.0.569; it was an in-page canvas, which a page CSP refused)
    let caption = format!("pup · {owner}'s window · {short_url} · {}", &when[11..16]);
    let cap = if ctx.settings.get("annCaptionBar", "on") != "off" { Some(caption.as_str()) } else { None };
    let (mut png, preview_rgba, pw, ph) = pup_identity::badge::annotate_compose(&shot, cap).ok_or("could not compose the capture")?;
    let preview = json!({ "b64": pup_identity::base64_encode(&preview_rgba), "w": pw, "h": ph });   // raw RGBA: the overlay paints it, never loads it
    let sidecar = json!({ "kind": "pup-annotation", "thread": owner, "container": s.owner_container, "sessionId": sid, "url": url, "ts": when, "shapes": res["shapes"], "snip": if is_snip { res["snip"].clone() } else { Value::Null }, "bbox": res["bbox"], "viewport": res["viewport"] });
    png = png_add_itxt(&png, "pup-annotation", &sidecar.to_string());
    let dir = ctx.shots_dir();
    let fp = dir.join(format!("annotate-{}{}-{}.png", if is_snip { "snip-" } else { "" }, pup_session::sanitize_sid(sid), pup_core::now_ms()));
    std::fs::write(&fp, &png).map_err(|e| e.to_string())?;
    ctx.sessions.with_mut(sid, |x| x.last_annotation = Some(json!({ "path": fp.to_string_lossy(), "ts": pup_core::now_ms(), "shapes": res["shapes"], "snip": sidecar["snip"] })));
    let sidecar_text = format!("pup snippet — {owner}'s window — {url} — {when}\n{}", serde_json::to_string_pretty(&sidecar).unwrap_or_default());
    let clipped = match ctx.ab.call("desktop_clipboard_set", json!({ "image": fp.to_string_lossy().replace('\\', "/"), "text": sidecar_text }), caller, Duration::from_secs(8)) { Ok(_) => true, Err(e) => { pup_core::log(format!("[annotate] \"{sid}\" clipboard copy failed: {e}")); false } };
    let line1 = format!("{} copied to your clipboard", if is_snip { "Snippet" } else { "Annotation" });
    let _ = crate::transport::call_window_fn(lane, t, "__pupCaptureDone", &[if clip.is_object() { clip.clone() } else { Value::Null }, preview, json!(line1)]);
    pup_core::log(format!("[annotate] \"{sid}\" CAPTURED {} -> {} | clipboard {}", if is_snip { format!("SNIP {}x{}", res["snip"]["w"], res["snip"]["h"]) } else if clip.is_object() { format!("tight-crop {}x{}", clip["w"], clip["h"]) } else { "viewport".to_string() }, fp.display(), if clipped { "yes" } else { "no" }));
    Ok(())
}

/// Insert an iTXt chunk (keyword, uncompressed UTF-8 text) before IEND.
pub fn png_add_itxt(png: &[u8], keyword: &str, text: &str) -> Vec<u8> {
    let Some(iend) = png.windows(4).rposition(|w| w == b"IEND") else { return png.to_vec() };
    let start = iend - 4;
    let mut data = Vec::new();
    data.extend_from_slice(keyword.as_bytes()); data.push(0); data.push(0); data.push(0); data.push(0); data.push(0); data.extend_from_slice(text.as_bytes());
    let mut chunk = Vec::new();
    chunk.extend_from_slice(&(data.len() as u32).to_be_bytes()); chunk.extend_from_slice(b"iTXt"); chunk.extend_from_slice(&data);
    let mut crc_input = b"iTXt".to_vec(); crc_input.extend_from_slice(&data);
    chunk.extend_from_slice(&crc32(&crc_input).to_be_bytes());
    let mut out = png[..start].to_vec(); out.extend_from_slice(&chunk); out.extend_from_slice(&png[start..]); out
}

fn crc32(data: &[u8]) -> u32 {
    let mut table = [0u32; 256];
    for i in 0..256u32 { let mut c = i; for _ in 0..8 { c = if c & 1 != 0 { 0xEDB88320 ^ (c >> 1) } else { c >> 1 }; } table[i as usize] = c; }
    let mut crc = 0xFFFF_FFFFu32; for &b in data { crc = table[((crc ^ b as u32) & 0xFF) as usize] ^ (crc >> 8); } crc ^ 0xFFFF_FFFF
}

/// pup_floaty_click / the floaty binding: arm the overlay from the corner pencil.
pub fn floaty_click(ctx: &Ctx, caller: &Caller, args: &Value) -> Value { let mut a = args.clone(); a["tool"] = json!("pen"); start(ctx, caller, &a) }

#[cfg(test)]
mod tests {
    use super::*;
    // 2.0.704: Trusted Types pages (Gmail, Google Docs) throw on any HTML-string assignment and the overlay dies with
    // it. Both in-page scripts must build their DOM from nodes; this fails the build if either one regresses.
    #[test] fn no_html_string_sinks() {
        for (name, src) in [("annotate-overlay.js", OVERLAY_JS), ("floaty.js", FLOATY_JS)] {
            for sink in [".innerHTML =", ".innerHTML=", ".outerHTML", "insertAdjacentHTML", "document.write("] {
                assert!(!src.contains(sink), "{name} uses {sink}: Trusted Types pages (Gmail, Google Docs) refuse it and the overlay dies");
            }
        }
    }
    #[test] fn itxt() { assert_eq!(crc32(b"123456789"), 0xCBF43926); let png = b"\x89PNG\r\n\x1a\n\0\0\0\0IEND\xaeB`\x82".to_vec(); let out = png_add_itxt(&png, "k", "v"); assert!(out.windows(4).any(|w| w == b"iTXt")); assert!(out.ends_with(b"IEND\xaeB`\x82")); }
}