//! generate_image AI 工具 — 调 provider 的 OpenAI 兼容图像生成端点生成图片并下载落地。
//!
//! 设计目标:让 LLM 调用图像生成模型(默认 SenseNova U1 Fast,也支持 DALL-E 等 OpenAI 兼容
//! 图像端点)生成图片并落到工作区磁盘。图像生成模型不是 chat 模型,走 `/v1/images/generations`
//! 端点(非 chat completions),故不能经主对话链路调用,必须用本工具显式触发。
//!
//! ## 与 http_request / fetch_url / download_file 的分工
//!
//! - `http_request`:结构化 API 调用,返回 body(截断 50KB),不落盘 —— 通用 HTTP,不感知图像端点。
//! - `fetch_url`:网页 → markdown 文档嗅探,不落盘。
//! - `download_file`:URL → 落盘文件(已知 URL)。
//! - `generate_image`:**调 provider 图像端点 + 自动下载落地**(封装图像生成的 provider 凭证
//!   解析 + 端点拼接 + 响应解析 + 图片下载全链路,LLM 只需给 prompt/model/output_path)。
//!
//! ## 风险:High
//!
//! ① 付费 API(图像生成按张计费);② 写磁盘(落盘 output_path);③ 外发网络(provider 端点
//! + 图片 URL 下载)。统一 High 须人工批准。
//!
//! ## 安全
//!
//! 1. **provider 凭证**:从 db.ai_providers 筛 `enabled && provider_type=="openai_compat"` 的
//!    provider,经 `df_storage::secret::resolve_provider_secret` 解析 key(DB 优先 fallback keyring)。
//!    anthropic provider 不参与(图像端点是 OpenAI 风格)。
//! 2. **端点拼接**:`build_images_url` 智能 base_url(`/v1`/`/v4` 后缀直接补 `/images/generations`,
//!    否则补 `/v1/images/generations`),对齐 `model_fetch_helpers::build_models_url` 思路。
//! 3. **provider 端点域名**:base_url 来自 DB 用户配置(设置页可填任意 URL),provider 攻陷 /
//!    配置错误 / 恶意 base_url 即可打内网(127.0.0.1 / 169.254.169.254 元数据)。故 provider
//!    POST 端点同样走 SSRF 防护(validate_url + resolve_and_check_host + build_client),
//!    与图片 URL 下载同源,口径一致(2026-08-02 走查修复)。
//! 4. **图片 URL 下载**:图片 URL 来自 provider 响应,**可能被恶意 provider 篡改指向内网**
//!    (provider 域名虽可信但响应内容不可信),故下载图片 URL 复用 SSRF 防护
//!    (validate_url + resolve_and_check_host + build_client,与 download_file 同源)。
//! 5. **路径校验**:output_path 走 `resolve_workspace_path_with_allowed`(白名单 + symlink 逃逸
//!    拦截 + 路径遍历拒),与 write_file/download_file 同源;默认路径锚定首个持久授权目录。
//! 6. **大小上限**:50MB(图片够用),流式边下边累计,超限中止 + 删半成品(对齐 download_file)。

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use base64::{engine::general_purpose::STANDARD, Engine as _};
use serde_json::{json, Value};
use tokio::io::AsyncWriteExt;
use tokio::sync::RwLock;

use df_storage::crud::AiProviderRepo;
use df_storage::db::Database;
use df_storage::models::AiProviderRecord;
use df_storage::secret::resolve_provider_secret;

use crate::commands::ai::http::{
    build_client, execute_with_redirects, resolve_and_check_host, validate_url, MAX_REDIRECTS,
};
use crate::commands::ai::tool_registry::resolve_workspace_path_with_allowed;
use crate::state::AllowedDirs;

/// 下载图片大小硬上限(字节)。50MB —— 覆盖高分辨率图片(4K壁纸级 < 20MB),超限中止。
/// 图片 URL 来自 provider 响应可能被篡改指向超大文件,流式边下边累计防磁盘炸。
const MAX_IMAGE_BYTES: u64 = 50 * 1024 * 1024;

/// 生成图片 POST 请求超时(秒)。图像生成模型耗时较高(高分辨率 10-30s),给 120s 余量。
const GENERATE_TIMEOUT_SECS: u64 = 120;

