dm-pal/src-tauri/src/commands/image_commands.rs

use crate::llm::AppState;
use futures_util::StreamExt;
use serde::{Deserialize, Serialize};
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use tauri::ipc::Channel;

// ponytail: DefaultHasher is fine for a cache filename — not crypto, just a stable key.

#[derive(Debug, Deserialize)]
pub struct ImageRequest {
    pub prompt: String,
    /// Override the configured image model for this call.
    pub model: Option<String>,
}

/// One line of Ollama's NDJSON image-generation response.
/// `step`/`total` carry progress; the final `done: true` line carries the
/// singular `image` base64 field. All fields optional per-line.
#[derive(Debug, Default, Deserialize)]
struct OllamaImageLine {
    #[serde(default)]
    done: bool,
    #[serde(default)]
    step: Option<u32>,
    #[serde(default)]
    total: Option<u32>,
    #[serde(default)]
    image: Option<String>,
}

/// Parsed view of one NDJSON line used by both the streaming and the
/// buffered paths. Extracted so the parsing logic is unit-testable without
/// touching the network.
#[derive(Debug, PartialEq)]
struct ImageProgress {
    done: bool,
    step: Option<u32>,
    total: Option<u32>,
    image: Option<String>,
}

fn parse_image_line(line: &str) -> Option<ImageProgress> {
    let line = line.trim();
    if line.is_empty() {
        return None;
    }
    let parsed: OllamaImageLine = serde_json::from_str(line).ok()?;
    Some(ImageProgress {
        done: parsed.done,
        step: parsed.step,
        total: parsed.total,
        image: parsed.image,
    })
}

/// Channel events for streaming image generation.
#[derive(Clone, Serialize)]
#[serde(tag = "type", content = "data", rename_all = "camelCase")]
pub enum ImageEvent {
    /// Progress update: (step, total). Either may be None if Ollama omits it.
    Progress { step: Option<u32>, total: Option<u32> },
    /// Final result: a `data:image/png;base64,...` URL ready for `<img src>`.
    Done(String),
    /// Fatal error.
    Error(String),
}

/// Generate (or fetch from disk cache) an image for `prompt` via the configured
/// Ollama image model. Returns a `data:image/png;base64,...` URL ready for `<img src>`.
///
/// Ollama image models are macOS-only today; on other platforms we return an error
/// so the front-end can fall back to a placeholder instead of a confusing timeout.
#[tauri::command]
pub async fn generate_image(
    state: tauri::State<'_, AppState>,
    req: ImageRequest,
) -> Result<String, String> {
    if cfg!(not(target_os = "macos")) {
        return Err("image generation is macOS-only via Ollama (for now)".into());
    }

    let config = state.config.lock().map_err(|e| e.to_string())?.clone();
    let model = req.model.unwrap_or(config.image_model.clone());

    let (cache_path, cache_hit) = prepare_cache(&state.data_dir, &model, &req.prompt)?;
    if cache_hit {
        let bytes = std::fs::read(&cache_path).map_err(|e| e.to_string())?;
        return Ok(data_url(&bytes));
    }

    let png_bytes = request_image_bytes(&config, &model, &req.prompt, None).await?;
    std::fs::write(&cache_path, &png_bytes).map_err(|e| e.to_string())?;
    Ok(data_url(&png_bytes))
}

/// Streaming variant: emits `ImageEvent::Progress` as Ollama reports `step`/`total`,
/// then `ImageEvent::Done` with the data URL (or `Error`). Reuses the same disk cache
/// as `generate_image`. The DM gets a real progress bar for the multi-second wait.
#[tauri::command]
pub async fn generate_image_stream(
    state: tauri::State<'_, AppState>,
    req: ImageRequest,
    channel: Channel<ImageEvent>,
) -> Result<(), String> {
    if cfg!(not(target_os = "macos")) {
        let _ = channel.send(ImageEvent::Error(
            "image generation is macOS-only via Ollama (for now)".into(),
        ));
        return Ok(());
    }

    let config = state.config.lock().map_err(|e| e.to_string())?.clone();
    let model = req.model.unwrap_or(config.image_model.clone());

    let (cache_path, cache_hit) = prepare_cache(&state.data_dir, &model, &req.prompt)?;
    if cache_hit {
        let bytes = std::fs::read(&cache_path).map_err(|e| e.to_string())?;
        let _ = channel.send(ImageEvent::Done(data_url(&bytes)));
        return Ok(());
    }

    // Drive the request on a background task so the command returns immediately
    // and progress flows through the channel. Errors become ImageEvent::Error.
    let channel = std::sync::Arc::new(channel);
    let ch = channel.clone();
    tauri::async_runtime::spawn(async move {
        match request_image_bytes(&config, &model, &req.prompt, Some(ch)).await {
            Ok(png_bytes) => {
                let _ = std::fs::write(&cache_path, &png_bytes);
                let _ = channel.send(ImageEvent::Done(data_url(&png_bytes)));
            }
            Err(e) => {
                let _ = channel.send(ImageEvent::Error(e));
            }
        }
    });

    Ok(())
}

