captchaforge 0.2.9

use super::*;

// ─── Ollama / VLM helpers ────────────────────────────────────────────────────

const DEFAULT_OLLAMA_BASE: &str = "http://localhost:11434";
const DEFAULT_OLLAMA_MODEL: &str = "qwen3-vl:30b";
const ENV_VLM_ENDPOINT: &str = "CAPTCHAFORGE_VLM_ENDPOINT";
const ENV_VLM_MODEL: &str = "CAPTCHAFORGE_VLM_MODEL";

/// Raw Ollama `/api/generate` request body.
#[derive(Debug, Serialize)]
struct OllamaRequest<'a> {
    model: &'a str,
    prompt: &'a str,
    images: Vec<String>,
    stream: bool,
}

/// Relevant fields from Ollama `/api/generate` response.
#[derive(Debug, Deserialize)]
struct OllamaResponse {
    response: String,
}

/// Send a screenshot (base64 PNG/JPEG) and a text prompt to a VLM endpoint.
/// Returns the model's raw text reply.
async fn vlm_query(
    client: &reqwest::Client,
    endpoint: &str,
    model: &str,
    image_b64: &str,
    prompt: &str,
    timeout_ms: u64,
) -> Result<String> {
    let req = OllamaRequest {
        model,
        prompt,
        images: vec![image_b64.to_string()],
        stream: false,
    };

    let resp = client
        .post(format!("{}/api/generate", endpoint.trim_end_matches('/')))
        .json(&req)
        .timeout(Duration::from_millis(timeout_ms))
        .send()
        .await
        .map_err(|e| anyhow!("vlm request failed: {}", e))?;

    if !resp.status().is_success() {
        let status = resp.status();
        let body = resp.text().await.unwrap_or_default();
        return Err(anyhow!("vlm returned {}: {}", status, body));
    }

    let ollama_resp: OllamaResponse = resp
        .json()
        .await
        .map_err(|e| anyhow!("vlm json parse: {}", e))?;

    Ok(ollama_resp.response)
}

/// Extract the first JSON object embedded in free-form text (handles markdown fences).
pub(crate) fn extract_json(text: &str) -> Option<&str> {
    if let Some(start) = text.find("```json") {
        let rest = &text[start + 7..];
        if let Some(end) = rest.find("```") {
            return Some(rest[..end].trim());
        }
    }
    if let Some(start) = text.find('{') {
        let rest = &text[start..];
        let mut depth = 0;
        let mut end = None;
        for (i, c) in rest.char_indices() {
            match c {
                '{' => depth += 1,
                '}' => {
                    depth -= 1;
                    if depth == 0 {
                        end = Some(i + 1);
                        break;
                    }
                }
                _ => {}
            }
        }
        if let Some(e) = end {
            return Some(&rest[..e]);
        }
    }
    None
}

// ─── VlmCaptchaSolver ─────────────────────────────────────────────────────────

/// Solves image-based CAPTCHAs via a vision-language model (VLM) such as
/// Ollama-hosted qwen3-vl.
///
/// Supported sub-types:
///   * Image-grid   (reCAPTCHA / hCaptcha "select all traffic lights")
///   * Text-based   (type the characters shown)
///   * Click-target (click on a specific object / button)
pub struct VlmCaptchaSolver {
    pub(crate) client: reqwest::Client,
    pub(crate) endpoint: String,
    pub(crate) model: String,
    pub(crate) config: SolveConfig,
}

impl Default for VlmCaptchaSolver {
    fn default() -> Self {
        Self::new()
    }
}

impl VlmCaptchaSolver {
    pub fn new() -> Self {
        Self::new_with_env(|name| std::env::var(name).ok())
    }

    /// Construct a solver, using the supplied closure to look up env
    /// variables. The public [`Self::new`] uses [`std::env::var`];
    /// tests use this overload to assert env precedence without
    /// mutating real process env (which is `unsafe` in the 2024
    /// edition and forbidden by the crate's `#![forbid(unsafe_code)]`).
    pub(crate) fn new_with_env<F>(env: F) -> Self
    where
        F: Fn(&str) -> Option<String>,
    {
        let config = SolveConfig::default();
        let endpoint = env(ENV_VLM_ENDPOINT)
            .filter(|s| !s.is_empty())
            .unwrap_or_else(|| DEFAULT_OLLAMA_BASE.to_string());
        let model = env(ENV_VLM_MODEL)
            .filter(|s| !s.is_empty())
            .unwrap_or_else(|| DEFAULT_OLLAMA_MODEL.to_string());
        Self {
            client: reqwest::Client::builder()
                .timeout(Duration::from_millis(config.client_http_timeout_ms))
                .build()
                .unwrap_or_else(|_| reqwest::Client::new()),
            endpoint,
            model,
            config,
        }
    }