/// 按 base64 长度估算解码后字节数(base64 每 4 字符编码 3 字节)。
///
/// 用于 b64_json 解码前的 OOM 防护:恶意 provider 返超长 b64_json(如 200MB base64 → ~150MB
/// 解码字节),若先 `STANDARD.decode` 全载入内存再检查上限,瞬时 OOM。故解码前先用本函数
/// 估算,超 MAX_IMAGE_BYTES 直接 bail 不解码。
///
/// 公式 `len * 3 / 4`:忽略 padding(`=`)和 whitespace 误差,估值**略小于实际** b64 长度对应
/// 的理论解码字节(实际含 padding 会更少)。安全方向:估算偏保守(偏小),配合解码后兜底校验
/// (`decoded.len() > MAX_IMAGE_BYTES`)双保险。`saturating_mul` 防 usize 溢出(64 位平台
/// b64.len() 不可能逼近 usize::MAX,但防御性编程)。
///
/// 抽为独立纯函数供单测覆盖边界(2026-08-02 走查修复,对齐 test-strategy 防回归)。
fn estimate_decoded_bytes(b64_len: usize) -> u64 {
    (b64_len as u64).saturating_mul(3) / 4
}

/// generate_image 工具 handler 入口(供 tools/generate_image.rs register 调用)。
///
/// 参数:
/// - prompt: 必填,图片描述(中英文均可,空串报错)
/// - model: 默认 `sensenova-u1-fast`(SenseNova U1 信息图生成,走 /v1/images/generations)
/// - size: 可选,图片尺寸字符串(SenseNova U1 常用 `2752x1536`/`1536x2752`,DALL-E 常用 `1024x1024`)
/// - n: 默认 1,生成数量,clamp [1, 4](多张图取首张 url/b64 落地,余下仅返回 url 列表)
/// - provider_id: 可选,指定 provider(多 provider 时);不传则按 model 名匹配厂商 host
/// - output_path: 可选,落盘路径;不传则默认 `{首个持久授权目录}/generated_images/{model}-{时间戳}.png`
///
/// 返回 {path, url(或 null 若仅 b64), model, provider_id, bytes_written, elapsed_ms}
///
/// 注:handler 接收 db + allowed_dirs(经 declare_tool! 闭包捕获 tuple Arc 传入)。
/// 路径校验在本函数内完成(需 allowed_dirs 快照,与 http_request output_file 同款)。
pub(crate) async fn execute_generate_image(
    args: Value,
    db: &Arc<Database>,
    allowed_dirs: &Arc<RwLock<AllowedDirs>>,
) -> anyhow::Result<Value> {
    let started = Instant::now();

    // ── 参数解析 ──
    let prompt = args.get("prompt").and_then(|v| v.as_str())
        .ok_or_else(|| anyhow::anyhow!("缺少 prompt 参数"))?
        .trim()
        .to_string();
    if prompt.is_empty() {
        anyhow::bail!("prompt 不能为空");
    }
    let model = args.get("model").and_then(|v| v.as_str())
        .unwrap_or("sensenova-u1-fast")
        .trim()
        .to_string();
    if model.is_empty() {
        anyhow::bail!("model 不能为空");
    }
    let size = args.get("size").and_then(|v| v.as_str())
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty());
    let n = args.get("n").and_then(|v| v.as_u64())
        .unwrap_or(1)
        .clamp(1, 4) as u32;
    let provider_id = args.get("provider_id").and_then(|v| v.as_str())
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty());
    let output_path_raw = args.get("output_path").and_then(|v| v.as_str())
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty());

    // ── 选 provider(enabled && openai_compat,按 model 名匹配厂商 host 兜底首个) ──
    let providers = AiProviderRepo::new(db).list_all().await?;
    let provider = select_image_provider(&providers, &model, provider_id.as_deref())?;
    let provider_id_resolved = provider.id.clone();

    // ── 解析 api_key(DB 优先 fallback keyring,空则报错) ──
    let api_key = resolve_provider_secret(&provider);
    let api_key = api_key.trim();
    if api_key.is_empty() {
        anyhow::bail!(
            "provider「{}」未配置 api_key,请在设置填入并保存",
            provider.name
        );
    }

    // ── 拼端点 URL ──
    let endpoint = build_images_url(&provider.base_url);

    // ── POST 请求 ──
    // base_url 来自 DB 用户配置(可填任意 URL),须走 SSRF 防护(协议白名单 + 私网 IP +
    // DNS resolve + 重定向每跳校验),与图片 URL 下载同源(2026-08-02 走查修复)。
    let body = {
        let mut m = serde_json::Map::new();
        m.insert("model".into(), json!(model));
        m.insert("prompt".into(), json!(prompt));
        m.insert("n".into(), json!(n));
        if let Some(sz) = &size {
            m.insert("size".into(), json!(sz));
        }
        Value::Object(m)
    };
    let body_str = serde_json::to_string(&body)
        .map_err(|e| anyhow::anyhow!("序列化请求 body 失败: {}", e))?;
    // SSRF 三层校验:词法 URL → DNS resolve IP(后续 build_client + execute_with_redirects
    // 对每跳重定向重复校验,防 302 绕过到内网)
    let (_scheme, host, port) = validate_url(&endpoint)?;
    resolve_and_check_host(&host, port).await?;
    // build_client 已含 timeout/connect_timeout + 关闭自动重定向(每跳手动校验)
    let client = build_client(Duration::from_secs(GENERATE_TIMEOUT_SECS))?;
    let mut headers = HashMap::new();
    headers.insert("Authorization".into(), format!("Bearer {}", api_key));
    headers.insert("Content-Type".into(), "application/json".into());
    let resp = execute_with_redirects(
        &client,
        reqwest::Method::POST,
        endpoint.clone(),
        &headers,
        &Some(body_str),
        MAX_REDIRECTS,
    )
    .await
    .map_err(|e| anyhow::anyhow!("调用图像生成端点失败 ({}): {}", endpoint, e))?;
    let status = resp.status();
    if !status.is_success() {
        // 错误响应读 body 摘要(截断 500 chars)帮助定位
        let err_text = resp.text().await.unwrap_or_default();
        let snippet: String = err_text.chars().take(500).collect();
        anyhow::bail!(
            "图像生成端点返回 HTTP {}: {}",
            status.as_u16(),
            snippet
        );
    }
    let resp_body: Value = resp.json().await
        .map_err(|e| anyhow::anyhow!("解析图像生成响应 JSON 失败: {}", e))?;

    // ── 解析响应:优先 data[0].url,次选 data[0].b64_json ──
    let data0 = resp_body.get("data").and_then(|d| d.get(0))
        .ok_or_else(|| {
            let snippet = body_snippet(&resp_body);
            anyhow::anyhow!("图像生成响应缺 data[0](响应摘要: {})", snippet)
        })?;
    let image_url_opt = data0.get("url").and_then(|u| u.as_str())
        .filter(|s| !s.trim().is_empty())
        .map(|s| s.trim().to_string());
    let b64_opt = data0.get("b64_json").and_then(|b| b.as_str())
        .filter(|s| !s.trim().is_empty())
        .map(|s| s.trim().to_string());

    if image_url_opt.is_none() && b64_opt.is_none() {
        let snippet = body_snippet(&resp_body);
        anyhow::bail!(
            "图像生成响应 data[0] 既无 url 也无 b64_json(响应摘要: {})",
            snippet
        );
    }

    // ── 解析 output_path(LLM 传则走白名单校验;不传则默认 generated_images/ 时间戳名) ──
    let snap = allowed_dirs.read().await.clone();
    let final_path = if let Some(raw) = &output_path_raw {
        resolve_workspace_path_with_allowed(raw, &snap)?
            .to_str()
            .ok_or_else(|| anyhow::anyhow!("output_path 含非法字符"))?
            .to_string()
    } else {
        // 默认:首个持久授权目录/generated_images/{model}-{YYYYMMDD-HHMMSS}.png
        let root = snap.first_persistent_dir()
            .ok_or_else(|| anyhow::anyhow!(
                "未传 output_path 且无持久授权目录,请先绑定项目目录或显式传 output_path"
            ))?;
        let ts = chrono::Local::now().format("%Y%m%d-%H%M%S");
        // model 含 / 等路径分隔符风险:替换为 - 防穿透到子目录
        let safe_model = model.replace(['/', '\\', ':', '*', '?', '"', '<', '>', '|'], "-");
        let fname = format!("{}-{}.png", safe_model, ts);
        root.join("generated_images").join(fname)
            .to_str()
            .ok_or_else(|| anyhow::anyhow!("默认 output_path 含非法字符"))?
            .to_string()
    };

    // ── 落盘:url 走 SSRF 防护下载流式写;b64 直接写解码后字节 ──
    let written_bytes = if let Some(image_url) = &image_url_opt {
        download_image_with_ssrf(image_url, &final_path).await?
    } else {
        // b64 路径:解码后直接写
        let b64 = b64_opt.as_ref().unwrap();
        // OOM 防护(2026-08-02 走查修复):解码前先按 base64 长度估算解码后字节数,超 MAX_IMAGE_BYTES
        // 直接 bail 不解码。恶意 provider 返超长 b64_json(如 200MB base64 → ~150MB 解码字节),
        // 若先 STANDARD.decode 全载入内存再检查,瞬时 OOM。详见 `estimate_decoded_bytes` 文档。
        let estimated_decoded = estimate_decoded_bytes(b64.len());
        if estimated_decoded > MAX_IMAGE_BYTES {
            anyhow::bail!(
                "b64_json 估算解码后约 {} 字节超过 {} 上限(原始 base64 长度 {})",
                estimated_decoded,
                MAX_IMAGE_BYTES,
                b64.len()
            );
        }
        let decoded = STANDARD.decode(b64)
            .map_err(|e| anyhow::anyhow!("b64_json 解码失败: {}", e))?;
        // 解码后再用实际长度兜底校验(防估算偏差,如 base64 含大量空白/padding)
        if (decoded.len() as u64) > MAX_IMAGE_BYTES {
            anyhow::bail!(
                "b64_json 解码后 {} 字节超过 {} 上限",
                decoded.len(),
                MAX_IMAGE_BYTES
            );
        }
        write_atomic(&final_path, &decoded).await?;
        decoded.len() as u64
    };

    let elapsed_ms = started.elapsed().as_millis() as u64;

    Ok(json!({
        "path": final_path,
        "url": image_url_opt,
        "model": model,
        "provider_id": provider_id_resolved,
        "bytes_written": written_bytes,
        "elapsed_ms": elapsed_ms,
    }))
}

