codecraft/server/src/tools/craft.rs

//! Craft mode tools: study a target in a headless browser and compare a build
//! with the original visually.

use super::{Tool, ToolCtx, ToolOutput, arg_str};
use crate::types::ImageData;
use anyhow::{Context, Result, bail};
use async_trait::async_trait;
use base64::Engine;
use image::{GenericImageView, ImageFormat, Rgba, RgbaImage};
use serde_json::{Value, json};
use std::io::Cursor;
use std::path::{Path, PathBuf};
use std::time::Duration;
use tokio::process::Command;

const SNAPSHOT_HEIGHT: u32 = 2000;
/// Long-side limit for images sent to the model.
const MAX_IMAGE_SIDE: u32 = 1568;

fn png_for_model(img: &image::DynamicImage) -> Result<ImageData> {
    let (w, h) = img.dimensions();
    let img = if w.max(h) > MAX_IMAGE_SIDE {
        img.resize(MAX_IMAGE_SIDE, MAX_IMAGE_SIDE, image::imageops::FilterType::Triangle)
    } else {
        img.clone()
    };
    let mut buf = Vec::new();
    img.write_to(&mut Cursor::new(&mut buf), ImageFormat::Png)?;
    Ok(ImageData {
        media_type: "image/png".into(),
        data: base64::engine::general_purpose::STANDARD.encode(&buf),
    })
}

fn slug(s: &str) -> String {
    let mut out: String = s
        .chars()
        .map(|c| if c.is_ascii_alphanumeric() { c.to_ascii_lowercase() } else { '-' })
        .collect();
    while out.contains("--") {
        out = out.replace("--", "-");
    }
    out.trim_matches('-').chars().take(60).collect()
}

async fn chromium(ctx: &ToolCtx, args: &[String]) -> Result<std::process::Output> {
    let bin = ctx.app.config().craft.chromium.clone();
    let fut = Command::new(&bin)
        .args([
            "--headless=new",
            "--no-sandbox",
            "--disable-gpu",
            "--hide-scrollbars",
            "--mute-audio",
            "--disable-dev-shm-usage",
            "--virtual-time-budget=8000",
        ])
        .args(args)
        .kill_on_drop(true)
        .output();
    tokio::select! {
        r = tokio::time::timeout(Duration::from_secs(60), fut) => {
            r.context("chromium timed out")?.with_context(|| format!("could not start '{bin}' (set craft.chromium in the config)"))
        }
        _ = ctx.cancel.cancelled() => bail!("aborted"),
    }
}

pub struct BrowserSnapshot;

#[async_trait]
impl Tool for BrowserSnapshot {
    fn name(&self) -> &str {
        "browser_snapshot"
    }
    fn description(&self) -> String {
        "Open a URL in a headless browser and capture what a visitor sees: screenshots at the \
         configured screen widths (or the given ones) plus the visible text. Works for public \
         sites and for local builds (http://localhost:PORT or file:///path). Screenshots are saved \
         under .craft/snapshots/<label>-<width>.png in the working directory for image_compare."
            .into()
    }
    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "url": {"type": "string"},
                "label": {"type": "string", "description": "Name for the saved files, e.g. 'original-home' or 'build-home'"},
                "widths": {"type": "array", "items": {"type": "integer"}},
                "text": {"type": "boolean", "description": "Include visible text (default true)"}
            },
            "required": ["url"]
        })
    }
    fn read_only(&self) -> bool {
        true
    }
    fn title(&self, input: &Value) -> String {
        format!("Snapshot {}", input["url"].as_str().unwrap_or("?"))
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let url = arg_str(&input, "url")?.to_string();
        // Local files only from the user's own folders.
        if let Some(path) = url.strip_prefix("file://") {
            ctx.resolve(path)?;
        } else if !(url.starts_with("http://") || url.starts_with("https://")) {
            bail!("only http(s) and file:// URLs are supported");
        }
        let label = slug(input["label"].as_str().unwrap_or(&url));
        let widths: Vec<u32> = input["widths"]
            .as_array()
            .map(|a| a.iter().filter_map(|w| w.as_u64().map(|w| w as u32)).collect())
            .filter(|v: &Vec<u32>| !v.is_empty())
            .unwrap_or_else(|| ctx.app.config().craft.widths.clone());
        let dir = ctx.cwd.join(".craft/snapshots");
        tokio::fs::create_dir_all(&dir).await?;

        let mut images = Vec::new();
        let mut files = Vec::new();
        for w in widths.iter().take(4) {
            let path = dir.join(format!("{label}-{w}.png"));
            let out = chromium(
                ctx,
                &[
                    format!("--window-size={w},{SNAPSHOT_HEIGHT}"),
                    format!("--screenshot={}", path.display()),
                    url.clone(),
                ],
            )
            .await?;
            if !path.exists() {
                bail!("screenshot failed: {}", String::from_utf8_lossy(&out.stderr).chars().take(800).collect::<String>());
            }
            let img = image::open(&path)?;
            images.push(png_for_model(&img)?);
            files.push(path.strip_prefix(&ctx.cwd).unwrap_or(&path).display().to_string());
        }

        let mut text = format!("Snapshots of {url}:\n");
        for (f, w) in files.iter().zip(&widths) {
            text.push_str(&format!("- {w}px → {f}\n"));
        }
        if input["text"].as_bool().unwrap_or(true) {
            let out = chromium(ctx, &["--dump-dom".into(), url.clone()]).await?;
            let dom = String::from_utf8_lossy(&out.stdout);
            let visible = html2text::from_read(dom.as_bytes(), 120).unwrap_or_default();
            text.push_str("\nVisible text:\n");
            text.push_str(&super::truncate(&visible, 20_000));
        }
        Ok(ToolOutput { content: text, images, meta: Some(json!({"files": files})), ..Default::default() })
    }
}

