searchi/crates/searchi-semantic/src/gemma.rs

use crate::TextWindow;
use anyhow::{ensure, Context, Result};
use base64::Engine;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::{
    collections::BTreeMap,
    fs::File,
    io::{BufRead, BufReader, Read, Write},
    path::{Path, PathBuf},
    process::{Child, ChildStdin, Command, Stdio},
    sync::mpsc::{self, Receiver},
    thread,
    time::{Duration, Instant},
};

pub const GEMMA_DIMENSIONS: usize = 768;
pub const GEMMA_FINGERPRINT: &str = "google/embeddinggemma-2:bfcd298762cc34d0357ece5ebdd31791a3a374d8:q8_0:llama-b11461-no-logits-v1:mean:l2:query512:doc2048:windows-overlap128:prompts-v1";
pub const GEMMA_MEDIA_FINGERPRINT: &str = "google/embeddinggemma-2:bfcd298762cc34d0357ece5ebdd31791a3a374d8:q8_0:llama-b11461-no-logits-v1:mtmd:mean:l2:rgb-catmullrom48-exif-largejpeg-scaled-v2:pcm16k30s-v1";

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageEmbedding {
    pub vector: Vec<f32>,
    pub encoded_width: u32,
    pub encoded_height: u32,
    pub token_budget: u32,
}
const MODEL_FILE: &str = "embeddinggemma-2-Q8_0.gguf";
const MODEL_BYTES: u64 = 309855456;
const MODEL_SHA256: &str = "2188ac1deca4b77dffefd603c2776a9d76d9d74ec01841392982ebb840b09135";
const MAX_REQUEST: usize = 16 * 1024 * 1024;
const MAX_RESPONSE: usize = 4 * 1024 * 1024;

