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,
    #[serde(default)]
    pub code_query_supported: bool,
}

fn query_operation(info: &WorkerInfo, code: bool) -> Result<&'static str> {
    ensure!(!code || info.code_query_supported,"Installed embedding worker does not support code retrieval; repair/rebuild the native runtime");
    Ok(if code { "code_query" } else { "query" })
}

/// Persistent isolated native worker. No HTTP or remote model fallback.
pub struct GemmaEmbedder {
    child: Child,
    requests: Option<mpsc::SyncSender<Vec<u8>>>,
    writes: Receiver<std::result::Result<(), String>>,
    writer_thread: Option<thread::JoinHandle<()>>,
    reader_thread: Option<thread::JoinHandle<()>>,
    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) {
        self.stop_worker();
    }
}

type WriterTransport = (
    mpsc::SyncSender<Vec<u8>>,
    Receiver<std::result::Result<(), String>>,
    thread::JoinHandle<()>,
);

fn spawn_writer(mut stdin: ChildStdin) -> WriterTransport {
    let (requests, pending) = mpsc::sync_channel::<Vec<u8>>(1);
    let (completed, writes) = mpsc::sync_channel(1);
    let writer = thread::spawn(move || {
        while let Ok(encoded) = pending.recv() {
            let result = stdin
                .write_all(&encoded)
                .and_then(|_| stdin.write_all(b"\n"))
                .and_then(|_| stdin.flush())
                .map_err(|error| format!("Cannot write local Gemma worker request: {error}"));
            let failed = result.is_err();
            if completed.try_send(result).is_err() || failed {
                break;
            }
        }
    });
    (requests, writes, writer)
}

#[derive(Debug)]
struct WorkerTransportFailure(String);
impl std::fmt::Display for WorkerTransportFailure {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(&self.0)
    }
}
impl std::error::Error for WorkerTransportFailure {}