/// Resolve the on-disk cache path for (model, prompt) and report a cache hit.
/// `base` is the configured data dir (AppState.data_dir).
fn prepare_cache(
    base: &std::path::Path,
    model: &str,
    prompt: &str,
) -> Result<(std::path::PathBuf, bool), String> {
    let cache_dir = base.join("images");
    std::fs::create_dir_all(&cache_dir).map_err(|e| e.to_string())?;

    let mut hasher = DefaultHasher::new();
    model.hash(&mut hasher);
    prompt.hash(&mut hasher);
    let cache_path = cache_dir.join(format!("{:016x}.png", hasher.finish()));
    let hit = cache_path.exists();
    Ok((cache_path, hit))
}

/// POST the image request to Ollama and collect the PNG bytes. When a channel
/// is given, parse the NDJSON body incrementally and emit progress; otherwise
/// read the whole body at once (legacy buffered path).
async fn request_image_bytes(
    config: &crate::llm::LlmConfig,
    model: &str,
    prompt: &str,
    channel: Option<std::sync::Arc<Channel<ImageEvent>>>,
) -> Result<Vec<u8>, String> {
    let client = reqwest::Client::new();
    let url = format!("{}/api/generate", config.api_url.trim_end_matches('/'));
    let body = serde_json::json!({ "model": model, "prompt": prompt, "stream": false });

    let res = client
        .post(&url)
        .json(&body)
        .send()
        .await
        .map_err(|e| format!("image request failed: {e}"))?;

    if !res.status().is_success() {
        let status = res.status();
        let text = res.text().await.unwrap_or_default();
        return Err(format!("image error {status}: {text}"));
    }

    // ponytail: Ollama image models emit NDJSON even with stream:false —
    // progress lines (step/total) then a final done:true carrying the image.
    // The buffered path reads it all at once; the streaming path splits lines
    // as they arrive so the bar animates.
    let mut png_b64: Option<String> = None;
    let mut buffer = String::new();

    let apply_line = |line: &str, png: &mut Option<String>, ch: Option<&Channel<ImageEvent>>| {
        if let Some(p) = parse_image_line(line) {
            if let Some(c) = ch {
                let _ = c.send(ImageEvent::Progress { step: p.step, total: p.total });
            }
            if let Some(b64) = p.image {
                *png = Some(b64);
            }
        }
    };

    if let Some(ch) = channel {
        let mut stream = res.bytes_stream();
        while let Some(chunk) = stream.next().await {
            let chunk = chunk.map_err(|e| format!("image read error: {e}"))?;
            buffer.push_str(&String::from_utf8_lossy(&chunk));
            // Process complete lines; keep the trailing partial line in buffer.
            while let Some(idx) = buffer.find('\n') {
                let line = buffer.split_off(idx + 1);
                let complete = std::mem::replace(&mut buffer, line);
                apply_line(&complete, &mut png_b64, Some(&ch));
            }
        }
        if !buffer.trim().is_empty() {
            apply_line(&buffer, &mut png_b64, Some(&ch));
        }
    } else {
        let body_text = res.text().await.map_err(|e| format!("image read error: {e}"))?;
        for line in body_text.lines() {
            apply_line(line, &mut png_b64, None);
            // ponytail: stop after the done line in the buffered path — the
            // final image is the one we want.
            if let Some(p) = parse_image_line(line) {
                if p.done && p.image.is_some() {
                    break;
                }
            }
        }
    }

    let b64 = png_b64.ok_or_else(|| "no image data in Ollama response".to_string())?;
    let png_bytes = base64_decode(&b64)?;
    Ok(png_bytes)
}