#[derive(Debug, Clone, Copy)]
pub enum Device {
    Auto,
    Cpu,
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct WorkerInfo {
    pub ready: bool,
    pub model: String,
    pub dimensions: usize,
    pub device: String,
    pub threads: usize,
    pub runtime: String,
    #[serde(default)]
    pub media_supported: bool,
}

/// Persistent isolated native worker. No HTTP or remote model fallback.
pub struct GemmaEmbedder {
    child: Child,
    stdin: ChildStdin,
    responses: Receiver<std::result::Result<Value, String>>,
    pub info: WorkerInfo,
    healthy: bool,
    force_cpu: bool,
    model_dir: PathBuf,
    worker_path: PathBuf,
    projector_verified: bool,
    query_profile: Option<BTreeMap<String, u64>>,
}

impl Drop for GemmaEmbedder {
    fn drop(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

impl GemmaEmbedder {
    /// Load the verified Google model and explicitly supplied local worker.
    /// Production selects available acceleration with CPU fallback. Tests set
    /// `SEARCHI_FORCE_CPU=1` or call `load_with_device(..., Device::Cpu)`.
    pub fn load(
        model_dir: impl AsRef<Path>,
        worker: impl AsRef<Path>,
        threads: usize,
    ) -> Result<Self> {
        Self::load_with_device(model_dir, worker, threads, Device::Auto)
    }

    pub fn load_with_device(
        model_dir: impl AsRef<Path>,
        worker: impl AsRef<Path>,
        threads: usize,
        device: Device,
    ) -> Result<Self> {
        ensure!(
            (1..=8).contains(&threads),
            "CPU threads must be between 1 and 8"
        );
        let model_dir = model_dir.as_ref().to_path_buf();
        let model = model_dir.join(MODEL_FILE);
        verify_model(&model)?;
        let worker = worker
            .as_ref()
            .canonicalize()
            .context("Cannot resolve local Gemma worker executable")?;
        ensure!(
            worker.is_file(),
            "Gemma worker must be a local executable file"
        );
        let force_cpu = std::env::var("SEARCHI_FORCE_CPU").is_ok_and(|value| value == "1");
        let cpu = force_cpu || matches!(device, Device::Cpu);
        let mut command = Command::new(&worker);
        command
            .args(["--model"])
            .arg(model)
            .args([
                "--threads",
                &threads.to_string(),
                "--device",
                if cpu { "cpu" } else { "auto" },
            ])
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::inherit());
        // Background helpers must not create a visible console on Windows.
        #[cfg(windows)]
        {
            use std::os::windows::process::CommandExt;
            command.creation_flags(0x08000000); // CREATE_NO_WINDOW
        }
        let mut child = command.spawn().context("Cannot start local Gemma worker")?;
        let stdin = child.stdin.take().context("Worker stdin missing")?;
        let stdout = child.stdout.take().context("Worker stdout missing")?;
        let (sender, responses) = mpsc::sync_channel(1);
        thread::spawn(move || {
            let mut reader = BufReader::new(stdout);
            loop {
                let response = read_response(&mut reader);
                let failed = response.is_err();
                if sender.send(response).is_err() || failed {
                    break;
                }
            }
        });
        let ready = match responses.recv_timeout(Duration::from_secs(120)) {
            Ok(Ok(value)) => value,
            other => {
                let _ = child.kill();
                let _ = child.wait();
                anyhow::bail!("Gemma worker did not initialize: {other:?}");
            }
        };
        let info: WorkerInfo = match serde_json::from_value(ready) {
            Ok(info) => info,
            Err(error) => {
                let _ = child.kill();
                let _ = child.wait();
                return Err(error).context("Invalid Gemma worker handshake");
            }
        };
        if !info.ready
            || info.model != "google/embeddinggemma-2"
            || info.dimensions != GEMMA_DIMENSIONS
            || info.runtime != "b11461-no-logits-v1"
            || info.threads != threads
            || !["cpu", "gpu"].contains(&info.device.as_str())
            || (cpu && info.device != "cpu")
        {
            let _ = child.kill();
            let _ = child.wait();
            anyhow::bail!("Gemma worker handshake violates configured model/device contract");
        }
        Ok(Self {
            child,
            stdin,
            responses,
            info,
            healthy: true,
            force_cpu: cpu,
            model_dir,
            worker_path: worker,
            projector_verified: false,
            query_profile: None,
        })
    }

    fn request(&mut self, request: Value) -> Result<Value> {
        ensure!(
            self.healthy,
            "Gemma worker session failed; reopen it explicitly"
        );
        let encode_started = self.query_profile.as_ref().map(|_| Instant::now());
        let encoded = serde_json::to_vec(&request)?;
        self.record_query_duration("request_encode_us", encode_started);
        ensure!(
            encoded.len() < MAX_REQUEST,
            "Gemma request exceeds 16 MiB protocol budget"
        );
        let response = match self.exchange(&encoded) {
            Ok(response) => response,
            Err(error) if can_retry_cpu(&self.info.device, self.force_cpu) => {
                // GPU runtimes can abort rather than return an encoder error.
                // Restart once as CPU and retry this same bounded request.
                let _ = self.child.wait();
                let restart_started = self.query_profile.as_ref().map(|_| Instant::now());
                let replacement = Self::load_with_device(
                    &self.model_dir,
                    &self.worker_path,
                    self.info.threads,
                    Device::Cpu,
                )
                .with_context(|| {
                    format!("CPU worker restart failed after GPU transport error: {error}")
                })?;
                let diagnostics = self.query_profile.take();
                *self = replacement;
                self.query_profile = diagnostics;
                self.record_query_duration("cpu_restart_us", restart_started);
                if let Some(profile) = &mut self.query_profile {
                    profile.insert("cpu_restart_count".into(), 1);
                }
                self.exchange(&encoded)?
            }
            Err(error) => return Err(error),
        };
        if let Some(device) = response["device"].as_str() {
            ensure!(
                ["cpu", "gpu"].contains(&device) && (!self.force_cpu || device == "cpu"),
                "Worker response violates configured device policy"
            );
            self.info.device = device.into();
        }
        if response["fallback"].as_str() == Some("gpu_inference_failure") {
            if let Some(profile) = &mut self.query_profile {
                profile.insert("worker_cpu_retry_count".into(), 1);
            }
        }
        if response["fatal"] == true {
            self.healthy = false;
            self.info.ready = false;
            let _ = self.child.kill();
        }
        ensure!(
            response["ok"] == true,
            "Gemma worker error: {}",
            response["error"].as_str().unwrap_or("unknown error")
        );
        Ok(response)
    }

    fn exchange(&mut self, encoded: &[u8]) -> Result<Value> {
        let write_started = self.query_profile.as_ref().map(|_| Instant::now());
        let sent = (|| {
            self.stdin.write_all(encoded)?;
            self.stdin.write_all(b"\n")?;
            self.stdin.flush()
        })()
        .context("Cannot write local Gemma worker request");
        self.record_query_duration("request_write_us", write_started);
        if let Err(error) = sent {
            self.healthy = false;
            self.info.ready = false;
            let _ = self.child.kill();
            return Err(error);
        }
        let wait_started = self.query_profile.as_ref().map(|_| Instant::now());
        let response = match self.responses.recv_timeout(Duration::from_secs(180)) {
            Ok(Ok(value)) => value,
            other => {
                self.healthy = false;
                self.info.ready = false;
                let _ = self.child.kill();
                anyhow::bail!("Gemma worker request failed: {other:?}");
            }
        };
        self.record_query_duration("worker_wait_and_response_decode_us", wait_started);
        Ok(response)
    }

    pub fn split_windows(&mut self, text: &str, title: &str) -> Result<Vec<TextWindow>> {
        ensure!(
            !text.trim().is_empty() && text.len() <= 1024 * 1024,
            "Passage must be nonempty and at most 1 MiB"
        );
        let response = self.request(json!({"op":"windows","text":text,"title":title}))?;
        let windows: Vec<TextWindow> = serde_json::from_value(response["windows"].clone())?;
        ensure!(
            !windows.is_empty() && windows.len() <= 256,
            "Invalid Gemma window count"
        );
        for window in &windows {
            ensure!(
                text.get(window.start_byte..window.end_byte) == Some(window.text.as_str()),
                "Gemma window does not match original source range"
            );
        }
        ensure!(
            windows
                .last()
                .is_some_and(|window| window.end_byte == text.len()),
            "Gemma windows did not cover the passage tail"
        );
        Ok(windows)
    }

    pub fn embed_documents(&mut self, texts: &[&str], title: &str) -> Result<Vec<Vec<f32>>> {
        ensure!(
            !texts.is_empty() && texts.len() <= 8,
            "Document batch size must be between 1 and 8"
        );
        ensure!(
            texts
                .iter()
                .all(|text| !text.trim().is_empty() && text.len() <= 65536),
            "Document window must be nonempty and at most 64 KiB"
        );
        let response = self.request(json!({"op":"documents","texts":texts,"title":title}))?;
        extract_vectors(response, texts.len())
    }

    pub fn embed_query(&mut self, query: &str) -> Result<Vec<f32>> {
        self.query_profile = None;
        self.embed_query_inner(query, false)
    }

    /// Opt-in diagnostics contain only durations and counts, never query text.
    pub fn embed_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
        self.query_profile = Some(BTreeMap::new());
        self.embed_query_inner(query, true)
    }

