eoka-cli 0.10.3

//! State save/load and console capture.

use std::collections::HashMap;

use eoka::Page;
use eoka_server::saved_state::StorageRestoreReport;
use serde::{Deserialize, Serialize};
use serde_json::json;

use super::state::TabState;

// ---------------------------------------------------------------------------
// Saved state types
// ---------------------------------------------------------------------------

#[derive(Debug, Serialize, Deserialize)]
pub struct SavedState {
    pub url: String,
    pub cookies: Vec<SavedCookie>,
    pub local_storage: HashMap<String, String>,
    pub session_storage: HashMap<String, String>,
    pub user_agent: String,
    #[serde(default)]
    pub saved_at: String,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct SavedCookie {
    pub name: String,
    pub value: String,
    pub domain: String,
    pub path: String,
    pub expires: f64,
    pub http_only: bool,
    pub secure: bool,
    pub same_site: Option<String>,
}

impl From<eoka::SessionCookie> for SavedCookie {
    fn from(c: eoka::SessionCookie) -> Self {
        Self {
            name: c.name,
            value: c.value,
            domain: c.domain,
            path: c.path,
            expires: c.expires.unwrap_or(0.0),
            http_only: c.http_only,
            secure: c.secure,
            same_site: c.same_site,
        }
    }
}

impl SavedCookie {
    fn to_session_cookie(&self) -> eoka::SessionCookie {
        eoka::SessionCookie {
            name: self.name.clone(),
            value: self.value.clone(),
            domain: self.domain.clone(),
            path: self.path.clone(),
            secure: self.secure,
            http_only: self.http_only,
            same_site: self.same_site.clone(),
            expires: if self.expires > 0.0 {
                Some(self.expires)
            } else {
                None
            },
        }
    }
}

// ---------------------------------------------------------------------------
// Capture / Restore
// ---------------------------------------------------------------------------

/// Capture full browser state: cookies (via CDP), localStorage, sessionStorage, URL, UA.
pub async fn capture_state(page: &Page) -> Result<SavedState, String> {
    let cookies: Vec<SavedCookie> = page
        .cookies()
        .await
        .map_err(|e| e.to_string())?
        .into_iter()
        .map(SavedCookie::from)
        .collect();

    let local_storage: HashMap<String, String> = page
        .evaluate_sync(
            "(() => { try { return Object.fromEntries(Object.entries(localStorage)) } catch(e) { return {} } })()",
        )
        .await
        .map_err(|e| e.to_string())?;

    let session_storage: HashMap<String, String> = page
        .evaluate_sync(
            "(() => { try { return Object.fromEntries(Object.entries(sessionStorage)) } catch(e) { return {} } })()",
        )
        .await
        .map_err(|e| e.to_string())?;

    let url: String = page
        .evaluate_sync("location.href")
        .await
        .map_err(|e| e.to_string())?;

    let user_agent: String = page
        .evaluate_sync("navigator.userAgent")
        .await
        .map_err(|e| e.to_string())?;

    let saved_at = {
        let d = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default();
        format!("{}", d.as_secs())
    };

    Ok(SavedState {
        url,
        cookies,
        local_storage,
        session_storage,
        user_agent,
        saved_at,
    })
}

/// Restore browser state: replace the covered sites' cookies, then clear + set
/// localStorage/sessionStorage.
pub async fn restore_state(
    page: &Page,
    state: &SavedState,
) -> Result<StorageRestoreReport, String> {
    restore_cookies(page, state).await?;
    eoka_server::saved_state::restore_storage(
        page,
        &state.url,
        &state.local_storage,
        &state.session_storage,
    )
    .await
    .map_err(|e| e.to_string())
}

/// Replace the cookies of the sites this state covers, including their parent domains
/// and subdomains, leaving unrelated sites' alone.
pub async fn restore_cookies(page: &Page, state: &SavedState) -> Result<(), String> {
    eoka_server::saved_state::clear_covered_cookies(
        page,
        &state.url,
        state.cookies.iter().map(|c| c.domain.as_str()),
    )
    .await
    .map_err(|e| e.to_string())?;
    let set_cookies: Vec<eoka::SessionCookie> = state
        .cookies
        .iter()
        .map(|c| c.to_session_cookie())
        .collect();
    if !set_cookies.is_empty() {
        page.set_cookies_bulk(set_cookies)
            .await
            .map_err(|e| e.to_string())?;
    }
    Ok(())
}

/// Connect to a running Chrome (port number or ws:// URL), capture its
/// front-tab state via CDP, and disconnect. The `Browser::disconnect()` call
/// leaves the user's Chrome running.
pub async fn clone_state_from_source(source: &str) -> Result<SavedState, String> {
    use crate::launch_spec::resolve_cdp_spec;

    let ws_url = resolve_cdp_spec(source)?;
    let browser = eoka::Browser::connect(&ws_url)
        .await
        .map_err(|e| format!("connect: {}", e))?;

    let tabs = browser.tabs().await.map_err(|e| e.to_string())?;
    // Skip DevTools/extension targets — we want a real user page.
    let target = tabs
        .iter()
        .find(|t| !t.url.starts_with("devtools://") && !t.url.starts_with("chrome-extension://"))
        .ok_or_else(|| "No user-page tab found in the running Chrome".to_string())?;

    let page = browser
        .attach_page(&target.id)
        .await
        .map_err(|e| e.to_string())?;

    let state = capture_state(&page).await?;
    if let Err(e) = browser.disconnect().await {
        eprintln!(
            "[eoka] warning: failed to disconnect from cloned Chrome session: {}",
            e
        );
    }
    Ok(state)
}

// ---------------------------------------------------------------------------
// Console capture
// ---------------------------------------------------------------------------

const CONSOLE_CAPTURE_JS: &str = include_str!("../js/console_capture.js");

/// Inject console capture into the current page and register for future navigations.
pub async fn ensure_console_capture(tab: &mut TabState) -> Result<(), String> {
    if tab.console_injected {
        return Ok(());
    }
    tab.page
        .session()
        .send::<_, serde_json::Value>(
            "Page.addScriptToEvaluateOnNewDocument",
            &json!({ "source": CONSOLE_CAPTURE_JS }),
        )
        .await
        .map_err(|e| e.to_string())?;
    // Best-effort: this IIFE has no return value, so evaluate_sync::<String>
    // routinely "fails" to deserialize a String even when the injection
    // itself succeeded. The addScriptToEvaluateOnNewDocument call above is
    // what actually needs to succeed; this just seeds the *current* page too.
    let _: String = tab
        .page
        .evaluate_sync(CONSOLE_CAPTURE_JS)
        .await
        .unwrap_or_default();
    tab.console_injected = true;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::{restore_cookies, SavedCookie, SavedState};
    use serde_json::{json, Value};
    use std::collections::HashMap;

    fn cookie(name: &str, value: &str, domain: &str) -> eoka::SessionCookie {
        eoka::SessionCookie {
            name: name.into(),
            value: value.into(),
            domain: domain.into(),
            path: "/".into(),
            secure: false,
            http_only: false,
            same_site: None,
            expires: None,
        }
    }

    #[tokio::test]
    #[ignore = "requires Chrome"]
    async fn restore_replaces_covered_cookies_and_keeps_other_sites() {
        let browser = eoka::Browser::launch_with(|config| {
            config.patch_binary = false;
            config.chrome_path = std::env::var("EOKA_TEST_CHROME_PATH").ok();
        })
        .await
        .unwrap();
        let page = browser.new_blank_page().await.unwrap();
        page.set_cookies_bulk(vec![
            cookie("session", "current", "app.test"),
            cookie("added_later", "1", "app.test"),
            cookie("other_account", "1", "auth.app.test"),
            cookie("sso", "keep-me", "sso.test"),
        ])
        .await
        .unwrap();

        let state = SavedState {
            url: "http://app.test/".into(),
            cookies: vec![SavedCookie {
                name: "session".into(),
                value: "saved".into(),
                domain: "app.test".into(),
                path: "/".into(),
                expires: 0.0,
                http_only: false,
                secure: false,
                same_site: None,
            }],
            local_storage: HashMap::new(),
            session_storage: HashMap::new(),
            user_agent: String::new(),
            saved_at: String::new(),
        };
        restore_cookies(&page, &state).await.unwrap();

        let jar: Value = page
            .session()
            .send("Storage.getCookies", &json!({}))
            .await
            .unwrap();
        let mut cookies: Vec<_> = jar["cookies"]
            .as_array()
            .unwrap()
            .iter()
            .map(|c| (c["name"].as_str().unwrap(), c["value"].as_str().unwrap()))
            .collect();
        cookies.sort();
        assert_eq!(cookies, [("session", "saved"), ("sso", "keep-me")]);

        // A page in another browser context restores into its own cookies.
        let transport = page.session().transport();
        let context: Value = transport
            .send("Target.createBrowserContext", &json!({}))
            .await
            .unwrap();
        let target: Value = transport
            .send(
                "Target.createTarget",
                &json!({"url": "about:blank", "browserContextId": context["browserContextId"]}),
            )
            .await
            .unwrap();
        let isolated = browser
            .attach_page(target["targetId"].as_str().unwrap())
            .await
            .unwrap();
        isolated
            .set_cookies_bulk(vec![cookie("session", "isolated", "app.test")])
            .await
            .unwrap();
        restore_cookies(&isolated, &state).await.unwrap();
        let jar: Value = transport
            .send(
                "Storage.getCookies",
                &json!({"browserContextId": context["browserContextId"]}),
            )
            .await
            .unwrap();
        assert_eq!(jar["cookies"].as_array().unwrap().len(), 1);
        assert_eq!(jar["cookies"][0]["value"], "saved");

        browser.close().await.unwrap();
    }

    #[test]
    fn loads_saved_state_shape() {
        let state: SavedState = serde_json::from_str(
            r#"{
                "url": "https://www.recreation.gov/cart",
                "cookies": [{
                    "name": "session", "value": "redacted", "domain": ".recreation.gov",
                    "path": "/", "expires": 0.0, "http_only": true,
                    "secure": true, "same_site": null
                }],
                "local_storage": {"recaccount": "value"},
                "session_storage": {"key": "value"},
                "user_agent": "Mozilla/5.0"
            }"#,
        )
        .expect("saved state should be accepted");

        assert_eq!(state.cookies[0].expires, 0.0);
        assert_eq!(state.local_storage["recaccount"], "value");
        assert_eq!(state.session_storage["key"], "value");
        assert_eq!(state.user_agent, "Mozilla/5.0");
        assert!(state.saved_at.is_empty());
    }
}