/// 从 provider 列表筛选图像生成可用 provider。
///
/// 筛选条件:`enabled && provider_type == "openai_compat"`(图像端点是 OpenAI 风格,
/// anthropic provider 排除)。
///
/// 选择策略:
/// 1. LLM 传 `provider_id` → 按 id 精确匹配(找不到报错)。
/// 2. 否则按 model 名匹配厂商 host 提示:
///    - model 含 `sensenova` → 找 base_url 含 `sensenova`
///    - model 含 `dall-e` → 找 base_url 含 `openai`
///    - model 含 `google`/`imagen` → 找 base_url 含 `googleapis`
/// 3. host 命中 → 用首个匹配;host 未命中 → 兜底取筛选结果首个。
/// 4. 筛选后空 → 报错(无可用 openai_compat provider)。
fn select_image_provider(
    providers: &[AiProviderRecord],
    model: &str,
    provider_id: Option<&str>,
) -> anyhow::Result<AiProviderRecord> {
    // 筛 enabled + openai_compat
    let candidates: Vec<&AiProviderRecord> = providers.iter()
        .filter(|p| p.enabled && p.provider_type == "openai_compat")
        .collect();

    // 路径1:LLM 指定 provider_id,精确匹配
    if let Some(pid) = provider_id {
        let found = candidates.iter().find(|p| p.id == pid)
            .ok_or_else(|| anyhow::anyhow!(
                "未找到指定 provider(id={},且需 enabled + openai_compat),请检查 provider_id 参数",
                pid
            ))?;
        return Ok((*found).clone());
    }

    if candidates.is_empty() {
        anyhow::bail!("无可用 openai_compat provider,请先在设置配置支持图像生成的 provider");
    }

    // 路径2:按 model 名匹配厂商 host
    let model_lower = model.to_lowercase();
    let host_hint: Option<&str> = if model_lower.contains("sensenova") {
        Some("sensenova")
    } else if model_lower.contains("dall-e") {
        Some("openai")
    } else if model_lower.contains("google") || model_lower.contains("imagen") {
        Some("googleapis")
    } else {
        None
    };

    if let Some(hint) = host_hint {
        let hit = candidates.iter()
            .find(|p| p.base_url.to_lowercase().contains(hint));
        if let Some(p) = hit {
            return Ok((*p).clone());
        }
    }

    // 路径3:兜底首个候选
    Ok(candidates[0].clone())
}