    pub fn query_diagnostics(&self) -> BTreeMap<String, u64> {
        self.query_profile.clone().unwrap_or_default()
    }

    fn record_query_duration(&mut self, name: &str, started: Option<Instant>) {
        if let (Some(profile), Some(started)) = (&mut self.query_profile, started) {
            let micros = started.elapsed().as_micros().min(u64::MAX as u128) as u64;
            *profile.entry(name.into()).or_default() += micros;
        }
    }

    fn embed_query_inner(&mut self, query: &str, debug: bool) -> Result<Vec<f32>> {
        ensure!(
            !query.trim().is_empty() && query.len() <= 4096,
            "Query must be nonempty and at most 4096 UTF-8 bytes"
        );
        let response =
            self.request(json!({"op":"query","text":query,"debug_performance":debug}))?;
        if debug {
            if let Some(profile) = &mut self.query_profile {
                for name in [
                    "tokenization_us",
                    "batch_prepare_us",
                    "inference_us",
                    "pool_normalize_us",
                    "token_count",
                ] {
                    if let Some(value) = response["performance"][name].as_u64() {
                        profile.insert(format!("worker_{name}"), value);
                    }
                }
            }
        }
        let validation_started = self.query_profile.as_ref().map(|_| Instant::now());
        let result = extract_vectors(response, 1)?.remove(0);
        self.record_query_duration("response_validate_us", validation_started);
        Ok(result)
    }

