codecraft/server/src/images.rs

//! Image generation via fal.ai – as an agent tool and for the manual image page.

use crate::app::App;
use crate::config::{ImageModel, expand_env};
use crate::store::{ImageRecord, UsageRecord};
use crate::tools::{Tool, ToolCtx, ToolOutput, arg_str};
use crate::types::{ImageData, new_id, now_ms};
use anyhow::{Context, Result, anyhow, bail};
use async_trait::async_trait;
use base64::Engine;
use serde_json::{Value, json};
use std::path::{Path, PathBuf};
use std::time::Duration;

/// Size presets offered in the UI and to the model (FLUX-style names).
pub const SIZES: &[&str] = &["square_hd", "square", "landscape_4_3", "landscape_16_9", "portrait_4_3", "portrait_16_9"];

fn aspect_ratio(size: &str) -> &'static str {
    match size {
        "landscape_4_3" => "4:3",
        "landscape_16_9" => "16:9",
        "portrait_4_3" => "3:4",
        "portrait_16_9" => "9:16",
        _ => "1:1",
    }
}

pub struct GenerateRequest {
    pub prompt: String,
    pub model: Option<String>,
    pub size: Option<String>,
    pub count: u32,
    /// Directory the files are written to.
    pub dir: PathBuf,
    pub user_id: Option<String>,
    pub session_id: Option<String>,
    /// Images to refine (sent to models with a `reference_param`).
    pub references: Vec<PathBuf>,
}

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(40).collect::<String>().trim_end_matches('-').to_string()
}

/// Purposes the agent can ask for; models list theirs in `good_for`.
pub const PURPOSES: &[&str] = &["photo", "illustration", "text", "design"];

/// Defaults for well-known fal.ai models: (id, stage, good_for, reference_param, size_param, price).
const KNOWN: &[(&str, &str, &[&str], Option<&str>, Option<&str>, f64)] = &[
    ("fal-ai/flux/schnell", "preview", &["photo", "illustration"], None, None, 0.003),
    ("fal-ai/nano-banana-pro", "final", &["photo", "illustration", "text"], None, Some("aspect_ratio"), 0.15),
    ("fal-ai/nano-banana-pro/edit", "final", &["photo", "illustration", "text", "design"], Some("image_urls"), Some("aspect_ratio"), 0.15),
    ("fal-ai/flux-2-pro", "final", &["photo"], None, None, 0.03),
    ("fal-ai/ideogram/v3", "final", &["text", "design"], None, Some("aspect_ratio"), 0.06),
    ("fal-ai/recraft/v4/pro/text-to-image", "final", &["design", "illustration"], None, None, 0.05),
];

/// A configured model with the known defaults filled in where the config leaves them out.
fn with_defaults(mut m: ImageModel) -> ImageModel {
    if let Some((_, stage, good_for, reference, size, price)) = KNOWN.iter().find(|k| k.0 == m.id) {
        m.stage.get_or_insert_with(|| stage.to_string());
        if m.good_for.is_empty() {
            m.good_for = good_for.iter().map(|s| s.to_string()).collect();
        }
        if m.reference_param.is_none() {
            m.reference_param = reference.map(String::from);
        }
        if m.size_param.is_none() {
            m.size_param = size.map(String::from);
        }
        m.price.get_or_insert(*price);
    }
    m
}

fn models(app: &App) -> Vec<ImageModel> {
    app.config().images.models.iter().cloned().map(with_defaults).collect()
}

fn model_config(app: &App, id: &str) -> ImageModel {
    models(app).into_iter().find(|m| m.id == id).unwrap_or_else(|| with_defaults(ImageModel { id: id.to_string(), ..Default::default() }))
}

/// Pick the model for a request: an explicit model wins; previews use the first preview model
/// (else the default); final images use the best-matching final model, preferring one that can
/// take the reference image when there is one.
pub fn choose_model(app: &App, stage: &str, purpose: Option<&str>, with_reference: bool) -> ImageModel {
    let default = app.config().images.default_model.clone();
    pick(&models(app), stage, purpose, with_reference).unwrap_or_else(|| model_config(app, &default))
}

fn pick(all: &[ImageModel], stage: &str, purpose: Option<&str>, with_reference: bool) -> Option<ImageModel> {
    if stage == "preview" && !with_reference {
        return all.iter().find(|m| m.stage.as_deref() == Some("preview")).cloned();
    }
    let score = |m: &ImageModel| {
        let mut s = 0;
        if m.stage.as_deref() == Some("final") {
            s += 4;
        }
        match (with_reference, m.reference_param.is_some()) {
            (true, true) => s += 8,
            // Without a reference, plain text-to-image endpoints come first.
            (false, true) => s -= 1,
            _ => {}
        }
        if purpose.is_some_and(|p| m.good_for.iter().any(|g| g == p)) {
            s += 2;
        }
        s
    };
    // max_by_key keeps the last of equal scores; reverse so the config order wins ties.
    all.iter().rev().max_by_key(|m| score(m)).cloned()
}

