//! 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>,
}

pub(crate) 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).
pub(crate) 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)
}

pub(crate) 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?;
        // A copy in the image store keeps it viewable when the file is moved or deleted.
        let blob = app.store.put_blob(&ctype, &bytes).map_err(|e| tracing::warn!("image copy: {e:#}")).ok();
        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(),
            blob,
        };
        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)
}

/// A generated image whose file is gone: the same file name below its original folder
/// (or the folder above it), e.g. after an agent sorted the images into subfolders.
/// Only inside `roots`; looks at a few thousand entries at most.
pub fn find_moved(path: &Path, roots: &[PathBuf]) -> Option<PathBuf> {
    let name = path.file_name()?;
    search_near(path, roots, |p, _| p.file_name() == Some(name))
}

/// Like `find_moved`, but by content: an image file with the same bytes (renamed).
pub fn find_by_content(path: &Path, roots: &[PathBuf], hash: &str, size: u64) -> Option<PathBuf> {
    search_near(path, roots, |p, len| {
        let ext = p.extension().and_then(|e| e.to_str()).map(|e| e.to_ascii_lowercase());
        len == size && matches!(ext.as_deref(), Some("png" | "jpg" | "jpeg" | "webp")) && std::fs::read(p).is_ok_and(|b| crate::blobs::hash(&b) == hash)
    })
}

/// Files below the folder of `path` and the folder above it (3 levels, hidden folders
/// skipped, inside `roots`) for which `matches(path, size)` holds.
fn search_near(path: &Path, roots: &[PathBuf], matches: impl Fn(&Path, u64) -> bool) -> Option<PathBuf> {
    let dir = path.parent()?;
    let mut budget = 5000usize;
    for base in [Some(dir), dir.parent()].into_iter().flatten() {
        if !roots.iter().any(|r| base.starts_with(r)) {
            continue;
        }
        let mut stack = vec![(base.to_path_buf(), 0usize)];
        while let Some((d, depth)) = stack.pop() {
            let Ok(entries) = std::fs::read_dir(&d) else { continue };
            for e in entries.flatten() {
                if budget == 0 {
                    return None;
                }
                budget -= 1;
                let p = e.path();
                let Ok(ft) = e.file_type() else { continue };
                if ft.is_file() && p != path && matches(&p, e.metadata().map(|m| m.len()).unwrap_or(0)) {
                    return Some(p);
                }
                let hidden = p.file_name().and_then(|n| n.to_str()).is_some_and(|n| n.starts_with('.') || n == "node_modules");
                if ft.is_dir() && depth < 3 && !hidden {
                    stack.push((p, depth + 1));
                }
            }
        }
    }
    None
}

/// Where the file of a generated image is now: by name, then (with a copy in the image
/// store) by content. Corrects the stored path when found.
pub fn relocate(app: &App, rec: &ImageRecord) -> Option<PathBuf> {
    let path = Path::new(&rec.path);
    let roots = app.user_roots(rec.user_id.as_deref());
    let found = find_moved(path, &roots).or_else(|| {
        let hash = rec.blob.as_deref()?;
        let (_, bytes) = app.store.blob(hash).ok()??;
        find_by_content(path, &roots, hash, bytes.len() as u64)
    })?;
    let _ = app.store.set_image_path(&rec.id, &found.display().to_string());
    Some(found)
}

/// After an agent moved or renamed files: find the generated images of `owner` from the
/// last two weeks whose files are gone and correct their paths. Returns how many.
pub fn relocate_missing(app: &App, owner: Option<&str>) -> usize {
    let since = now_ms() - 14 * 86_400_000;
    let Ok(list) = app.store.images(owner, 500) else { return 0 };
    list.iter().filter(|r| r.created_at >= since && !Path::new(&r.path).exists()).filter(|r| relocate(app, r).is_some()).count()
}

/// Copies in the image store for generated images from before they were kept (once per
/// image, files up to 25 MB): later moves and renames are then found by content.
pub fn backfill_copies(app: &App) -> Result<usize> {
    let mut n = 0;
    for (id, path) in app.store.images_without_copy()? {
        let Ok(meta) = std::fs::metadata(&path) else { continue };
        if !meta.is_file() || meta.len() > 25 * 1024 * 1024 {
            continue;
        }
        let Ok(bytes) = std::fs::read(&path) else { continue };
        let mime = match Path::new(&path).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",
        };
        let hash = app.store.put_blob(mime, &bytes)?;
        app.store.set_image_blob(&id, &hash)?;
        n += 1;
    }
    Ok(n)
}

/// Downscaled copy for the model's context.
pub(crate) 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 { blob: None, size: None, 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::*;

    #[test]
    fn finds_moved_files_and_keeps_their_copies() {
        let dir = std::env::temp_dir().join(format!("cc-img-moved-{}", std::process::id()));
        let art = dir.join("art");
        std::fs::create_dir_all(art.join("older")).unwrap();
        std::fs::write(art.join("older/20261110-a.png"), b"x").unwrap();
        // Moved into a subfolder: found; renamed or outside the roots: not.
        assert_eq!(find_moved(&art.join("20261110-a.png"), std::slice::from_ref(&dir)), Some(art.join("older/20261110-a.png")));
        assert_eq!(find_moved(&art.join("20261110-b.png"), std::slice::from_ref(&dir)), None);
        assert_eq!(find_moved(&art.join("20261110-a.png"), &[dir.join("elsewhere")]), None);
        // Renamed: found by content.
        std::fs::write(art.join("older/level-5.png"), b"renamed bytes").unwrap();
        let h = crate::blobs::hash(b"renamed bytes");
        assert_eq!(find_by_content(&art.join("20261110-c.png"), std::slice::from_ref(&dir), &h, 13), Some(art.join("older/level-5.png")));
        assert_eq!(find_by_content(&art.join("20261110-c.png"), std::slice::from_ref(&dir), &crate::blobs::hash(b"other"), 13), None);
        // The image store keeps the copy of a generated image.
        let store = crate::store::Store::open(&dir.join("t.db")).unwrap();
        let hash = store.put_blob("image/png", b"png bytes").unwrap();
        store
            .add_image(&ImageRecord { id: "i".into(), user_id: None, session_id: None, prompt: "p".into(), model: "m".into(), path: "/gone.png".into(), width: 1, height: 1, created_at: 0, blob: Some(hash.clone()) })
            .unwrap();
        assert_eq!(store.gc_blobs().unwrap(), 0);
        // Moved with the file manager: the entry follows (file or folder, not look-alikes).
        assert_eq!(store.move_image_paths("/gone.png", "/art/new.png").unwrap(), 1);
        assert_eq!(store.move_image_paths("/art", "/bilder/").unwrap(), 1);
        assert_eq!(store.move_image_paths("/bild", "/x").unwrap(), 0);
        assert_eq!(store.image("i").unwrap().unwrap().path, "/bilder/new.png");
        assert_eq!(store.blob(&hash).unwrap().unwrap().1, b"png bytes");
        std::fs::remove_dir_all(&dir).unwrap();
    }

    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());
    }
}