/// Retry only process/transport failure, never an incompatible protocol or model.
fn start_with_cpu_fallback<T>(cpu: bool, mut start: impl FnMut(bool) -> Result<T>) -> Result<T> {
    match start(cpu) {
        Err(error) if !cpu && error.downcast_ref::<WorkerTransportFailure>().is_some() => {
            start(true).with_context(|| {
                format!("CPU initialization failed after automatic worker startup failure: {error}")
            })
        }
        result => result,
    }
}

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);
        start_with_cpu_fallback(cpu, |cpu| {
            Self::start_verified(&model_dir, &worker, threads, cpu)
        })
    }

    // Both attempts share the already validated model/worker paths. CPU fallback
    // does not repeat the expensive hash check or weaken the handshake contract.
    fn start_verified(model_dir: &Path, worker: &Path, threads: usize, cpu: bool) -> Result<Self> {
        let model = model_dir.join(MODEL_FILE);
        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().map_err(|error| {
            WorkerTransportFailure(format!("Cannot start local Gemma worker: {error}"))
        })?;
        let (Some(stdin), Some(stdout)) = (child.stdin.take(), child.stdout.take()) else {
            let _ = child.kill();
            let _ = child.wait();
            anyhow::bail!("Gemma worker protocol pipes missing");
        };
        let (sender, responses) = mpsc::sync_channel(1);
        let reader_thread = thread::spawn(move || {
            let mut reader = BufReader::new(stdout);
            loop {
                let response = read_response(&mut reader);
                let failed = response.is_err();
                if sender.try_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();
                let _ = reader_thread.join();
                // Malformed/oversized JSON is a protocol violation, not evidence
                // that hardware acceleration failed. Do not mask it with retry.
                if matches!(&other, Ok(Err(error)) if error.starts_with("Invalid Gemma worker JSON:") || error == "Gemma worker response exceeds protocol byte budget")
                {
                    anyhow::bail!("Invalid Gemma worker startup protocol: {other:?}");
                }
                return Err(WorkerTransportFailure(format!(
                    "Gemma worker did not initialize: {other:?}"
                ))
                .into());
            }
        };
        let info: WorkerInfo = match serde_json::from_value(ready) {
            Ok(info) => info,
            Err(error) => {
                let _ = child.kill();
                let _ = child.wait();
                let _ = reader_thread.join();
                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();
            let _ = reader_thread.join();
            anyhow::bail!("Gemma worker handshake violates configured model/device contract");
        }
        let (requests, writes, writer_thread) = spawn_writer(stdin);
        Ok(Self {
            child,
            requests: Some(requests),
            writes,
            writer_thread: Some(writer_thread),
            reader_thread: Some(reader_thread),
            responses,
            info,
            healthy: true,
            force_cpu: cpu,
            model_dir: model_dir.to_path_buf(),
            worker_path: worker.to_path_buf(),
            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)
                    && error.downcast_ref::<WorkerTransportFailure>().is_some() =>
            {
                // 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 stop_worker(&mut self) {
        self.requests.take();
        // The shipped native worker does not spawn descendants. Killing/reaping
        // it closes its pipe ends and releases any blocked writer/read thread.
        let _ = self.child.kill();
        let _ = self.child.wait();
        if let Some(writer) = self.writer_thread.take() {
            let _ = writer.join();
        }
        if let Some(reader) = self.reader_thread.take() {
            let _ = reader.join();
        }
    }

    fn exchange(&mut self, encoded: &[u8]) -> Result<Value> {
        self.exchange_with_timeout(encoded, Duration::from_secs(180))
    }

    fn exchange_with_timeout(&mut self, encoded: &[u8], timeout: Duration) -> Result<Value> {
        let deadline = Instant::now() + timeout;
        let write_started = self.query_profile.as_ref().map(|_| Instant::now());
        let result = (|| {
            self.requests
                .as_ref()
                .ok_or_else(|| WorkerTransportFailure("Gemma worker writer stopped".into()))?
                .try_send(encoded.to_vec())
                .map_err(|_| {
                    WorkerTransportFailure("Gemma worker request queue unavailable".into())
                })?;
            match self
                .writes
                .recv_timeout(deadline.saturating_duration_since(Instant::now()))
            {
                Ok(Ok(())) => (),
                other => {
                    return Err(WorkerTransportFailure(format!(
                        "Gemma worker request write failed or timed out: {other:?}"
                    ))
                    .into())
                }
            }
            self.record_query_duration("request_write_us", write_started);
            let wait_started = self.query_profile.as_ref().map(|_| Instant::now());
            let response = match self
                .responses
                .recv_timeout(deadline.saturating_duration_since(Instant::now()))
            {
                Ok(Ok(value)) => value,
                Ok(Err(error))
                    if error.starts_with("Invalid Gemma worker JSON:")
                        || error == "Gemma worker response exceeds protocol byte budget" =>
                {
                    anyhow::bail!("Invalid Gemma worker response protocol: {error}")
                }
                other => {
                    return Err(WorkerTransportFailure(format!(
                        "Gemma worker request failed or timed out: {other:?}"
                    ))
                    .into())
                }
            };
            self.record_query_duration("worker_wait_and_response_decode_us", wait_started);
            Ok(response)
        })();
        if result.is_err() {
            self.healthy = false;
            self.info.ready = false;
            self.stop_worker();
        }
        result
    }

    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, 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, false)
    }

    /// Google code retrieval task; unavailable workers fail explicitly.
    pub fn embed_code_query(&mut self, query: &str) -> Result<Vec<f32>> {
        self.query_profile = None;
        self.embed_query_inner(query, false, true)
    }
    pub fn embed_code_query_debug(&mut self, query: &str) -> Result<Vec<f32>> {
        self.query_profile = Some(BTreeMap::new());
        self.embed_query_inner(query, true, 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, code: bool) -> Result<Vec<f32>> {
        ensure!(
            !query.trim().is_empty() && query.len() <= 4096,
            "Query must be nonempty and at most 4096 UTF-8 bytes"
        );
        let op = query_operation(&self.info, code)?;
        let response = self.request(json!({"op":op,"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::*;
    // This helper is an inert pipe-protocol fixture, never an embedding provider.
    fn stalled_protocol_helper(read_request: bool) -> GemmaEmbedder {
        #[cfg(windows)]
        let mut command = {
            use std::os::windows::process::CommandExt;
            let mut command = Command::new("powershell.exe");
            command.args([
                "-NoProfile",
                "-NonInteractive",
                "-Command",
                if read_request {
                    "[Console]::In.ReadLine() | Out-Null; Start-Sleep -Seconds 20"
                } else {
                    "Start-Sleep -Seconds 20"
                },
            ]);
            command.creation_flags(0x08000000);
            command
        };
        #[cfg(not(windows))]
        let mut command = {
            let mut command = Command::new("sh");
            command.args([
                "-c",
                if read_request {
                    "read line; exec sleep 20"
                } else {
                    "exec sleep 20"
                },
            ]);
            command
        };
        let mut child = command
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null())
            .spawn()
            .unwrap();
        let (requests, writes, writer_thread) = spawn_writer(child.stdin.take().unwrap());
        let stdout = child.stdout.take().unwrap();
        let (sender, responses) = mpsc::sync_channel(1);
        let reader_thread = thread::spawn(move || {
            let response = read_response(&mut BufReader::new(stdout));
            let _ = sender.try_send(response);
        });
        GemmaEmbedder {
            child,
            requests: Some(requests),
            writes,
            writer_thread: Some(writer_thread),
            reader_thread: Some(reader_thread),
            responses,
            info: WorkerInfo {
                ready: true,
                model: "pipe protocol fixture".into(),
                dimensions: 0,
                device: "cpu".into(),
                threads: 1,
                runtime: "fixture".into(),
                media_supported: false,
                code_query_supported: false,
            },
            healthy: true,
            force_cpu: true,
            model_dir: PathBuf::new(),
            worker_path: PathBuf::new(),
            projector_verified: false,
            query_profile: None,
        }
    }

    #[test]
    fn disconnected_writer_queue_is_transport_failure_and_reaps_helper() {
        let mut fixture = stalled_protocol_helper(false);
        let (sender, receiver) = mpsc::sync_channel(1);
        drop(receiver);
        fixture.requests = Some(sender);
        let error = fixture
            .exchange_with_timeout(b"{}", Duration::from_millis(300))
            .unwrap_err();
        assert!(error.downcast_ref::<WorkerTransportFailure>().is_some());
        assert!(fixture.child.try_wait().unwrap().is_some());
        assert!(!fixture.healthy && fixture.requests.is_none());
        assert!(fixture.writer_thread.is_none() && fixture.reader_thread.is_none());
    }

    #[test]
    fn entire_request_deadline_kills_and_joins_worker_that_never_reads_stdin() {
        let mut fixture = stalled_protocol_helper(false);
        let started = Instant::now();
        let error = fixture
            .exchange_with_timeout(&vec![b'x'; MAX_REQUEST - 1], Duration::from_millis(300))
            .unwrap_err();
        assert!(error.downcast_ref::<WorkerTransportFailure>().is_some());
        assert!(error.to_string().contains("write failed or timed out"));
        assert!(started.elapsed() < Duration::from_secs(5));
        assert!(fixture.child.try_wait().unwrap().is_some());
        assert!(!fixture.healthy && !fixture.info.ready);
        assert!(fixture.writer_thread.is_none() && fixture.reader_thread.is_none());
    }

    #[test]
    fn response_wait_uses_same_deadline_and_reaps_silent_reader() {
        let mut fixture = stalled_protocol_helper(true);
        let started = Instant::now();
        let error = fixture
            .exchange_with_timeout(b"{}", Duration::from_millis(300))
            .unwrap_err();
        assert!(error.downcast_ref::<WorkerTransportFailure>().is_some());
        assert!(error.to_string().contains("request failed or timed out"));
        assert!(started.elapsed() < Duration::from_secs(5));
        assert!(fixture.child.try_wait().unwrap().is_some());
        assert!(fixture.requests.is_none());
        assert!(fixture.writer_thread.is_none() && fixture.reader_thread.is_none());
    }

    #[test]
    fn automatic_startup_transport_failure_retries_cpu_once() {
        let mut attempts = Vec::new();
        let result = start_with_cpu_fallback(false, |cpu| {
            attempts.push(cpu);
            if cpu {
                Ok("cpu ready")
            } else {
                Err(WorkerTransportFailure("worker closed its output before ready".into()).into())
            }
        })
        .unwrap();
        assert_eq!(result, "cpu ready");
        assert_eq!(attempts, [false, true]);
    }

    #[test]
    fn startup_failure_never_loops_or_retries_explicit_cpu() {
        for cpu in [false, true] {
            let mut attempts = Vec::new();
            let result: Result<()> = start_with_cpu_fallback(cpu, |selected_cpu| {
                attempts.push(selected_cpu);
                Err(WorkerTransportFailure("worker exited".into()).into())
            });
            let error = format!("{:#}", result.unwrap_err());
            assert_eq!(attempts, if cpu { vec![true] } else { vec![false, true] });
            assert!(error.contains("worker exited"));
            if !cpu {
                assert!(error.contains("CPU initialization failed"));
            }
        }
    }

    #[test]
    fn malformed_startup_protocol_is_not_retried() {
        let mut attempts = 0;
        let result: Result<()> = start_with_cpu_fallback(false, |_| {
            attempts += 1;
            let mut response = BufReader::new(&b"not JSON\n"[..]);
            let error = read_response(&mut response).unwrap_err();
            assert!(error.starts_with("Invalid Gemma worker JSON:"));
            anyhow::bail!("Invalid Gemma worker startup protocol: {error}");
        });
        assert!(result.unwrap_err().to_string().contains("startup protocol"));
        assert_eq!(attempts, 1);
    }

    #[test]
    fn incompatible_startup_contract_is_not_retried() {
        let mut attempts = 0;
        let result: Result<()> = start_with_cpu_fallback(false, |_| {
            attempts += 1;
            let handshake = json!({"ready":true,"model":"other-model","dimensions":768,"device":"cpu","threads":1,"runtime":"other-runtime"});
            let info: WorkerInfo = serde_json::from_value(handshake)?;
            ensure!(
                info.model == "google/embeddinggemma-2",
                "Gemma worker handshake violates configured model/device contract"
            );
            Ok(())
        });
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("model/device contract"));
        assert_eq!(attempts, 1);
    }

    #[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);
        assert!(!old.code_query_supported);
        assert!(query_operation(&old, true).is_err());
        assert_eq!(query_operation(&old, false).unwrap(), "query");
        let mut new = old.clone();
        new.code_query_supported = true;
        assert_eq!(query_operation(&new, true).unwrap(), "code_query");
        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());
    }
}