/// Read local reference images as data URIs (fal.ai accepts them in place of URLs).
async fn reference_uris(paths: &[PathBuf]) -> Result<Vec<String>> {
    let mut out = Vec::new();
    for p in paths {
        let bytes = tokio::fs::read(p).await.with_context(|| format!("reference image {}", p.display()))?;
        let mime = match p.extension().and_then(|e| e.to_str()).map(|e| e.to_ascii_lowercase()).as_deref() {
            Some("jpg" | "jpeg") => "image/jpeg",
            Some("webp") => "image/webp",
            _ => "image/png",
        };
        out.push(format!("data:{mime};base64,{}", base64::engine::general_purpose::STANDARD.encode(bytes)));
    }
    Ok(out)
}

async fn download(url: &str) -> Result<(Vec<u8>, String)> {
    if let Some(rest) = url.strip_prefix("data:") {
        let (meta, data) = rest.split_once(',').context("malformed data URI")?;
        let ctype = meta.split(';').next().unwrap_or("image/png").to_string();
        return Ok((base64::engine::general_purpose::STANDARD.decode(data)?, ctype));
    }
    let resp = crate::provider::http_client().get(url).timeout(Duration::from_secs(120)).send().await?.error_for_status()?;
    let ctype = resp.headers().get(reqwest::header::CONTENT_TYPE).and_then(|v| v.to_str().ok()).unwrap_or("image/png").to_string();
    Ok((resp.bytes().await?.to_vec(), ctype))
}

/// Generate images, store them in `req.dir` and record them (history + cost).
pub async fn generate(app: &App, req: GenerateRequest) -> Result<Vec<ImageRecord>> {
    let cfg = app.config();
    let ic = &cfg.images;
    if !ic.enabled {
        bail!("image generation is disabled – enable [images] in the settings");
    }
    let key = expand_env(&ic.api_key);
    if key.is_empty() {
        bail!("no fal.ai API key configured (images.api_key, e.g. ${{FAL_KEY}})");
    }
    crate::budget::check(app, "fal", req.user_id.as_deref())?;
    let model_id = req.model.clone().filter(|m| !m.is_empty()).unwrap_or_else(|| ic.default_model.clone());
    let model = model_config(app, &model_id);
    let explicit_size = req.size.clone().filter(|s| SIZES.contains(&s.as_str()));
    let size = explicit_size.clone().unwrap_or_else(|| "landscape_4_3".into());
    let count = req.count.clamp(1, 4);

    let mut body = json!({"prompt": req.prompt, "num_images": count, "output_format": "png"});
    let refining = !req.references.is_empty() && model.reference_param.is_some();
    match model.size_param.as_deref() {
        // When refining without an explicit size, keep the reference's aspect ratio.
        Some("aspect_ratio") if refining && explicit_size.is_none() => body["aspect_ratio"] = json!("auto"),
        Some("aspect_ratio") => body["aspect_ratio"] = json!(aspect_ratio(&size)),
        Some("none") => {}
        _ => body["image_size"] = json!(size),
    }
    if refining && let Some(param) = &model.reference_param {
        body[param.as_str()] = json!(reference_uris(&req.references).await?);
    }
    for (k, v) in &model.options {
        body[k] = v.clone();
    }

    let url = format!("{}/{}", expand_env(&ic.base_url).trim_end_matches('/'), model_id.trim_matches('/'));
    let resp = crate::provider::http_client()
        .post(&url)
        .header("Authorization", format!("Key {key}"))
        .timeout(Duration::from_secs(300))
        .json(&body)
        .send()
        .await
        .context("fal.ai request failed")?;
    if !resp.status().is_success() {
        return Err(crate::provider::http_error(resp).await);
    }
    let out: Value = resp.json().await?;
    let images = out["images"].as_array().cloned().or_else(|| out["image"].as_object().map(|o| vec![Value::Object(o.clone())])).unwrap_or_default();
    if images.is_empty() {
        bail!("fal.ai returned no images: {}", out.to_string().chars().take(300).collect::<String>());
    }

    tokio::fs::create_dir_all(&req.dir).await?;
    let stamp = chrono::Local::now().format("%Y%m%d-%H%M%S");
    let base = slug(&req.prompt);
    let mut records = Vec::new();
    for (i, img) in images.iter().enumerate() {
        let src = img["url"].as_str().context("image without url")?;
        let (bytes, ctype) = download(src).await?;
        let ext = match ctype.as_str() {
            "image/jpeg" => "jpg",
            "image/webp" => "webp",
            _ => "png",
        };
        let name = if images.len() > 1 { format!("{stamp}-{base}-{}.{ext}", i + 1) } else { format!("{stamp}-{base}.{ext}") };
        let path = req.dir.join(name);
        tokio::fs::write(&path, &bytes).await?;
        let rec = ImageRecord {
            id: new_id(),
            user_id: req.user_id.clone(),
            session_id: req.session_id.clone(),
            prompt: req.prompt.clone(),
            model: model_id.clone(),
            path: path.display().to_string(),
            width: img["width"].as_i64().unwrap_or(0),
            height: img["height"].as_i64().unwrap_or(0),
            created_at: now_ms(),
        };
        app.store.add_image(&rec)?;
        records.push(rec);
    }

    // Costs go into the normal statistics (provider "fal").
    let cost = model.price.unwrap_or(0.0) * records.len() as f64;
    let session = req.session_id.clone().unwrap_or_default();
    let _ = app.store.add_usage(&UsageRecord {
        ts: now_ms(),
        session_id: session.clone(),
        root_session_id: if session.is_empty() { String::new() } else { app.root_session_id(&session) },
        provider: "fal".into(),
        model: model_id,
        agent: "image".into(),
        input_tokens: 0,
        output_tokens: 0,
        cache_read_tokens: 0,
        cache_write_tokens: 0,
        ttft_ms: 0,
        duration_ms: 0,
        cost,
        error: false,
        user_id: None,
        allowance: None,
    });
    Ok(records)
}

