tail-fin-gemini 0.7.3

//! `reqwest`-backed Gemini client — no Chrome, no DOM. Talks directly to
//! the same HTTP endpoints `gemini.google.com` calls from the browser.

use std::path::{Path, PathBuf};

use reqwest::{header, Client, ClientBuilder, StatusCode};
use serde_json::{json, Value};
use tail_fin_common::TailFinError;

use crate::cookies::{load_netscape, Cookies};
use crate::parsing::{
    extract_response_text, extract_session_tokens, extract_turn_ids, SessionTokens,
};
use crate::signing::{current_sapisidhash, ORIGIN};
use crate::types::GeminiResponse;

const APP_URL: &str = "https://gemini.google.com/app";
const STREAM_URL: &str =
    "https://gemini.google.com/_/BardChatUi/data/assistant.lamda.BardFrontendService/StreamGenerate";
const UPLOAD_URL: &str = "https://push.clients6.google.com/upload/";
const USER_AGENT: &str = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 \
     (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36";

pub struct GeminiClient {
    http: Client,
    cookies: Cookies,
}

impl GeminiClient {
    pub fn new(cookies: Cookies) -> Result<Self, TailFinError> {
        let http = ClientBuilder::new()
            .user_agent(USER_AGENT)
            .redirect(reqwest::redirect::Policy::none())
            .build()
            .map_err(|e| TailFinError::Api(format!("build HTTP client: {e}")))?;
        Ok(Self { http, cookies })
    }

    /// Convenience constructor: load Netscape-format cookie file and wrap.
    pub fn from_cookie_file(path: &Path) -> Result<Self, TailFinError> {
        let cookies = load_netscape(path)?;
        Self::new(cookies)
    }

    /// Start a fresh conversation: send a prompt (optionally with a
    /// file attachment) and return the reply. Use [`ask_continue`] with
    /// the returned `GeminiResponse` to send follow-up turns.
    ///
    /// [`ask_continue`]: Self::ask_continue
    pub async fn ask(
        &self,
        prompt: &str,
        attach: Option<&Path>,
    ) -> Result<GeminiResponse, TailFinError> {
        self.ask_inner(prompt, attach, None).await
    }

    /// Continue an existing conversation. `prev` must carry all three
    /// ids (`conversation_id`, `response_id`, `choice_id`) — typically
    /// the `GeminiResponse` from a prior call. If any id is missing or
    /// malformed (wrong prefix, literal "null", empty), errors *before*
    /// sending the request.
    pub async fn ask_continue(
        &self,
        prompt: &str,
        attach: Option<&Path>,
        prev: &GeminiResponse,
    ) -> Result<GeminiResponse, TailFinError> {
        let (cid, rid, rcid) = prev.require_continuation()?;
        self.ask_continue_with_ids(prompt, attach, cid, rid, rcid)
            .await
    }

    /// Same as `ask_continue` but takes the three ids as strings —
    /// convenience for callers that persist ids elsewhere (e.g. CLI
    /// flags or a database) instead of holding a `GeminiResponse`.
    ///
    /// Validates prefixes and common foot-guns (empty, literal "null")
    /// before dispatching.
    pub async fn ask_continue_with_ids(
        &self,
        prompt: &str,
        attach: Option<&Path>,
        cid: &str,
        rid: &str,
        rcid: &str,
    ) -> Result<GeminiResponse, TailFinError> {
        validate_continuation_ids(cid, rid, rcid)?;
        self.ask_inner(
            prompt,
            attach,
            Some((cid.to_string(), rid.to_string(), rcid.to_string())),
        )
        .await
    }

    async fn ask_inner(
        &self,
        prompt: &str,
        attach: Option<&Path>,
        continuation: Option<(String, String, String)>,
    ) -> Result<GeminiResponse, TailFinError> {
        let tokens = self.fetch_session_tokens().await?;
        let files = if let Some(path) = attach {
            let filename = filename_of(path)?;
            let url = self.upload_file(path, &filename, &tokens.push_id).await?;
            vec![(url, filename)]
        } else {
            vec![]
        };
        let body = self
            .post_stream_generate(prompt, &tokens, &files, continuation.as_ref())
            .await?;
        let response = extract_response_text(&body)?;
        let ids = extract_turn_ids(&body);

        // Stale-id guard: if we sent continuation ids but the server
        // returned a DIFFERENT conversation_id (or none at all), it
        // silently dropped our continuation and started a new session.
        // The caller almost certainly doesn't want that — they asked
        // a follow-up expecting prior context, now their prompt is
        // talking to a stranger. Fail loud.
        if let Some((sent_cid, _, _)) = continuation.as_ref() {
            // Trim before compare — Gemini has never returned padded
            // ids in observation, but byte-equality on a
            // defensive-check is a classic flake source.
            let sent_trim = sent_cid.trim();
            match ids.conversation_id.as_deref().map(str::trim) {
                Some(got) if got == sent_trim => {}
                Some(got) => {
                    return Err(TailFinError::Api(format!(
                        "continuation silently dropped: sent conversation_id={sent_cid} but \
                         server returned {got}. The prior turn may have expired — start a \
                         new conversation with `ask` instead of `ask_continue`."
                    )));
                }
                None => {
                    return Err(TailFinError::Api(format!(
                        "continuation silently dropped: sent conversation_id={sent_cid} but \
                         server response carried no conversation_id. Retry or start fresh."
                    )));
                }
            }
        }

        Ok(GeminiResponse {
            response,
            conversation_id: ids.conversation_id,
            response_id: ids.response_id,
            choice_id: ids.choice_id,
        })
    }

    async fn fetch_session_tokens(&self) -> Result<SessionTokens, TailFinError> {
        let resp = self
            .http
            .get(APP_URL)
            .header(header::COOKIE, self.cookies.to_header())
            .header(header::USER_AGENT, USER_AGENT)
            .send()
            .await
            .map_err(|e| TailFinError::Api(format!("fetch /app: {e}")))?;
        let status = resp.status();
        if !status.is_success() {
            return Err(classify_status(status, "fetch /app"));
        }
        let html = resp
            .text()
            .await
            .map_err(|e| TailFinError::Api(format!("/app body: {e}")))?;
        extract_session_tokens(&html)
    }

    /// Resumable upload, two-step — see docs/sites/gemini.md for the
    /// protocol trace. Caller supplies the already-validated UTF-8
    /// filename (same value is later threaded into the f.req payload).
    async fn upload_file(
        &self,
        path: &Path,
        filename: &str,
        push_id: &str,
    ) -> Result<String, TailFinError> {
        let bytes = tokio::fs::read(path)
            .await
            .map_err(|e| TailFinError::Io(format!("read {}: {e}", path.display())))?;
        let size = bytes.len();

        // Step 1 — initiate.
        let init_body = format!("File name={}", urlencoding::encode(filename));
        let init = self
            .http
            .post(UPLOAD_URL)
            .header("Push-ID", push_id)
            .header("X-Tenant-Id", "bard-storage")
            .header("X-Goog-Upload-Command", "start")
            .header("X-Goog-Upload-Protocol", "resumable")
            .header("X-Goog-Upload-Header-Content-Length", size.to_string())
            .header(header::USER_AGENT, USER_AGENT)
            .header(header::ORIGIN, ORIGIN)
            .header(header::REFERER, "https://gemini.google.com/")
            .header(header::COOKIE, self.cookies.to_header())
            .header(
                header::CONTENT_TYPE,
                "application/x-www-form-urlencoded;charset=UTF-8",
            )
            .body(init_body)
            .send()
            .await
            .map_err(|e| TailFinError::Api(format!("upload init: {e}")))?;

        let init_status = init.status();
        let session_url = init
            .headers()
            .get("X-Goog-Upload-URL")
            .and_then(|v| v.to_str().ok())
            .map(str::to_string);
        if !init_status.is_success() || session_url.is_none() {
            if !init_status.is_success() {
                return Err(classify_status(init_status, "upload init"));
            }
            let body = init.text().await.unwrap_or_default();
            return Err(TailFinError::Api(format!(
                "upload init succeeded but response lacked X-Goog-Upload-URL \
                 (Gemini protocol may have changed): {}",
                truncate(&body, 300)
            )));
        }
        let session_url = session_url.unwrap();

        // Step 2 — upload + finalize.
        let resp = self
            .http
            .post(&session_url)
            .header("X-Goog-Upload-Command", "upload, finalize")
            .header("X-Goog-Upload-Offset", "0")
            .header(header::USER_AGENT, USER_AGENT)
            .body(bytes)
            .send()
            .await
            .map_err(|e| TailFinError::Api(format!("upload finalize: {e}")))?;
        let status = resp.status();
        if !status.is_success() {
            return Err(classify_status(status, "upload finalize"));
        }
        let body = resp
            .text()
            .await
            .map_err(|e| TailFinError::Api(format!("upload finalize body: {e}")))?;
        let url = body.trim().to_string();
        if url.is_empty() {
            return Err(TailFinError::Parse(
                "upload finalize returned empty body".into(),
            ));
        }
        Ok(url)
    }

    async fn post_stream_generate(
        &self,
        prompt: &str,
        tokens: &SessionTokens,
        files: &[(String, String)],
        continuation: Option<&(String, String, String)>,
    ) -> Result<String, TailFinError> {
        let sapisid = self.cookies.require("SAPISID")?;
        let auth = current_sapisidhash(sapisid);

        // `files_data` in f.req inner: each attachment is `[[<url>,1],<name>]`.
        let files_data: Vec<Value> = files
            .iter()
            .map(|(url, name)| json!([[url, 1], name]))
            .collect();

        // Thread previous turn's ids when continuing a conversation;
        // `[null, null, null]` starts fresh.
        let turn_ids: Value = match continuation {
            Some((cid, rid, rcid)) => json!([cid, rid, rcid]),
            None => json!([null, null, null]),
        };

        let inner = json!([
            [prompt, 0, null, files_data, null, null, 0],
            ["en"],
            turn_ids,
            null,
            null,
            null,
            [1],
            0,
            [],
            [],
            1,
            0,
        ]);
        let inner_str = serde_json::to_string(&inner)
            .map_err(|e| TailFinError::Parse(format!("encode f.req inner: {e}")))?;
        let f_req = json!([null, inner_str]).to_string();

        let form = [("at", tokens.snlm0e.as_str()), ("f.req", f_req.as_str())];

        let resp = self
            .http
            .post(STREAM_URL)
            .query(&[
                ("bl", tokens.build_label.as_str()),
                ("_reqid", "0"),
                ("rt", "c"),
            ])
            .header(header::COOKIE, self.cookies.to_header())
            .header(header::USER_AGENT, USER_AGENT)
            .header(header::ORIGIN, ORIGIN)
            .header(header::REFERER, APP_URL)
            .header(header::AUTHORIZATION, auth)
            .header("X-Same-Domain", "1")
            .form(&form)
            .send()
            .await
            .map_err(|e| TailFinError::Api(format!("StreamGenerate: {e}")))?;

        let status = resp.status();
        if !status.is_success() {
            return Err(classify_status(status, "StreamGenerate"));
        }
        let body = resp
            .text()
            .await
            .map_err(|e| TailFinError::Api(format!("StreamGenerate body: {e}")))?;
        Ok(body)
    }
}

/// Turn a non-success HTTP status into an actionable `TailFinError`.
/// The messages hint at the specific fix (re-export cookies, back off,
/// retry) rather than just echoing the status code.
fn classify_status(status: StatusCode, op: &str) -> TailFinError {
    match status {
        StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => TailFinError::Api(format!(
            "{op}: {status} — cookies are stale or the account lacks access. \
             Re-export Gemini cookies from a logged-in browser."
        )),
        StatusCode::TOO_MANY_REQUESTS => TailFinError::Api(format!(
            "{op}: {status} — rate-limited by Gemini. Back off and retry in a few minutes."
        )),
        StatusCode::NOT_FOUND => TailFinError::Api(format!(
            "{op}: {status} — endpoint URL may have drifted. \
             Check STREAM_URL / UPLOAD_URL / APP_URL in crates/tail-fin-gemini/src/client.rs."
        )),
        s if s.is_server_error() => TailFinError::Api(format!(
            "{op}: {status} — Google-side error, usually transient. Retry."
        )),
        _ => TailFinError::Api(format!("{op}: {status}")),
    }
}

/// Same validation as `GeminiResponse::require_continuation`, but for
/// loose id strings (CLI / persisted state).
fn validate_continuation_ids(cid: &str, rid: &str, rcid: &str) -> Result<(), TailFinError> {
    let mut bad = Vec::new();
    if !is_valid_id(cid, "c_") {
        bad.push("conversation_id (expected `c_…`)");
    }
    if !is_valid_id(rid, "r_") {
        bad.push("response_id (expected `r_…`)");
    }
    if !is_valid_id(rcid, "rc_") {
        bad.push("choice_id (expected `rc_…`)");
    }
    if bad.is_empty() {
        Ok(())
    } else {
        Err(TailFinError::Api(format!(
            "continuation ids invalid: {}",
            bad.join(", ")
        )))
    }
}

fn is_valid_id(s: &str, prefix: &str) -> bool {
    let t = s.trim();
    !t.is_empty() && t != "null" && t.starts_with(prefix)
}

/// Extract the filename as a UTF-8 string, erroring cleanly instead of
/// silently renaming to `"upload"` — users need to know their file's
/// original name made it into the upload.
fn filename_of(path: &Path) -> Result<String, TailFinError> {
    path.file_name()
        .and_then(|n| n.to_str())
        .map(str::to_string)
        .ok_or_else(|| {
            TailFinError::Api(format!(
                "attachment filename is not valid UTF-8: {}",
                path.display()
            ))
        })
}

/// Default session file location: `~/.tail-fin/gemini-cookies.txt`.
pub fn default_cookies_path() -> PathBuf {
    let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
    PathBuf::from(home).join(".tail-fin/gemini-cookies.txt")
}

fn truncate(s: &str, max: usize) -> String {
    if s.len() <= max {
        s.to_string()
    } else {
        format!("{}… (truncated)", &s[..max])
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use reqwest::StatusCode;

    #[test]
    fn classify_401_hints_at_cookies() {
        let err = classify_status(StatusCode::UNAUTHORIZED, "op");
        let msg = format!("{err}");
        assert!(msg.contains("stale"), "got: {msg}");
        assert!(msg.contains("Re-export"), "got: {msg}");
    }

    #[test]
    fn classify_429_hints_at_rate_limit() {
        let err = classify_status(StatusCode::TOO_MANY_REQUESTS, "op");
        assert!(format!("{err}").contains("rate-limited"));
    }

    #[test]
    fn classify_500_hints_at_google_side() {
        let err = classify_status(StatusCode::INTERNAL_SERVER_ERROR, "op");
        assert!(format!("{err}").contains("Google-side"));
    }

    #[test]
    fn classify_404_hints_at_endpoint_drift() {
        let err = classify_status(StatusCode::NOT_FOUND, "op");
        let msg = format!("{err}");
        assert!(msg.contains("endpoint URL"), "got: {msg}");
        assert!(msg.contains("client.rs"), "got: {msg}");
    }

    #[cfg(unix)]
    #[test]
    fn filename_of_rejects_non_utf8_name() {
        use std::ffi::OsStr;
        use std::os::unix::ffi::OsStrExt;
        use std::path::PathBuf;

        // 0xff is invalid UTF-8 — constructs a real non-UTF-8 filename,
        // not an empty-path stand-in.
        let mut p = PathBuf::from("/tmp");
        p.push(OsStr::from_bytes(&[0xff, 0xfe, b'.', b'j', b'p', b'g']));
        let err = filename_of(&p).unwrap_err();
        assert!(format!("{err}").contains("not valid UTF-8"));
    }

    #[test]
    fn filename_of_accepts_normal_path() {
        assert_eq!(
            filename_of(Path::new("/tmp/image.jpg")).unwrap(),
            "image.jpg"
        );
    }

    #[test]
    fn default_cookies_path_ends_with_expected_file() {
        let p = default_cookies_path();
        assert!(p.ends_with("gemini-cookies.txt"));
    }

    #[test]
    fn validate_continuation_ids_accepts_well_shaped() {
        assert!(validate_continuation_ids("c_abc", "r_def", "rc_ghi").is_ok());
    }

    #[test]
    fn validate_continuation_ids_rejects_literal_null() {
        let err = validate_continuation_ids("null", "r_ok", "rc_ok").unwrap_err();
        assert!(format!("{err}").contains("conversation_id"));
    }

    #[test]
    fn validate_continuation_ids_rejects_wrong_prefix() {
        let err = validate_continuation_ids("r_wrongslot", "r_ok", "rc_ok").unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("conversation_id"), "got: {msg}");
    }

    #[test]
    fn validate_continuation_ids_lists_all_bad_fields() {
        let err = validate_continuation_ids("", "null", "wrong").unwrap_err();
        let msg = format!("{err}");
        assert!(msg.contains("conversation_id"));
        assert!(msg.contains("response_id"));
        assert!(msg.contains("choice_id"));
    }

    #[test]
    fn validate_continuation_ids_trims_whitespace() {
        // A shell variable might carry trailing space; we accept if the
        // trimmed value is well-shaped.
        assert!(validate_continuation_ids(" c_abc ", "r_def\n", "\trc_ghi").is_ok());
    }

    #[test]
    fn validate_continuation_ids_rejects_whitespace_only() {
        let err = validate_continuation_ids("   ", "r_ok", "rc_ok").unwrap_err();
        assert!(format!("{err}").contains("conversation_id"));
    }
}