    /// Override the VLM endpoint (default: `http://localhost:11434`).
    pub fn with_endpoint(mut self, url: impl Into<String>) -> Self {
        self.endpoint = url.into();
        self
    }

    /// Override the VLM model name (default: `qwen3-vl:30b`).
    pub fn with_model(mut self, model: impl Into<String>) -> Self {
        self.model = model.into();
        self
    }

    /// Provide a custom solve configuration.
    pub fn with_config(mut self, config: SolveConfig) -> Self {
        self.config = config;
        self
    }

    pub fn grid_prompt(task: &str) -> String {
        format!(
            r#"You are solving an image-grid CAPTCHA.
Task: {}
Respond with a JSON object exactly like this:
{{"selected": [0,1,2], "confidence": 0.95}}
where `selected` is a list of zero-based tile indices that match the task.
If none match, use an empty list.
Only respond with the JSON object, no extra text."#,
            task
        )
    }

    pub fn text_captcha_prompt() -> String {
        r#"You are solving a text CAPTCHA.
Read the text in the image and respond with a JSON object exactly like this:
{"text": "answer", "confidence": 0.95}
Only respond with the JSON object, no extra text."#
            .to_string()
    }

    pub fn click_target_prompt(target: &str) -> String {
        format!(
            r#"You are solving a click-target CAPTCHA.
Click the object described as: "{}"
Respond with a JSON object exactly like this:
{{"x": 123, "y": 456, "confidence": 0.95}}
where x and y are pixel coordinates in the screenshot.
Only respond with the JSON object, no extra text."#,
            target
        )
    }
}