/// Downscaled copy for the model's context.
fn preview_for_model(path: &Path) -> Option<ImageData> {
    let img = image::open(path).ok()?;
    let img = img.resize(1024, 1024, image::imageops::FilterType::Triangle);
    let mut buf = Vec::new();
    img.write_to(&mut std::io::Cursor::new(&mut buf), image::ImageFormat::Png).ok()?;
    Some(ImageData { media_type: "image/png".into(), data: base64::engine::general_purpose::STANDARD.encode(buf) })
}

pub struct GenerateImage;

#[async_trait]
impl Tool for GenerateImage {
    fn name(&self) -> &str {
        "generate_image"
    }
    fn description(&self) -> String {
        "Generate images from a text prompt (fal.ai) and save them in the project. Use for \
         illustrations, icons, backgrounds or photos the user needs. Write a detailed visual prompt \
         in English (subject, style, composition, colours, lighting). The server picks the model. \
         Work in two steps: first stage \"preview\" (cheap drafts, count 2-4); show them and ask the \
         user with ask_user which one to refine – include the final model and price the preview \
         result names. Then call again with stage \"final\", the chosen file in \"reference\" and \
         the same or an improved prompt. Skip the preview only if the user asks for a final image \
         directly. Each image costs money – generate only what is needed."
            .into()
    }
    fn schema(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "prompt": {"type": "string"},
                "stage": {"type": "string", "enum": ["preview", "final"], "description": "preview = cheap drafts (default), final = best model for the purpose"},
                "purpose": {"type": "string", "enum": PURPOSES, "description": "What the image is: photo (realistic), illustration, text (lettering, posters, signs), design (logos, icons, UI, vector style)"},
                "reference": {"type": "array", "items": {"type": "string"}, "maxItems": 4, "description": "Image files (relative paths) to refine, e.g. the preview the user chose"},
                "size": {"type": "string", "enum": SIZES, "description": "Default: landscape_4_3; when refining, the reference's aspect ratio"},
                "count": {"type": "integer", "minimum": 1, "maximum": 4},
                "directory": {"type": "string", "description": "Target directory relative to the working directory (default: images.output_dir)"},
                "model": {"type": "string", "description": "fal.ai model id – only if the user asks for a specific model"}
            },
            "required": ["prompt"]
        })
    }
    fn title(&self, input: &Value) -> String {
        let stage = if input["stage"].as_str() == Some("final") { "Generate final image" } else { "Generate image" };
        format!("{stage}: {}", input["prompt"].as_str().unwrap_or("").chars().take(80).collect::<String>())
    }
    async fn run(&self, input: Value, ctx: &ToolCtx) -> Result<ToolOutput> {
        let prompt = arg_str(&input, "prompt")?.to_string();
        let stage = if input["stage"].as_str() == Some("final") { "final" } else { "preview" };
        let purpose = input["purpose"].as_str().filter(|p| PURPOSES.contains(p));
        let dir_arg = input["directory"].as_str().map(String::from).unwrap_or_else(|| ctx.app.config().images.output_dir.clone());
        let dir = ctx.resolve(&dir_arg)?;
        let mut references = Vec::new();
        for r in input["reference"].as_array().into_iter().flatten().filter_map(|v| v.as_str()).take(4) {
            references.push(ctx.resolve(r)?);
        }
        let model = match input["model"].as_str().filter(|m| !m.is_empty()) {
            Some(m) => model_config(&ctx.app, m),
            None => choose_model(&ctx.app, stage, purpose, !references.is_empty()),
        };
        let owner = ctx.app.session(&ctx.session_id).ok().and_then(|s| s.owner_id);
        let refined = !references.is_empty() && model.reference_param.is_some();
        let records = generate(
            &ctx.app,
            GenerateRequest {
                prompt,
                model: Some(model.id.clone()),
                size: input["size"].as_str().map(String::from),
                count: input["count"].as_u64().unwrap_or(if stage == "preview" { 2 } else { 1 }) as u32,
                dir,
                user_id: owner,
                session_id: Some(ctx.session_id.clone()),
                references,
            },
        )
        .await
        .map_err(|e| anyhow!("{e:#}"))?;
        let price = |m: &ImageModel| m.price.map(|p| format!(" (${p:.3} per image)")).unwrap_or_default();
        let label = |m: &ImageModel| m.name.clone().unwrap_or_else(|| m.id.clone());
        let mut text = format!("Model: {}{}\nGenerated:\n", label(&model), price(&model));
        let mut images = Vec::new();
        for r in &records {
            let rel = Path::new(&r.path).strip_prefix(&ctx.cwd).map(|p| p.display().to_string()).unwrap_or_else(|_| r.path.clone());
            text.push_str(&format!("- {rel} ({}×{})\n", r.width, r.height));
            if let Some(p) = preview_for_model(Path::new(&r.path)) {
                images.push(p);
            }
        }
        if stage == "preview" && !refined {
            let fin = choose_model(&ctx.app, "final", purpose, true);
            if fin.id != model.id {
                let how = if fin.reference_param.is_some() { "refines the chosen preview" } else { "generates it again from the prompt" };
                text.push_str(&format!(
                    "\nThese are previews. Ask the user (ask_user) which one to refine. The final image would use {}{} and {how}.",
                    label(&fin),
                    price(&fin)
                ));
            }
        } else if !input["reference"].as_array().is_none_or(|a| a.is_empty()) && !refined {
            text.push_str("\nNote: this model takes no reference image, so it was generated from the prompt only.");
        }
        let ids: Vec<&str> = records.iter().map(|r| r.id.as_str()).collect();
        Ok(ToolOutput { content: text, images, meta: Some(json!({"image_ids": ids})), ..Default::default() })
    }
}

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

    fn m(id: &str) -> ImageModel {
        with_defaults(ImageModel { id: id.into(), ..Default::default() })
    }

    #[test]
    fn picks_by_stage_purpose_and_reference() {
        let all = vec![
            m("fal-ai/flux/schnell"),
            m("fal-ai/nano-banana-pro/edit"),
            m("fal-ai/nano-banana-pro"),
            m("fal-ai/flux-2-pro"),
            m("fal-ai/ideogram/v3"),
            m("fal-ai/recraft/v4/pro/text-to-image"),
        ];
        let id = |stage, purpose, r| pick(&all, stage, purpose, r).unwrap().id;
        assert_eq!(id("preview", Some("photo"), false), "fal-ai/flux/schnell");
        // Refining a preview: the edit model, whatever the purpose.
        assert_eq!(id("final", Some("photo"), true), "fal-ai/nano-banana-pro/edit");
        assert_eq!(id("final", Some("design"), true), "fal-ai/nano-banana-pro/edit");
        // From the prompt only: the first fitting text-to-image model (config order).
        assert_eq!(id("final", Some("photo"), false), "fal-ai/nano-banana-pro");
        assert_eq!(id("final", Some("text"), false), "fal-ai/nano-banana-pro");
        assert_eq!(id("final", Some("design"), false), "fal-ai/ideogram/v3");
        let no_edit: Vec<_> = all.iter().filter(|m| m.reference_param.is_none()).cloned().collect();
        assert_eq!(pick(&no_edit, "final", Some("design"), true).unwrap().id, "fal-ai/ideogram/v3");
        assert_eq!(pick(&no_edit, "final", Some("photo"), false).unwrap().id, "fal-ai/nano-banana-pro");
    }

    #[test]
    fn config_overrides_known_defaults() {
        let custom = with_defaults(ImageModel { id: "fal-ai/flux/schnell".into(), stage: Some("final".into()), price: Some(0.01), ..Default::default() });
        assert_eq!(custom.stage.as_deref(), Some("final"));
        assert_eq!(custom.price, Some(0.01));
        assert!(m("some/unknown-model").stage.is_none());
    }
}