fn data_url(png: &[u8]) -> String {
    // ponytail: no base64 crate dep — a 30-line encoder.
    let table: [u8; 64] = *b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut out = String::with_capacity((png.len() + 2) / 3 * 4);
    let mut chunks = png.chunks_exact(3);
    for c in &mut chunks {
        let n = (c[0] as usize) << 16 | (c[1] as usize) << 8 | (c[2] as usize);
        out.push(table[(n >> 18) & 63] as char);
        out.push(table[(n >> 12) & 63] as char);
        out.push(table[(n >> 6) & 63] as char);
        out.push(table[n & 63] as char);
    }
    let rem = chunks.remainder();
    match rem.len() {
        1 => {
            let n = (rem[0] as usize) << 16;
            out.push(table[(n >> 18) & 63] as char);
            out.push(table[(n >> 12) & 63] as char);
            out.push('=');
            out.push('=');
        }
        2 => {
            let n = (rem[0] as usize) << 16 | (rem[1] as usize) << 8;
            out.push(table[(n >> 18) & 63] as char);
            out.push(table[(n >> 12) & 63] as char);
            out.push(table[(n >> 6) & 63] as char);
            out.push('=');
        }
        _ => {}
    }
    format!("data:image/png;base64,{out}")
}

/// Minimal standard base64 decoder (no extra dependency).
fn base64_decode(input: &str) -> Result<Vec<u8>, String> {
    fn val(c: u8) -> Option<u8> {
        match c {
            b'A'..=b'Z' => Some(c - b'A'),
            b'a'..=b'z' => Some(c - b'a' + 26),
            b'0'..=b'9' => Some(c - b'0' + 52),
            b'+' => Some(62),
            b'/' => Some(63),
            _ => None,
        }
    }
    let input = input.trim();
    let bytes: Vec<u8> = input.bytes().filter(|b| !b.is_ascii_whitespace()).collect();
    if bytes.is_empty() {
        return Ok(Vec::new());
    }
    let mut out = Vec::with_capacity(bytes.len() * 3 / 4);
    let mut buf: u32 = 0;
    let mut bits: u32 = 0;
    for &b in &bytes {
        if b == b'=' {
            break;
        }
        let v = val(b).ok_or_else(|| format!("invalid base64 char: {b}"))? as u32;
        buf = (buf << 6) | v;
        bits += 6;
        if bits >= 8 {
            bits -= 8;
            out.push((buf >> bits) as u8);
            buf &= (1 << bits) - 1;
        }
    }
    Ok(out)
}

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

    #[test]
    fn roundtrip_base64() {
        let data = b"hello world \x00\xff\x10";
        let url = data_url(data);
        let b64 = url.strip_prefix("data:image/png;base64,").unwrap();
        let decoded = base64_decode(b64).unwrap();
        assert_eq!(decoded, data);
    }

    #[test]
    fn cache_key_is_stable() {
        // sanity: same inputs → same filename shape (16 hex digits)
        let mut h = DefaultHasher::new();
        "x/flux2-klein:4b".hash(&mut h);
        "prompt".hash(&mut h);
        let s = format!("{:016x}", h.finish());
        assert_eq!(s.len(), 16);
    }

    #[test]
    fn parse_progress_and_image_lines() {
        // intermediate progress line, no image
        let p = parse_image_line(r#"{"model":"x/flux2-klein:4b","done":false,"total":4,"step":1}"#).unwrap();
        assert_eq!(p, ImageProgress { done: false, step: Some(1), total: Some(4), image: None });

        // final done line carries the image
        let p = parse_image_line(r#"{"model":"x/flux2-klein:4b","done":true,"image":"iVBORw0KGgoAAAANSUhEUgAA"}"#).unwrap();
        assert!(p.done);
        assert_eq!(p.image.as_deref(), Some("iVBORw0KGgoAAAANSUhEUgAA"));
        assert!(p.total.is_none() && p.step.is_none());

        // blank/garbage lines are ignored, not errors
        assert!(parse_image_line("").is_none());
        assert!(parse_image_line("not json").is_none());
    }

    #[test]
    fn split_ndjson_buffer_keeps_trailing_partial() {
        // Simulate two chunks arriving separately where the split falls mid-line.
        let chunk1 = "{\"done\":false,\"step\":1,\"total\":4}\n{\"done\":tru";
        let chunk2 = "e,\"image\":\"abc\"}\n";
        let mut buffer = String::new();
        let mut png: Option<String> = None;
        let whole = format!("{chunk1}{chunk2}");
        // emulate the streaming loop over the concatenated body
        buffer.push_str(&whole);
        let mut lines = Vec::new();
        while let Some(idx) = buffer.find('\n') {
            let line = buffer.split_off(idx + 1);
            let complete = std::mem::replace(&mut buffer, line);
            lines.push(complete);
        }
        if !buffer.trim().is_empty() {
            lines.push(std::mem::take(&mut buffer));
        }
        for line in &lines {
            if let Some(p) = parse_image_line(line) {
                if let Some(b) = p.image { png = Some(b); }
            }
        }
        assert_eq!(png.as_deref(), Some("abc"));
    }
}