    fn projector(&mut self) -> Result<PathBuf> {
        ensure!(
            self.info.media_supported,
            "Installed embedding worker does not support media; repair/rebuild the native runtime"
        );
        let path = self.model_dir.join("mmproj-embeddinggemma-2-Q8_0.gguf");
        if !self.projector_verified {
            verify_artifact(
                &path,
                554821024,
                "c4a8a52691ecef40618438928bdf9e68379b854e24166f292592353db0aab64f",
            )?;
            self.projector_verified = true;
        }
        Ok(path)
    }

    /// Genuine Google pixel embeddings, serial bounded image operations.
    pub fn embed_images(&mut self, paths: &[&Path]) -> Result<Vec<Vec<f32>>> {
        self.embed_images_with_budget(paths, 280)
    }

    /// 70 tokens for video frames, 140 for photos, 280 or 1120 for screenshots.
    pub fn embed_images_with_budget(
        &mut self,
        paths: &[&Path],
        budget: u32,
    ) -> Result<Vec<Vec<f32>>> {
        Ok(self
            .embed_images_with_metadata(paths, budget)?
            .into_iter()
            .map(|image| image.vector)
            .collect())
    }

    pub fn embed_images_with_metadata(
        &mut self,
        paths: &[&Path],
        budget: u32,
    ) -> Result<Vec<ImageEmbedding>> {
        ensure!(
            !paths.is_empty() && paths.len() <= 8,
            "Image batch must contain 1–8 paths"
        );
        let projector = self.projector()?;
        let mut vectors = Vec::with_capacity(paths.len());
        for path in paths {
            let pixels = decode_image(path, budget)?;
            let response=self.request(json!({"op":"image","projector":projector,"width":pixels.width(),"height":pixels.height(),"data":base64::engine::general_purpose::STANDARD.encode(pixels.as_raw())}))?;
            vectors.push(ImageEmbedding {
                vector: extract_vectors(response, 1)?.remove(0),
                encoded_width: pixels.width(),
                encoded_height: pixels.height(),
                token_budget: budget,
            });
        }
        Ok(vectors)
    }