// ---------------------------------------------------------------------------

pub struct ImageCompare;

/// Similarity in [0, 1] and a diff image (differences in red over a faded original).
pub fn compare(a: &image::DynamicImage, b: &image::DynamicImage) -> (f64, RgbaImage) {
    let (w, h) = a.dimensions();
    let b = if b.dimensions() != (w, h) {
        b.resize_exact(w, h, image::imageops::FilterType::Triangle)
    } else {
        b.clone()
    };
    let a = a.to_rgba8();
    let b = b.to_rgba8();
    let mut diff = RgbaImage::new(w, h);
    let mut total = 0f64;
    for (x, y, pa) in a.enumerate_pixels() {
        let pb = b.get_pixel(x, y);
        let d = (0..3).map(|i| (pa[i] as f64 - pb[i] as f64).abs()).sum::<f64>() / (3.0 * 255.0);
        total += d;
        let gray = ((pa[0] as u32 + pa[1] as u32 + pa[2] as u32) / 3) as u8;
        let faded = 200 + gray / 5;
        diff.put_pixel(
            x,
            y,
            if d > 0.1 { Rgba([255, (255.0 * (1.0 - d)) as u8 / 2, 0, 255]) } else { Rgba([faded, faded, faded, 255]) },
        );
    }
    let px = (w as f64 * h as f64).max(1.0);
    (1.0 - total / px, diff)
}

#[async_trait]
impl Tool for ImageCompare {
    fn name(&self) -> &str {
        "image_compare"
    }
    fn description(&self) -> String {
        "Compare two images (e.g. the original's and the build's snapshot at the same width). \
         Returns a similarity score (0-100 %) and a diff image where differences are red."
            .into()
    }
    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "original": {"type": "string", "description": "Path of the reference image"},
                "candidate": {"type": "string", "description": "Path of the image to check"}
            },
            "required": ["original", "candidate"]
        })
    }
    fn read_only(&self) -> bool {
        true
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let pa: PathBuf = ctx.resolve(arg_str(&input, "original")?)?;
        let pb: PathBuf = ctx.resolve(arg_str(&input, "candidate")?)?;
        let (score, diff) = tokio::task::spawn_blocking(move || -> Result<_> {
            let a = image::open(&pa).with_context(|| format!("cannot open {}", pa.display()))?;
            let b = image::open(&pb).with_context(|| format!("cannot open {}", pb.display()))?;
            Ok(compare(&a, &b))
        })
        .await??;
        let diff_path = ctx.cwd.join(".craft/snapshots").join(format!(
            "diff-{}.png",
            Path::new(input["candidate"].as_str().unwrap_or("x")).file_stem().and_then(|s| s.to_str()).unwrap_or("x")
        ));
        if let Some(d) = diff_path.parent() {
            tokio::fs::create_dir_all(d).await?;
        }
        diff.save(&diff_path)?;
        let img = image::DynamicImage::ImageRgba8(diff);
        let pct = (score * 1000.0).round() / 10.0;
        Ok(ToolOutput {
            content: format!("Similarity: {pct} %\nDiff image: {}", diff_path.display()),
            images: vec![png_for_model(&img)?],
            meta: Some(json!({"similarity": pct})),
            ..Default::default()
        })
    }
}

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

    #[test]
    fn identical_images_are_fully_similar() {
        let img = image::DynamicImage::ImageRgba8(RgbaImage::from_pixel(10, 10, Rgba([10, 20, 30, 255])));
        let (s, _) = compare(&img, &img);
        assert!((s - 1.0).abs() < 1e-9);
        let other = image::DynamicImage::ImageRgba8(RgbaImage::from_pixel(10, 10, Rgba([255, 255, 255, 255])));
        let (s2, _) = compare(&img, &other);
        assert!(s2 < 0.2);
    }
}