#[async_trait]
impl CaptchaSolver for VlmCaptchaSolver {
    fn name(&self) -> &'static str {
        "VlmCaptchaSolver"
    }

    fn method(&self) -> SolveMethod {
        SolveMethod::VisionLLM
    }

    fn supports(&self, kind: &crate::captcha_detect::DetectedCaptcha) -> bool {
        use crate::captcha_detect::DetectedCaptcha;
        // VLM screenshot+vision can attempt any visual captcha,
        // including TOML-rule vendors whose providers recommend
        // VisionLLM as a fallback method.
        matches!(
            kind,
            DetectedCaptcha::RecaptchaV2
                | DetectedCaptcha::RecaptchaV3
                | DetectedCaptcha::HCaptcha
                | DetectedCaptcha::ImageCaptcha
                | DetectedCaptcha::Turnstile
                | DetectedCaptcha::CanvasCaptcha
                | DetectedCaptcha::ShadowDomCaptcha
                | DetectedCaptcha::MultiStepCaptcha
                | DetectedCaptcha::SliderCaptcha
                | DetectedCaptcha::Custom(_)
        )
    }

    async fn solve(
        &self,
        page: &Page,
        captcha_info: &crate::captcha_detect::CaptchaInfo,
    ) -> Result<CaptchaSolveResult> {
        let t0 = Instant::now();
        let image_b64 = screenshot_b64(page).await?;

        // Use the detected kind to pick a prompt strategy.
        match captcha_info.kind {
            crate::captcha_detect::DetectedCaptcha::RecaptchaV2
            | crate::captcha_detect::DetectedCaptcha::HCaptcha
            | crate::captcha_detect::DetectedCaptcha::ImageCaptcha => {
                // Image grid challenge.
                let prompt = Self::grid_prompt("Select all matching images");
                let raw = vlm_query(
                    &self.client,
                    &self.endpoint,
                    &self.model,
                    &image_b64,
                    &prompt,
                    self.config.vlm_http_timeout_ms,
                )
                .await?;

                let json_str = extract_json(&raw).unwrap_or(&raw);
                let val: serde_json::Value =
                    serde_json::from_str(json_str).unwrap_or(serde_json::Value::Null);
                let confidence = val["confidence"].as_f64().unwrap_or(0.5) as f32;
                let selected: Vec<usize> = val["selected"]
                    .as_array()
                    .unwrap_or(&vec![])
                    .iter()
                    .filter_map(|v| v.as_u64().map(|n| n as usize))
                    .collect();

                // Click each selected tile via direct DOM dispatch
                // — far more reliable than computing mouse coordinates
                // across iframe boundaries. Probe the grid first to
                // get the actual tile elements + dimensions, then
                // dispatch a click event on each by index.
                //
                // Grid probes per vendor:
                //   - reCAPTCHA v2:  .rc-imageselect-tile (img wrapper)
                //   - hCaptcha:      .task-image (sometimes .image)
                //   - Generic:       td > img inside the grid table
                let click_indices_js = format!(
                    r#"
                    ((indices) => {{
                        const probes = [
                            '.rc-imageselect-tile',
                            '.task-image',
                            '.image-task',
                            '.hcaptcha-checkbox-img',
                            '.rc-imageselect-table td',
                            '.hcaptcha-table td',
                            '.captcha-grid > .tile',
                        ];
                        let tiles = null;
                        for (const sel of probes) {{
                            const els = document.querySelectorAll(sel);
                            if (els.length > 0) {{ tiles = Array.from(els); break; }}
                        }}
                        if (!tiles || tiles.length === 0) return {{ ok: false, count: 0 }};
                        let clicked = 0;
                        for (const idx of indices) {{
                            const el = tiles[idx];
                            if (!el) continue;
                            const r = el.getBoundingClientRect();
                            const cx = r.left + r.width / 2;
                            const cy = r.top + r.height / 2;
                            ['mousedown','mouseup','click'].forEach(t => {{
                                el.dispatchEvent(new MouseEvent(t, {{
                                    bubbles: true, button: 0,
                                    clientX: cx, clientY: cy
                                }}));
                            }});
                            clicked++;
                        }}
                        /* Find + click the verify/submit button. */
                        const verifyProbes = [
                            '#recaptcha-verify-button',
                            '.rc-button-default',
                            '.button-submit',
                            '[data-pp="submit"]',
                            'button[type="submit"]',
                        ];
                        for (const sel of verifyProbes) {{
                            const b = document.querySelector(sel);
                            if (b) {{ b.click(); break; }}
                        }}
                        return {{ ok: true, count: clicked, total: tiles.length }};
                    }})({selected})
                    "#,
                    selected = serde_json::to_string(&selected).unwrap_or_else(|_| "[]".into())
                );
                let _ = page.evaluate(click_indices_js).await;

                let success = confidence > 0.5 && !selected.is_empty();
                let cookies = if success {
                    crate::cookies::capture_from_page(page)
                        .await
                        .unwrap_or_default()
                } else {
                    Vec::new()
                };
                Ok(CaptchaSolveResult {
                    solution: serde_json::to_string(&selected).unwrap_or_default(),
                    confidence,
                    method: SolveMethod::VisionLLM,
                    time_ms: t0.elapsed().as_millis() as u64,
                    success,
                    screenshot: None,
                    cookies,
                })
            }
            crate::captcha_detect::DetectedCaptcha::RecaptchaV3 => {
                // Text-based challenge.
                let prompt = Self::text_captcha_prompt();
                let raw = vlm_query(
                    &self.client,
                    &self.endpoint,
                    &self.model,
                    &image_b64,
                    &prompt,
                    self.config.vlm_http_timeout_ms,
                )
                .await?;

                let json_str = extract_json(&raw).unwrap_or(&raw);
                let val: serde_json::Value =
                    serde_json::from_str(json_str).unwrap_or(serde_json::Value::Null);
                let text = val["text"].as_str().unwrap_or("").trim().to_string();
                let confidence = val["confidence"].as_f64().unwrap_or(0.5) as f32;

                // Type the answer into the response field if present.
                if let Ok(input) = page
                    .find_element("input[type=text], textarea, .rc-response-input")
                    .await
                {
                    input.click().await.ok();
                    input.type_str(&text).await.ok();
                }

                let success = confidence > 0.5 && !text.is_empty();
                let cookies = if success {
                    crate::cookies::capture_from_page(page)
                        .await
                        .unwrap_or_default()
                } else {
                    Vec::new()
                };
                Ok(CaptchaSolveResult {
                    solution: text.clone(),
                    confidence,
                    method: SolveMethod::VisionLLM,
                    time_ms: t0.elapsed().as_millis() as u64,
                    success,
                    screenshot: None,
                    cookies,
                })
            }
            _ => {
                // Generic click-target fallback.
                let prompt = Self::click_target_prompt("the CAPTCHA checkbox or submit button");
                let raw = vlm_query(
                    &self.client,
                    &self.endpoint,
                    &self.model,
                    &image_b64,
                    &prompt,
                    self.config.vlm_http_timeout_ms,
                )
                .await?;

                let json_str = extract_json(&raw).unwrap_or(&raw);
                let val: serde_json::Value =
                    serde_json::from_str(json_str).unwrap_or(serde_json::Value::Null);
                let x = val["x"].as_f64().unwrap_or(0.0);
                let y = val["y"].as_f64().unwrap_or(0.0);
                let confidence = val["confidence"].as_f64().unwrap_or(0.0) as f32;

                let viewport = page
                    .evaluate("window.innerWidth + ',' + window.innerHeight")
                    .await?
                    .into_value::<String>()
                    .unwrap_or_default();
                let mut parts = viewport.split(',');
                let vw: f64 = parts.next().and_then(|s| s.parse().ok()).unwrap_or(1920.0);
                let vh: f64 = parts.next().and_then(|s| s.parse().ok()).unwrap_or(1080.0);

                if x > 0.0 && y > 0.0 && x < vw && y < vh {
                    crate::behavior::click_realistic(page, x, y).await?;
                }

                let success = confidence > 0.5;
                let cookies = if success {
                    crate::cookies::capture_from_page(page)
                        .await
                        .unwrap_or_default()
                } else {
                    Vec::new()
                };
                Ok(CaptchaSolveResult {
                    solution: format!("click:{},{}:{:.2}", x as i64, y as i64, confidence),
                    confidence,
                    method: SolveMethod::VisionLLM,
                    time_ms: t0.elapsed().as_millis() as u64,
                    success,
                    screenshot: None,
                    cookies,
                })
            }
        }
    }
}

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

    #[test]
    fn vlm_solver_defaults_when_env_returns_none() {
        let s = VlmCaptchaSolver::new_with_env(|_| None);
        assert_eq!(s.endpoint, "http://localhost:11434");
        assert_eq!(s.model, "qwen3-vl:30b");
    }

    #[test]
    fn vlm_solver_env_overrides_defaults() {
        let s = VlmCaptchaSolver::new_with_env(|k| match k {
            "CAPTCHAFORGE_VLM_ENDPOINT" => Some("http://env-host:11434".into()),
            "CAPTCHAFORGE_VLM_MODEL" => Some("env-model:latest".into()),
            _ => None,
        });
        assert_eq!(s.endpoint, "http://env-host:11434");
        assert_eq!(s.model, "env-model:latest");
    }

    #[test]
    fn vlm_solver_builder_overrides_env() {
        let s = VlmCaptchaSolver::new_with_env(|k| match k {
            "CAPTCHAFORGE_VLM_ENDPOINT" => Some("http://env-host:11434".into()),
            "CAPTCHAFORGE_VLM_MODEL" => Some("env-model:latest".into()),
            _ => None,
        })
        .with_endpoint("http://builder:11434")
        .with_model("builder-model:7b");
        assert_eq!(s.endpoint, "http://builder:11434");
        assert_eq!(s.model, "builder-model:7b");
    }

    #[test]
    fn vlm_solver_empty_env_falls_through_to_defaults() {
        let s = VlmCaptchaSolver::new_with_env(|_| Some(String::new()));
        assert_eq!(s.endpoint, "http://localhost:11434");
        assert_eq!(s.model, "qwen3-vl:30b");
    }

    #[test]
    fn vlm_solver_env_constants_match_documented_names() {
        // The CLAUDE.md / README contract is the env-var names. Pin
        // them so a rename is a deliberate breaking change.
        assert_eq!(ENV_VLM_ENDPOINT, "CAPTCHAFORGE_VLM_ENDPOINT");
        assert_eq!(ENV_VLM_MODEL, "CAPTCHAFORGE_VLM_MODEL");
    }

    #[test]
    fn vlm_solver_custom_endpoint() {
        let s = VlmCaptchaSolver::new().with_endpoint("http://ollama.internal:11434");
        assert_eq!(s.endpoint, "http://ollama.internal:11434");
    }

    #[test]
    fn vlm_solver_custom_model() {
        let s = VlmCaptchaSolver::new().with_model("llava:13b");
        assert_eq!(s.model, "llava:13b");
    }

    #[test]
    fn vlm_solver_chained_builders() {
        let s = VlmCaptchaSolver::new()
            .with_endpoint("http://gpu-box:11434")
            .with_model("qwen3-vl:72b");
        assert_eq!(s.endpoint, "http://gpu-box:11434");
        assert_eq!(s.model, "qwen3-vl:72b");
    }

    #[test]
    fn grid_prompt_includes_task() {
        let prompt = VlmCaptchaSolver::grid_prompt("Select all traffic lights");
        assert!(prompt.contains("Select all traffic lights"));
        assert!(prompt.contains("selected"));
        assert!(prompt.contains("confidence"));
    }

    #[test]
    fn text_captcha_prompt_is_valid() {
        let prompt = VlmCaptchaSolver::text_captcha_prompt();
        assert!(prompt.contains("text"));
        assert!(prompt.contains("confidence"));
    }

    #[test]
    fn click_target_prompt_includes_target() {
        let prompt = VlmCaptchaSolver::click_target_prompt("submit button");
        assert!(prompt.contains("submit button"));
        assert!(prompt.contains('"'));
    }

    #[test]
    fn extract_json_from_plain_text() {
        let s = r#"Some preamble {"key": "value"} trailing"#;
        assert_eq!(extract_json(s), Some(r#"{"key": "value"}"#));
    }

    #[test]
    fn extract_json_from_markdown_fence() {
        let s = "```json\n{\"answer\": 42}\n```";
        assert_eq!(extract_json(s), Some("{\"answer\": 42}"));
    }

    #[test]
    fn extract_json_no_json_returns_none() {
        assert_eq!(extract_json("no json here"), None);
    }
}