    /// One mono 16 kHz normalized PCM window, at most 30 seconds. Callers
    /// segment entire recordings and preserve real timestamps; no truncation.
    pub fn embed_audio_pcm(&mut self, samples: &[f32]) -> Result<Vec<f32>> {
        ensure!(
            !samples.is_empty() && samples.len() <= 480000,
            "Audio window must contain 1–480000 mono 16 kHz samples"
        );
        ensure!(
            samples.iter().all(|v| v.is_finite() && v.abs() <= 1.0),
            "Invalid normalized audio PCM"
        );
        ensure!(
            cfg!(target_endian = "little"),
            "Native PCM worker requires little-endian host"
        );
        let projector = self.projector()?;
        let bytes: Vec<u8> = samples.iter().flat_map(|v| v.to_le_bytes()).collect();
        let response=self.request(json!({"op":"audio","projector":projector,"data":base64::engine::general_purpose::STANDARD.encode(bytes)}))?;
        Ok(extract_vectors(response, 1)?.remove(0))
    }
}

fn decode_image(path: &Path, budget: u32) -> Result<image::RgbImage> {
    ensure!(
        (70..=1120).contains(&budget),
        "Image token budget must be 70–1120"
    );
    ensure!(
        std::fs::metadata(path)?.len() <= 64 * 1024 * 1024,
        "Image file exceeds 64 MiB"
    );
    let mut reader = image::ImageReader::open(path)?.with_guessed_format()?;
    let is_jpeg = reader.format() == Some(image::ImageFormat::Jpeg);
    let mut limits = image::Limits::default();
    limits.max_alloc = Some(256 * 1024 * 1024);
    limits.max_image_width = Some(16384);
    limits.max_image_height = Some(16384);
    reader.limits(limits);
    use image::ImageDecoder;
    let mut decoder = reader
        .into_decoder()
        .context("Cannot initialize bounded image decoder")?;
    let (source_width, source_height) = decoder.dimensions();
    ensure!(
        u64::from(source_width) * u64::from(source_height) <= 64 * 1024 * 1024
            && decoder.total_bytes() <= 256 * 1024 * 1024,
        "Decoded image exceeds memory/pixel budget"
    );
    let orientation = decoder
        .orientation()
        .context("Cannot read image EXIF orientation")?;
    if is_jpeg && u64::from(source_width) * u64::from(source_height) > 16 * 1024 * 1024 {
        drop(decoder);
        return decode_large_jpeg(path, budget, source_width, source_height, orientation);
    }
    ensure!(u64::from(source_width)*u64::from(source_height)<=16*1024*1024,"Non-JPEG image exceeds16megapixels; bounded scaled decoding is currently available for JPEG");
    let mut decoded = image::DynamicImage::from_decoder(decoder)
        .context("Cannot decode supported image pixels")?;
    decoded.apply_orientation(orientation);
    ensure!(
        u64::from(decoded.width()) * u64::from(decoded.height()) <= 16 * 1024 * 1024,
        "Image exceeds 16 megapixels"
    );
    let (w, h) = image_target(decoded.width(), decoded.height(), budget)?;
    Ok(image::imageops::resize(
        &decoded.to_rgb8(),
        w,
        h,
        image::imageops::FilterType::CatmullRom,
    ))
}

fn decode_large_jpeg(
    path: &Path,
    budget: u32,
    width: u32,
    height: u32,
    orientation: image::metadata::Orientation,
) -> Result<image::RgbImage> {
    let (target_width, target_height) = image_target(width, height, budget)?;
    let mut decoder = jpeg_decoder::Decoder::new(BufReader::new(File::open(path)?));
    decoder.set_max_decoding_buffer_size(256 * 1024 * 1024);
    decoder
        .read_info()
        .context("Cannot read scaled JPEG header")?;
    let source = decoder.info().context("Missing scaled JPEG header")?;
    ensure!(
        u32::from(source.width) == width && u32::from(source.height) == height,
        "JPEG changed during decoding"
    );
    ensure!(
        matches!(
            source.pixel_format,
            jpeg_decoder::PixelFormat::RGB24 | jpeg_decoder::PixelFormat::L8
        ),
        "Large JPEG requires supported RGB24 or grayscale scaled decoding"
    );
    let (w, h) = decoder.scale(u16::try_from(target_width)?, u16::try_from(target_height)?)?;
    ensure!(
        u64::from(w) * u64::from(h) * 4 <= 256 * 1024 * 1024,
        "Scaled JPEG output exceeds memory budget"
    );
    let pixels = decoder
        .decode()
        .context("Cannot decode large JPEG with scaled IDCT")?;
    let mut decoded = match source.pixel_format {
        jpeg_decoder::PixelFormat::RGB24 => image::DynamicImage::ImageRgb8(
            image::RgbImage::from_raw(u32::from(w), u32::from(h), pixels)
                .context("Invalid scaled JPEG RGB length")?,
        ),
        jpeg_decoder::PixelFormat::L8 => image::DynamicImage::ImageLuma8(
            image::GrayImage::from_raw(u32::from(w), u32::from(h), pixels)
                .context("Invalid scaled JPEG grayscale length")?,
        ),
        other => anyhow::bail!(
            "Large JPEG pixel format {other:?} is not supported by bounded scaled decoder"
        ),
    };
    decoded.apply_orientation(orientation);
    let (w, h) = image_target(decoded.width(), decoded.height(), budget)?;
    Ok(image::imageops::resize(
        &decoded.into_rgb8(),
        w,
        h,
        image::imageops::FilterType::CatmullRom,
    ))
}

fn image_target(width: u32, height: u32, budget: u32) -> Result<(u32, u32)> {
    ensure!(
        width > 0 && height > 0 && (70..=1120).contains(&budget),
        "Invalid image dimensions/budget"
    );
    let scale = (f64::from(budget) * 2304.0 / (f64::from(width) * f64::from(height))).sqrt();
    let mut w = ((f64::from(width) * scale / 48.0).floor() as u32) * 48;
    let mut h = ((f64::from(height) * scale / 48.0).floor() as u32) * 48;
    if w == 0 {
        w = 48;
        h = (height / width * 48).min(budget * 48);
    }
    if h == 0 {
        h = 48;
        w = (width / height * 48).min(budget * 48);
    }
    ensure!(
        w >= 48 && h >= 48 && u64::from(w) * u64::from(h) <= u64::from(budget) * 2304,
        "Image target outside token budget"
    );
    Ok((w, h))
}

fn can_retry_cpu(actual_device: &str, force_cpu: bool) -> bool {
    actual_device == "gpu" && !force_cpu
}

fn read_response(reader: &mut impl BufRead) -> std::result::Result<Value, String> {
    let mut line = Vec::new();
    match reader
        .take((MAX_RESPONSE + 1) as u64)
        .read_until(b'\n', &mut line)
    {
        Ok(0) => Err("Gemma worker closed its output".into()),
        Ok(_) if line.len() > MAX_RESPONSE || line.last() != Some(&b'\n') => {
            Err("Gemma worker response exceeds protocol byte budget".into())
        }
        Ok(_) => serde_json::from_slice(&line)
            .map_err(|error| format!("Invalid Gemma worker JSON: {error}")),
        Err(error) => Err(format!("Gemma worker read failed: {error}")),
    }
}

fn extract_vectors(response: Value, count: usize) -> Result<Vec<Vec<f32>>> {
    let vectors: Vec<Vec<f32>> = serde_json::from_value(response["vectors"].clone())?;
    ensure!(
        vectors.len() == count,
        "Gemma returned incorrect embedding batch length"
    );
    for vector in &vectors {
        ensure!(
            vector.len() == GEMMA_DIMENSIONS && vector.iter().all(|value| value.is_finite()),
            "Invalid Gemma embedding shape or values"
        );
        let norm = vector.iter().map(|v| v * v).sum::<f32>().sqrt();
        ensure!(
            (norm - 1.0).abs() < 1e-4,
            "Gemma embedding is not L2 normalized"
        );
    }
    Ok(vectors)
}

fn verify_model(path: &Path) -> Result<()> {
    verify_artifact(path, MODEL_BYTES, MODEL_SHA256)
}

fn verify_artifact(path: &Path, bytes: u64, sha: &str) -> Result<()> {
    let file = File::open(path).with_context(|| {
        format!(
            "Missing Google EmbeddingGemma 2 model {}; run explicit model setup first",
            path.display()
        )
    })?;
    ensure!(
        file.metadata()?.len() == bytes,
        "EmbeddingGemma 2 model size mismatch"
    );
    let mut reader = BufReader::new(file);
    let mut hash = Sha256::new();
    let mut buffer = [0; 65536];
    loop {
        let count = reader.read(&mut buffer)?;
        if count == 0 {
            break;
        }
        hash.update(&buffer[..count]);
    }
    ensure!(
        format!("{:x}", hash.finalize()) == sha,
        "EmbeddingGemma 2 model SHA256 mismatch"
    );
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn old_text_worker_handshake_does_not_claim_media_support() {
        let mut handshake = json!({"ready":true,"model":"google/embeddinggemma-2","dimensions":768,"device":"cpu","threads":1,"runtime":"b11461-no-logits-v1"});
        let old: WorkerInfo = serde_json::from_value(handshake.clone()).unwrap();
        assert!(!old.media_supported);
        handshake["media_supported"] = json!(true);
        assert!(
            serde_json::from_value::<WorkerInfo>(handshake)
                .unwrap()
                .media_supported
        );
    }
    #[test]
    fn media_image_grid_is_bounded_for_extreme_aspects_and_full_budget() {
        for budget in [70, 140, 280, 1120] {
            for (width, height) in [(1, 16384), (16384, 1), (1920, 1080), (500, 500)] {
                let (w, h) = image_target(width, height, budget).unwrap();
                assert_eq!(w % 48, 0);
                assert_eq!(h % 48, 0);
                assert!(u64::from(w) * u64::from(h) <= u64::from(budget) * 2304);
            }
        }
        assert!(image_target(0, 100, 280).is_err());
        assert!(image_target(100, 100, 1121).is_err());
    }
    #[test]
    fn artifact_verification_rejects_projector_size_and_checksum_mismatch() {
        let path = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/LICENSE-DigUp.txt");
        let bytes = std::fs::read(&path).unwrap();
        let sha = format!("{:x}", Sha256::digest(&bytes));
        assert!(verify_artifact(&path, bytes.len() as u64, &sha).is_ok());
        assert!(verify_artifact(&path, bytes.len() as u64 + 1, &sha).is_err());
        assert!(verify_artifact(&path, bytes.len() as u64, &"0".repeat(64)).is_err());
    }
    #[test]
    #[ignore = "Explicit48MP JPEG resource test; no model needed, run release CPU"]
    fn large_phone_jpeg_scaled_decode_preserves_exif_and_bounds() {
        use image::ImageEncoder;
        let image = image::RgbImage::from_fn(8000, 6000, |x, _| {
            if x < 4000 {
                image::Rgb([240, 10, 10])
            } else {
                image::Rgb([10, 10, 240])
            }
        });
        let mut encoded = Vec::new();
        let mut encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut encoded, 80);
        // Little-endian TIFF orientation tag6: rotate90clockwise.
        encoder
            .set_exif_metadata(vec![
                b'I', b'I', 42, 0, 8, 0, 0, 0, 1, 0, 0x12, 1, 3, 0, 1, 0, 0, 0, 6, 0, 0, 0, 0, 0,
                0, 0,
            ])
            .unwrap();
        encoder.encode_image(&image).unwrap();
        drop(image);
        let path =
            std::env::temp_dir().join(format!("searchi-phone48mp-{}.jpg", std::process::id()));
        std::fs::write(&path, &encoded).unwrap();
        drop(encoded);
        let started = std::time::Instant::now();
        let result = decode_image(&path, 280);
        let _ = std::fs::remove_file(&path);
        let pixels = result.unwrap();
        println!(
            "48MPactualJPEG scaleddecode+EXIF+resize {:?}; output{}x{} RGB{}bytes",
            started.elapsed(),
            pixels.width(),
            pixels.height(),
            pixels.as_raw().len()
        );
        assert!(pixels.height() > pixels.width());
        assert!(pixels.get_pixel(pixels.width() / 2, pixels.height() / 4)[0] > 200);
        assert!(pixels.get_pixel(pixels.width() / 2, pixels.height() * 3 / 4)[2] > 200);
        assert!(u64::from(pixels.width()) * u64::from(pixels.height()) <= 280 * 2304);
    }
    #[test]
    fn oversized_jpeg_header_rejected_before_pixel_decode() {
        let mut encoded = Vec::new();
        image::codecs::jpeg::JpegEncoder::new(&mut encoded)
            .encode_image(&image::RgbImage::new(16, 16))
            .unwrap();
        let sof = encoded.windows(2).position(|p| p == [0xff, 0xc0]).unwrap();
        encoded[sof + 5..sof + 7].copy_from_slice(&9000u16.to_be_bytes());
        encoded[sof + 7..sof + 9].copy_from_slice(&9000u16.to_be_bytes());
        let path =
            std::env::temp_dir().join(format!("searchi-too-large-{}.jpg", std::process::id()));
        std::fs::write(&path, encoded).unwrap();
        let result = decode_image(&path, 280);
        let _ = std::fs::remove_file(&path);
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("memory/pixel budget"));
    }
    #[test]
    fn missing_google_model_fails_without_starting_worker() {
        assert!(GemmaEmbedder::load("missing-google-model", "missing-worker", 1).is_err());
    }

    #[test]
    fn retry_policy_never_enables_gpu_or_retries_cpu_as_gpu() {
        assert!(can_retry_cpu("gpu", false));
        assert!(!can_retry_cpu("gpu", true));
        assert!(!can_retry_cpu("cpu", false));
        assert!(!can_retry_cpu("cpu", true));
    }

    #[test]
    fn protocol_rejects_oversized_truncated_and_malformed_frames() {
        assert!(read_response(&mut std::io::Cursor::new(vec![b'x'; MAX_RESPONSE + 1])).is_err());
        assert!(read_response(&mut std::io::Cursor::new(b"{\"ok\":true}")).is_err());
        assert!(read_response(&mut std::io::Cursor::new(b"not-json\n")).is_err());
        assert_eq!(
            read_response(&mut std::io::Cursor::new(b"{\"ok\":true}\n")).unwrap()["ok"],
            true
        );
    }

    #[test]
    fn protocol_rejects_wrong_embedding_shape_count_and_normalization() {
        assert!(extract_vectors(json!({"vectors":[]}), 1).is_err());
        assert!(extract_vectors(json!({"vectors":[[1.0,0.0]]}), 1).is_err());
        assert!(extract_vectors(json!({"vectors":[vec![0.0;GEMMA_DIMENSIONS]]}), 1).is_err());
        let mut invalid = vec![0.0; GEMMA_DIMENSIONS];
        invalid[0] = 1.0;
        invalid[1] = 1.0;
        assert!(extract_vectors(json!({"vectors":[invalid]}), 1).is_err());
    }
}