/// 拼接图像生成端点 URL。
///
/// 规则(base_url 先去尾 `/`,对齐 `model_fetch_helpers::build_models_url` 思路,
/// 但目标是 `/images/generations`):
/// - 以 `/v1` 或 `/v4` 结尾 → 直接补 `/images/generations`(避免 `/v1/v1/images/...`)
/// - 否则 → 补 `/v1/images/generations`(最常见:用户填 `https://api.sensenova.cn`)
fn build_images_url(base_url: &str) -> String {
    let url = base_url.trim_end_matches('/');
    if url.ends_with("/v1") || url.ends_with("/v4") {
        format!("{url}/images/generations")
    } else {
        format!("{url}/v1/images/generations")
    }
}

/// 流式下载图片 URL 到 output_path,带 SSRF 防护 + 大小上限 + 原子写。
///
/// 图片 URL 来自 provider 响应可能被恶意 provider 篡改指向内网,故复用 SSRF 防护
/// (validate_url + resolve_and_check_host + build_client + execute_with_redirects)。
async fn download_image_with_ssrf(image_url: &str, output_path: &str) -> anyhow::Result<u64> {
    // SSRF 三层校验:词法 URL → DNS resolve IP → 每跳重定向重复
    let (_scheme, host, port) = validate_url(image_url)?;
    resolve_and_check_host(&host, port).await?;

    let client = build_client(Duration::from_secs(GENERATE_TIMEOUT_SECS))?;
    let resp = execute_with_redirects(
        &client,
        reqwest::Method::GET,
        image_url.to_string(),
        &HashMap::new(), // 图片下载不带自定义头
        &None,
        MAX_REDIRECTS,
    )
    .await?;

    if !resp.status().is_success() {
        anyhow::bail!(
            "下载图片失败:HTTP {} {}({})",
            resp.status().as_u16(),
            resp.status().canonical_reason().unwrap_or(""),
            image_url
        );
    }

    // 流式写文件(避免大图 OOM),边下边累计,超 MAX_IMAGE_BYTES 立即中止 + 删半成品。
    // 父目录不存在自动创建。原子写:tmp → rename(同目录不跨卷)。
    let target = std::path::Path::new(output_path);
    if let Some(parent) = target.parent() {
        tokio::fs::create_dir_all(parent).await
            .map_err(|e| anyhow::anyhow!("创建目录失败: {}", e))?;
    }

    let tmp = format!("{}.tmp-img", output_path);
    let mut file = tokio::fs::File::create(&tmp).await
        .map_err(|e| anyhow::anyhow!("创建临时文件失败 ({}): {}", tmp, e))?;

    let mut stream = resp.bytes_stream();
    let mut written: u64 = 0;
    let mut exceeded = false;
    use futures::StreamExt;
    while let Some(chunk) = stream.next().await {
        let chunk = chunk.map_err(|e| anyhow::anyhow!("读取图片流失败: {}", e))?;
        written += chunk.len() as u64;
        if written > MAX_IMAGE_BYTES {
            exceeded = true;
            break;
        }
        file.write_all(&chunk).await
            .map_err(|e| anyhow::anyhow!("写入图片失败: {}", e))?;
    }
    file.flush().await
        .map_err(|e| anyhow::anyhow!("flush 图片文件失败: {}", e))?;
    drop(file);

    if exceeded {
        let _ = tokio::fs::remove_file(&tmp).await;
        anyhow::bail!(
            "图片超过 {} 字节上限,已中止并清理临时文件",
            MAX_IMAGE_BYTES
        );
    }

    if let Err(e) = tokio::fs::rename(&tmp, output_path).await {
        let _ = tokio::fs::remove_file(&tmp).await;
        return Err(anyhow::anyhow!("原子替换失败: {}", e));
    }

    Ok(written)
}

/// 原子写字节到指定路径(tmp → rename,同目录不跨卷)。
/// b64 路径用此函数(已全部解码在内存,无需流式)。
async fn write_atomic(output_path: &str, data: &[u8]) -> anyhow::Result<()> {
    let target = std::path::Path::new(output_path);
    if let Some(parent) = target.parent() {
        tokio::fs::create_dir_all(parent).await
            .map_err(|e| anyhow::anyhow!("创建目录失败: {}", e))?;
    }
    let tmp = format!("{}.tmp-img", output_path);
    tokio::fs::write(&tmp, data).await
        .map_err(|e| {
            let _ = std::fs::remove_file(&tmp);
            anyhow::anyhow!("写入图片失败: {}", e)
        })?;
    tokio::fs::rename(&tmp, output_path).await
        .map_err(|e| {
            let _ = std::fs::remove_file(&tmp);
            anyhow::anyhow!("原子替换失败: {}", e)
        })?;
    Ok(())
}

/// 响应 body 截断为最多 500 chars 的字符串(用于错误信息),供定位用。
fn body_snippet(v: &Value) -> String {
    let s = v.to_string();
    s.chars().take(500).collect()
}

// ============================================================
// 单元测试
//
// 覆盖:
// ① build_images_url:端点拼接规则(纯函数,零依赖)
// ② select_image_provider:provider 筛选逻辑(纯函数,内存构造 record)
// ③ estimate_decoded_bytes:b64 解码字节数估算(纯函数,OOM 防护核心)
// ④ handler 参数边界(prompt 缺失/空、provider_id 不存在)
// ⑤ handler SSRF 防护(endpoint 走 validate_url 前置拒,不发真实请求,CI 可跑)
// 真实 API 调用走 #[ignore](CI 无凭证/无网时跳过,本地手跑)。
// ============================================================

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

    // ── build_images_url:端点拼接规则 ──

    #[test]
    fn build_images_url_v1_suffix() {
        // /v1 结尾 → 直接补 /images/generations
        assert_eq!(
            build_images_url("https://api.openai.com/v1"),
            "https://api.openai.com/v1/images/generations"
        );
        // 不应产生 /v1/v1/images/...
        assert!(!build_images_url("https://api.openai.com/v1").contains("/v1/v1/"));
    }

    #[test]
    fn build_images_url_v4_suffix() {
        // GLM v4 风格 base
        assert_eq!(
            build_images_url("https://open.bigmodel.cn/api/paas/v4"),
            "https://open.bigmodel.cn/api/paas/v4/images/generations"
        );
    }

    #[test]
    fn build_images_url_bare_base() {
        // 用户只填根 → 补 /v1/images/generations
        assert_eq!(
            build_images_url("https://api.sensenova.cn"),
            "https://api.sensenova.cn/v1/images/generations"
        );
    }

    #[test]
    fn build_images_url_strips_trailing_slash() {
        assert_eq!(
            build_images_url("https://api.x.com/v1/"),
            "https://api.x.com/v1/images/generations"
        );
    }

    // ── select_image_provider:筛选 + 匹配逻辑 ──

    /// 构造测试用 AiProviderRecord(只填筛选关键字段)。
    fn mk_provider(id: &str, enabled: bool, ptype: &str, base_url: &str) -> AiProviderRecord {
        AiProviderRecord {
            id: id.into(),
            name: id.into(),
            provider_type: ptype.into(),
            api_key: String::new(),
            base_url: base_url.into(),
            default_model: String::new(),
            models: None,
            model_configs: Vec::new(),
            is_default: false,
            config: None,
            created_at: "0".into(),
            updated_at: "0".into(),
            enabled,
            weight: 50,
        }
    }

    #[test]
    fn select_provider_no_candidates_errors() {
        // 全 disabled 或非 openai_compat → 报错
        let providers = vec![
            mk_provider("p1", false, "openai_compat", "https://x"),
            mk_provider("p2", true, "anthropic", "https://y"),
        ];
        let err = select_image_provider(&providers, "any-model", None).unwrap_err();
        assert!(format!("{}", err).contains("无可用"));
    }

    #[test]
    fn select_provider_provider_id_match() {
        let providers = vec![
            mk_provider("p1", true, "openai_compat", "https://api.sensenova.cn"),
            mk_provider("p2", true, "openai_compat", "https://api.openai.com"),
        ];
        let got = select_image_provider(&providers, "any", Some("p2")).unwrap();
        assert_eq!(got.id, "p2");
    }

    #[test]
    fn select_provider_provider_id_not_found_errors() {
        let providers = vec![mk_provider("p1", true, "openai_compat", "https://x")];
        let err = select_image_provider(&providers, "any", Some("nonexistent")).unwrap_err();
        assert!(format!("{}", err).contains("未找到指定 provider"));
    }

    #[test]
    fn select_provider_host_match_sensenova() {
        // model 含 sensenova → 匹配 base_url 含 sensenova 的 provider
        let providers = vec![
            mk_provider("openai", true, "openai_compat", "https://api.openai.com"),
            mk_provider("sense", true, "openai_compat", "https://api.sensenova.cn"),
        ];
        let got = select_image_provider(&providers, "sensenova-u1-fast", None).unwrap();
        assert_eq!(got.id, "sense");
    }

    #[test]
    fn select_provider_host_match_dalle() {
        // model 含 dall-e → 匹配 base_url 含 openai 的 provider
        let providers = vec![
            mk_provider("glm", true, "openai_compat", "https://open.bigmodel.cn"),
            mk_provider("oai", true, "openai_compat", "https://api.openai.com"),
        ];
        let got = select_image_provider(&providers, "dall-e-3", None).unwrap();
        assert_eq!(got.id, "oai");
    }

    #[test]
    fn select_provider_host_unmatched_falls_back_first() {
        // model 名无厂商提示 → 兜底取筛选结果首个
        let providers = vec![
            mk_provider("first", true, "openai_compat", "https://api.unknown.com"),
            mk_provider("second", true, "openai_compat", "https://api.other.com"),
        ];
        let got = select_image_provider(&providers, "my-custom-model", None).unwrap();
        assert_eq!(got.id, "first");
    }

    // ── estimate_decoded_bytes:b64 解码字节数估算(OOM 防护核心纯函数) ──

    #[test]
    fn estimate_decoded_bytes_normal_small() {
        // 1MB base64 → 估算 ~0.75MB(base64 每 4 字符编码 3 字节)
        let one_mb = 1024 * 1024;
        assert_eq!(estimate_decoded_bytes(one_mb), (one_mb as u64) * 3 / 4);
        assert_eq!(estimate_decoded_bytes(one_mb), 786_432); // 1MB * 3 / 4
    }

    #[test]
    fn estimate_decoded_bytes_over_limit() {
        // 200MB 解码后场景:b64 长约 267MB → 估算 ~200MB > 50MB 上限,OOM 防护应挡
        let two_hundred_mb_b64 = (200 * 1024 * 1024) as usize; // 假设这是 b64 长度
        let estimated = estimate_decoded_bytes(two_hundred_mb_b64);
        assert!(
            estimated > MAX_IMAGE_BYTES,
            "200MB b64 估算 {} 应超 {} 上限,OOM 防护靠此拒",
            estimated,
            MAX_IMAGE_BYTES
        );
    }

    #[test]
    fn estimate_decoded_bytes_boundary_at_max() {
        // 边界:估算正好等于 MAX_IMAGE_BYTES → 比较用 `>`(超过才拒),正好等于算合法通过
        // 反推 b64 长度:要使 len*3/4 == MAX_IMAGE_BYTES,len = MAX_IMAGE_BYTES*4/3
        let b64_len = (MAX_IMAGE_BYTES as u128 * 4 / 3) as usize;
        let estimated = estimate_decoded_bytes(b64_len);
        // 估算值应 ≤ MAX_IMAGE_BYTES(因整数除法向下取整,可能略小于上限)
        assert!(
            estimated <= MAX_IMAGE_BYTES,
            "估算 {} 应 ≤ 上限 {}（正好或略小于，`>` 比较下不拒）",
            estimated,
            MAX_IMAGE_BYTES
        );
    }

    #[test]
    fn estimate_decoded_bytes_just_over_max() {
        // 边界:估算值 = MAX_IMAGE_BYTES + 1 → 应被拒
        // 反推 b64 长度:要使 len*3/4 == MAX+1,len = (MAX+1)*4/3
        let target = MAX_IMAGE_BYTES + 1;
        let b64_len = (target as u128 * 4 / 3) as usize;
        let estimated = estimate_decoded_bytes(b64_len);
        // 因整数除法,估算可能略小于 target,但应保证 > MAX_IMAGE_BYTES
        // 用更精确的 b64_len:target*4/3 + 4(向上补一个 base64 块),保证估算严格 > MAX
        let b64_len_safe = ((target as u128 * 4 + 2) / 3) as usize;
        let estimated_safe = estimate_decoded_bytes(b64_len_safe);
        assert!(
            estimated_safe > MAX_IMAGE_BYTES,
            "b64_len={} 估算 {} 应严格 > 上限 {}（超限应拒）",
            b64_len_safe,
            estimated_safe,
            MAX_IMAGE_BYTES
        );
    }

    #[test]
    fn estimate_decoded_bytes_saturating_no_overflow() {
        // 极大值(u64::MAX 的 usize 等价)→ saturating_mul 防 panic,返回 u64 内合理值
        // 注:usize 在 64 位平台 = u64,saturating_mul(u64::MAX, 3) 会 saturate 到 u64::MAX
        let huge = usize::MAX;
        let estimated = estimate_decoded_bytes(huge);
        // saturating_mul(MAX, 3) = MAX,再 / 4 = MAX/4,不 panic 即通过
        assert_eq!(estimated, u64::MAX / 4);
    }

    #[test]
    fn estimate_decoded_bytes_zero() {
        // 空输入 → 0(边界,虽 b64_opt 已 filter 空串不进此路径,纯函数仍应优雅)
        assert_eq!(estimate_decoded_bytes(0), 0);
    }

    // ── handler 参数边界(走 db 需 in-memory db) ──

    #[tokio::test]
    async fn handler_missing_prompt_errors() {
        // 缺 prompt → 参数解析阶段报错(不发请求)
        let db = Arc::new(Database::open_in_memory().await.unwrap());
        let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
        let args = json!({ "model": "sensenova-u1-fast" });
        let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
        assert!(format!("{}", err).contains("prompt"));
    }

    #[tokio::test]
    async fn handler_empty_prompt_errors() {
        let db = Arc::new(Database::open_in_memory().await.unwrap());
        let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
        let args = json!({ "prompt": "   " });
        let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
        assert!(format!("{}", err).contains("prompt 不能为空"));
    }

    #[tokio::test]
    async fn handler_no_provider_errors() {
        // 无 openai_compat provider → 报错(不发请求)
        let db = Arc::new(Database::open_in_memory().await.unwrap());
        let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
        let args = json!({ "prompt": "一只猫" });
        let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
        assert!(format!("{}", err).contains("无可用") || format!("{}", err).contains("provider"));
    }

    // ── handler SSRF 防护(endpoint 走 validate_url 前置拒,不发真实请求) ──
    //
    // base_url 来自 DB 用户配置(设置页可填任意 URL),provider 攻陷 / 配置错误 / 恶意 base_url
    // 即可打内网。endpoint POST 须走 SSRF 防护(validate_url + resolve_and_check_host),与图片
    // URL 下载同源(2026-08-02 走查修复)。本组测试构造 enabled + openai_compat + api_key 非空
    // 的 provider 写入 in-memory db,handler 走到 validate_url 阶段拒,不发真实请求(CI 可跑)。
    // 风格对齐 http.rs 的 test_handler_rejects_localhost_url / test_handler_rejects_private_ip_url。

    /// 构造带 api_key 的 openai_compat provider(复用 mk_provider 但填 api_key,跳过 keyring)。
    fn mk_provider_with_key(id: &str, base_url: &str) -> AiProviderRecord {
        AiProviderRecord {
            api_key: "sk-test-dummy-not-real".into(),
            ..mk_provider(id, true, "openai_compat", base_url)
        }
    }

    #[tokio::test]
    async fn handler_rejects_localhost_base_url() {
        // base_url = http://localhost:8080 → validate_url L117-119 拒 localhost(SSRF 防护)
        // 在 build_client 前拒,不发真实请求
        let db = Arc::new(Database::open_in_memory().await.unwrap());
        AiProviderRepo::new(&db)
            .insert(mk_provider_with_key("p1", "http://localhost:8080"))
            .await
            .unwrap();
        let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
        let args = json!({ "prompt": "一只猫", "provider_id": "p1" });
        let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
        let msg = format!("{}", err);
        assert!(
            msg.contains("localhost") || msg.contains("SSRF"),
            "期望 SSRF 拒绝 localhost,实际: {}",
            msg
        );
    }

    #[tokio::test]
    async fn handler_rejects_private_ip_base_url() {
        // base_url 含 169.254.169.254(云元数据服务,SSRF 头号目标)→ validate_url 字面量 IP 私网拒
        let db = Arc::new(Database::open_in_memory().await.unwrap());
        AiProviderRepo::new(&db)
            .insert(mk_provider_with_key("p1", "http://169.254.169.254/latest"))
            .await
            .unwrap();
        let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
        let args = json!({ "prompt": "一只猫", "provider_id": "p1" });
        let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
        let msg = format!("{}", err);
        assert!(
            msg.contains("私网") || msg.contains("SSRF") || msg.contains("169.254"),
            "期望 SSRF 拒绝元数据 IP,实际: {}",
            msg
        );
    }

    #[tokio::test]
    async fn handler_rejects_loopback_base_url() {
        // base_url = http://127.0.0.1 → validate_url 字面量 IP 环回拒
        let db = Arc::new(Database::open_in_memory().await.unwrap());
        AiProviderRepo::new(&db)
            .insert(mk_provider_with_key("p1", "http://127.0.0.1:9000"))
            .await
            .unwrap();
        let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
        let args = json!({ "prompt": "一只猫", "provider_id": "p1" });
        let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
        let msg = format!("{}", err);
        assert!(
            msg.contains("私网") || msg.contains("SSRF") || msg.contains("127.0.0.1"),
            "期望 SSRF 拒绝环回 IP,实际: {}",
            msg
        );
    }

    #[tokio::test]
    async fn handler_rejects_non_http_scheme_base_url() {
        // base_url = file:///etc/passwd → validate_url L104-109 拒非 http 协议
        let db = Arc::new(Database::open_in_memory().await.unwrap());
        AiProviderRepo::new(&db)
            .insert(mk_provider_with_key("p1", "file:///etc/passwd"))
            .await
            .unwrap();
        let allowed = Arc::new(RwLock::new(AllowedDirs::default()));
        let args = json!({ "prompt": "一只猫", "provider_id": "p1" });
        let err = execute_generate_image(args, &db, &allowed).await.unwrap_err();
        let msg = format!("{}", err);
        assert!(
            msg.contains("协议"),
            "期望拒绝非 http 协议,实际: {}",
            msg
        );
    }
}
