use crate::{auth::jwt::validate_jwt, db::models::downstream_api_key::DownstreamApiKey, AppState};
use axum::{
    extract::{DefaultBodyLimit, State},
    http::HeaderMap,
    response::sse::{Event, Sse},
    routing::{get, post},
    Json, Router,
};
use futures::stream::Stream;
use regex::Regex;
use reqwest::Client;
use serde::Deserialize;
use serde_json::{json, Value};
use std::convert::Infallible;
use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
use tokio_stream::{wrappers::UnboundedReceiverStream, StreamExt};
use tracing::info;

const SYSTEM_PROMPT: &str = r#"你是 BoosAPI 的 AI 视频生成助手，帮用户通过 AI 制作视频。

## 规则
1. 先用中文与用户讨论脚本构思、视觉风格、角色设定
2. 用户确认后（说"开始制作"或"生成"），再分步输出结构化内容
3. **分两步走：第一步输出角色+场景+TTS，第二步输出 WORKFLOW**（不要一次全出）

## 结构化标记格式
严格按照以下格式输出，系统会自动解析并生成图片/TTS：

【CHARACTER】
name: 林澈
description: 女宇航员，30岁，冷静坚毅，短发收在头盔内
prompt_front: cinematic full body front view of a female astronaut, white EVA suit, photorealistic
prompt_back: cinematic full body back view of a female astronaut, life support backpack
prompt_left: cinematic full body left side view of a female astronaut, side utility pouches
prompt_right: cinematic full body right side view of a female astronaut, helmet visor reflection
【/CHARACTER】

【SCENE】
id: 1
description: 飞船驾驶舱内景
prompt: interior of a sci-fi spacecraft cockpit, holographic displays, cold blue lighting, cinematic
【/SCENE】

【TTS】
character: 林澈
text: 我们准备出舱检查
emotion: 平静
【/TTS】

【WORKFLOW】
```json
{"comfyui_workflow": {}}
```
【/WORKFLOW】

## 关键要求
- prompt_* 用英文，同时包含主体 + 视角 + 细节 + 画质标签
- 不要一次输出太多，先输角色和场景，下次再输工作流
- 如果没有用户明确要求，切勿自作主张输出全部内容
"#;

const AGENT_CHAT_BODY_LIMIT_BYTES: usize = 64 * 1024 * 1024;

pub fn routes() -> Router<AppState> {
    Router::new()
        .route("/api/comfyui-agent/status", get(status_handler))
        .route("/api/comfyui-agent/models", get(models_handler))
        .route("/api/comfyui-agent/image-models", get(image_models_handler))
        .route("/api/comfyui-agent/chat", post(chat_handler))
        .layer(DefaultBodyLimit::max(AGENT_CHAT_BODY_LIMIT_BYTES))
}

#[derive(Debug, Deserialize)]
struct AgentChatRequest {
    messages: Vec<AgentMessage>,
    model: Option<String>,
    #[serde(rename = "imageModel")]
    image_model: Option<String>,
}

#[derive(Debug, Deserialize)]
struct AgentMessage {
    role: String,
    content: String,
    #[serde(default)]
    attachments: Vec<AgentAttachment>,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct AgentAttachment {
    name: String,
    mime_type: Option<String>,
    #[serde(default, rename = "mediaKind")]
    media_kind: Option<String>,
    #[serde(default, rename = "referenceLabel")]
    reference_label: Option<String>,
    text: Option<String>,
    data_url: Option<String>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum AgentMode {
    /// Route through this server's own `/v1/*` using the caller's downstream key.
    Local,
    /// Legacy fallback: external upstream configured via `BOOSAPI_UPSTREAM_BASE` / `BOOSAPI_UPSTREAM_KEY`.
    FallbackRemote,
}

#[derive(Debug, Clone)]
struct AgentAuthContext {
    user_id: Option<i64>,
    role: String,
    outbound_key: String,
    key_name: String,
}

impl AgentMode {
    fn as_str(self) -> &'static str {
        match self {
            AgentMode::Local => "local",
            AgentMode::FallbackRemote => "fallback_remote",
        }
    }
}

fn detect_agent_mode() -> AgentMode {
    let base_override = std::env::var("BOOSAPI_UPSTREAM_BASE")
        .ok()
        .map(|v| v.trim().to_string())
        .filter(|v| !v.is_empty());
    let key_override = std::env::var("BOOSAPI_UPSTREAM_KEY")
        .ok()
        .map(|v| v.trim().to_string())
        .filter(|v| !v.is_empty());
    if base_override.is_some() || key_override.is_some() {
        AgentMode::FallbackRemote
    } else {
        AgentMode::Local
    }
}

fn local_base(state: &AppState) -> String {
    let host = if state.config.server.host == "0.0.0.0" || state.config.server.host == "::" {
        "127.0.0.1".to_string()
    } else {
        state.config.server.host.clone()
    };
    format!("http://{}:{}/v1", host, state.config.server.port)
}

/// Resolve `(mode, api_base, api_key, key_name)` for the agent's outbound calls.
///
/// Defaults to `AgentMode::Local`, which targets this server's own `/v1`
/// endpoints with the current caller's downstream key — every call then runs
/// through the regular accounting / routing pipeline and is billed/logged
/// against that user key.
///
/// Setting `BOOSAPI_UPSTREAM_BASE` or `BOOSAPI_UPSTREAM_KEY` switches to
/// `FallbackRemote` for operators that still want to point straight at an
/// external aggregator.
fn upstream_config(
    state: &AppState,
    auth: &AgentAuthContext,
) -> (AgentMode, String, String, String) {
    match detect_agent_mode() {
        AgentMode::Local => {
            let base = local_base(state);
            (
                AgentMode::Local,
                base,
                auth.outbound_key.clone(),
                auth.key_name.clone(),
            )
        }
        AgentMode::FallbackRemote => {
            let base = std::env::var("BOOSAPI_UPSTREAM_BASE")
                .ok()
                .map(|value| value.trim().trim_end_matches('/').to_string())
                .filter(|value| !value.is_empty())
                .unwrap_or_else(|| "https://api.boos.lat/v1".to_string());
            let key = std::env::var("BOOSAPI_UPSTREAM_KEY")
                .ok()
                .map(|value| value.trim().to_string())
                .filter(|value| !value.is_empty())
                .unwrap_or_else(|| auth.outbound_key.clone());
            let key_name = if std::env::var("BOOSAPI_UPSTREAM_KEY")
                .ok()
                .map(|value| !value.trim().is_empty())
                .unwrap_or(false)
            {
                "external".to_string()
            } else {
                auth.key_name.clone()
            };
            (AgentMode::FallbackRemote, base, key, key_name)
        }
    }
}

fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
    let auth_header = headers
        .get("authorization")
        .and_then(|value| value.to_str().ok())
        .unwrap_or_default();

    auth_header
        .strip_prefix("Bearer ")
        .or_else(|| auth_header.strip_prefix("bearer "))
        .map(str::trim)
        .filter(|token| !token.is_empty())
}

async fn user_downstream_key(state: &AppState, user_id: i64) -> Result<DownstreamApiKey, String> {
    sqlx::query_as::<_, DownstreamApiKey>(
        "SELECT *
         FROM downstream_api_keys
         WHERE user_id = ?
           AND COALESCE(enabled, 1) = 1
           AND (
             expires_at IS NULL
             OR expires_at = ''
             OR datetime(expires_at) > datetime('now')
           )
         ORDER BY updated_at DESC, created_at DESC, id DESC
         LIMIT 1",
    )
    .bind(user_id)
    .fetch_optional(&*state.db)
    .await
    .map_err(|err| err.to_string())?
    .ok_or_else(|| "当前用户没有可用的 API Key，请先在用户中心创建或启用 API Key".to_string())
}

async fn resolve_agent_auth(
    headers: &HeaderMap,
    state: &AppState,
) -> Result<AgentAuthContext, (axum::http::StatusCode, Json<Value>)> {
    let token = extract_bearer_token(headers).ok_or_else(|| {
        (
            axum::http::StatusCode::UNAUTHORIZED,
            Json(json!({ "error": "Unauthorized" })),
        )
    })?;

    let proxy_token = state.runtime_proxy_token();
    if token == state.config.auth.auth_token || token == proxy_token {
        return Ok(AgentAuthContext {
            user_id: None,
            role: "admin".to_string(),
            outbound_key: proxy_token,
            key_name: "global-proxy".to_string(),
        });
    }

    let claims = validate_jwt(&state.config.auth.jwt_secret, token).map_err(|_| {
        (
            axum::http::StatusCode::UNAUTHORIZED,
            Json(json!({ "error": "Unauthorized" })),
        )
    })?;
    if claims.role != "admin" && claims.role != "user" {
        return Err((
            axum::http::StatusCode::FORBIDDEN,
            Json(json!({ "error": "Forbidden" })),
        ));
    }

    let user_id = claims.sub.parse::<i64>().map_err(|_| {
        (
            axum::http::StatusCode::UNAUTHORIZED,
            Json(json!({ "error": "Invalid user ID in token" })),
        )
    })?;

    match user_downstream_key(state, user_id).await {
        Ok(key) => Ok(AgentAuthContext {
            user_id: Some(user_id),
            role: claims.role,
            outbound_key: key.key,
            key_name: key.name,
        }),
        Err(_) if claims.role == "admin" => Ok(AgentAuthContext {
            user_id: Some(user_id),
            role: claims.role,
            outbound_key: state.runtime_proxy_token(),
            key_name: "global-proxy".to_string(),
        }),
        Err(message) => Err((
            axum::http::StatusCode::FORBIDDEN,
            Json(json!({ "error": message })),
        )),
    }
}

fn is_exact_model_pattern(value: &str) -> bool {
    let normalized = value.trim();
    !normalized.is_empty()
        && !normalized.to_ascii_lowercase().starts_with("re:")
        && !normalized.contains('*')
        && !normalized.contains('?')
}

fn is_image_model_pattern(value: &str) -> bool {
    if !is_exact_model_pattern(value) {
        return false;
    }

    let normalized = value.trim().to_ascii_lowercase();
    let image_keywords = [
        "gpt-image",
        "dall-e",
        "dalle",
        "midjourney",
        "recraft",
        "kolors",
        "cogview",
        "seedream",
        "sdxl",
        "stable-diffusion",
        "stablediffusion",
        "pixart",
        "playground",
        "firefly",
        "imagen",
        "flux",
        "image",
        "img",
        "canvas",
        "painting",
        "draw",
    ];
    if image_keywords
        .iter()
        .any(|keyword| normalized.contains(keyword))
    {
        return true;
    }

    let chat_keywords = [
        "claude", "chatgpt", "gpt-4", "gpt-3", "gpt-5", "o1", "o3", "o4", "gemini", "deepseek",
        "qwen", "llama", "mistral", "command", "kimi", "doubao", "glm", "hunyuan", "spark",
        "baichuan", "ernie", "wenxin", "grok",
    ];
    !chat_keywords
        .iter()
        .any(|keyword| normalized.contains(keyword))
}

async fn status_handler(
    State(state): State<AppState>,
    headers: HeaderMap,
) -> Result<Json<Value>, (axum::http::StatusCode, Json<Value>)> {
    let auth = resolve_agent_auth(&headers, &state).await?;
    let (mode, upstream_base, _, key_name) = upstream_config(&state, &auth);
    let upstream_label = match mode {
        AgentMode::Local => "BoosAPI 本地路由".to_string(),
        AgentMode::FallbackRemote => {
            let parsed = reqwest::Url::parse(&upstream_base).ok();
            parsed
                .and_then(|url| url.host_str().map(ToOwned::to_owned))
                .unwrap_or_else(|| "外部上游".to_string())
        }
    };
    Ok(Json(json!({
        "success": true,
        "model": state.config.comfyui.agent_model,
        "mode": mode.as_str(),
        "keyName": key_name,
        "userId": auth.user_id,
        "role": auth.role,
        "apiBase": upstream_base,
        "upstreamBase": upstream_base,
        "upstreamLabel": upstream_label,
        "comfyui": state.comfyui.status_response().await,
    })))
}

async fn models_handler(
    State(state): State<AppState>,
    headers: HeaderMap,
) -> Result<Json<Value>, (axum::http::StatusCode, Json<Value>)> {
    let _auth = resolve_agent_auth(&headers, &state).await?;

    let rows: Vec<String> = sqlx::query_scalar(
        "SELECT model_pattern
         FROM token_routes
         WHERE COALESCE(enabled, 1) = 1
         ORDER BY model_pattern COLLATE NOCASE",
    )
    .fetch_all(&*state.db)
    .await
    .map_err(|err| {
        (
            axum::http::StatusCode::INTERNAL_SERVER_ERROR,
            Json(json!({ "error": err.to_string() })),
        )
    })?;

    let mut seen = std::collections::HashSet::new();
    let models: Vec<String> = rows
        .into_iter()
        .map(|value| value.trim().to_string())
        .filter(|value| is_exact_model_pattern(value))
        .filter(|value| seen.insert(value.clone()))
        .collect();

    Ok(Json(json!({
        "success": true,
        "models": models,
        "defaultModel": state.config.comfyui.agent_model,
    })))
}

async fn image_models_handler(
    State(state): State<AppState>,
    headers: HeaderMap,
) -> Result<Json<Value>, (axum::http::StatusCode, Json<Value>)> {
    let _auth = resolve_agent_auth(&headers, &state).await?;

    let rows: Vec<String> = sqlx::query_scalar(
        "SELECT model_pattern
         FROM token_routes
         WHERE COALESCE(enabled, 1) = 1
         ORDER BY model_pattern COLLATE NOCASE",
    )
    .fetch_all(&*state.db)
    .await
    .map_err(|err| {
        (
            axum::http::StatusCode::INTERNAL_SERVER_ERROR,
            Json(json!({ "error": err.to_string() })),
        )
    })?;

    let mut seen = std::collections::HashSet::new();
    let models: Vec<String> = rows
        .into_iter()
        .map(|value| value.trim().to_string())
        .filter(|value| is_image_model_pattern(value))
        .filter(|value| seen.insert(value.clone()))
        .collect();

    Ok(Json(json!({
        "success": true,
        "models": models,
        "defaultModel": models
            .iter()
            .find(|model| model.to_ascii_lowercase().contains("image"))
            .or_else(|| models.first())
            .cloned()
            .unwrap_or_else(|| "gpt-image-2".to_string()),
    })))
}

#[derive(Debug, Default)]
struct ParsedSections {
    characters: Vec<CharacterSection>,
    scenes: Vec<SceneSection>,
    tts_list: Vec<TtsSection>,
    workflow_json: Option<String>,
}

#[derive(Debug, Default)]
struct CharacterSection {
    name: String,
    description: String,
    prompt_front: String,
    prompt_back: String,
    prompt_left: String,
    prompt_right: String,
}

#[derive(Debug)]
struct SceneSection {
    id: u32,
    description: String,
    prompt: String,
}

#[derive(Debug)]
struct TtsSection {
    character: String,
    text: String,
    emotion: String,
}

fn parse_sections(full_text: &str) -> ParsedSections {
    let mut sections = ParsedSections::default();

    // Parse CHARACTER blocks
    let char_re = Regex::new(r"(?s)【CHARACTER】\s*(.*?)【/CHARACTER】").unwrap();
    for cap in char_re.captures_iter(full_text) {
        let block = cap.get(1).unwrap().as_str();
        let mut char_section = CharacterSection::default();
        for line in block.lines() {
            let line = line.trim();
            let line_lower = line.to_lowercase();
            if let Some(val) = line
                .strip_prefix("name:")
                .or_else(|| line.strip_prefix("name :"))
            {
                char_section.name = val.trim().to_string();
            } else if let Some(val) = line
                .strip_prefix("description:")
                .or_else(|| line.strip_prefix("description :"))
            {
                char_section.description = val.trim().to_string();
            } else if line_lower.starts_with("prompt_front:")
                || line_lower.starts_with("prompt_front :")
            {
                let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
                char_section.prompt_front = val.to_string();
            } else if line_lower.starts_with("prompt_back:")
                || line_lower.starts_with("prompt_back :")
            {
                let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
                char_section.prompt_back = val.to_string();
            } else if line_lower.starts_with("prompt_left:")
                || line_lower.starts_with("prompt_left :")
            {
                let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
                char_section.prompt_left = val.to_string();
            } else if line_lower.starts_with("prompt_right:")
                || line_lower.starts_with("prompt_right :")
            {
                let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
                char_section.prompt_right = val.to_string();
            }
        }
        if !char_section.name.is_empty() {
            sections.characters.push(char_section);
        }
    }

    // Parse SCENE blocks
    let scene_re = Regex::new(r"(?s)【SCENE】\s*(.*?)【/SCENE】").unwrap();
    for cap in scene_re.captures_iter(full_text) {
        let block = cap.get(1).unwrap().as_str();
        let mut scene = SceneSection {
            id: 0,
            description: String::new(),
            prompt: String::new(),
        };
        for line in block.lines() {
            let line = line.trim();
            let line_lower = line.to_lowercase();
            if line_lower.starts_with("id:") || line_lower.starts_with("id :") {
                let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
                scene.id = val.parse().unwrap_or(0);
            } else if let Some(val) = line
                .strip_prefix("description:")
                .or_else(|| line.strip_prefix("description :"))
            {
                scene.description = val.trim().to_string();
            } else if line_lower.starts_with("prompt:") || line_lower.starts_with("prompt :") {
                let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
                scene.prompt = val.to_string();
            }
        }
        if !scene.prompt.is_empty() {
            sections.scenes.push(scene);
        }
    }

    // Parse TTS blocks
    let tts_re = Regex::new(r"(?s)【TTS】\s*(.*?)【/TTS】").unwrap();
    for cap in tts_re.captures_iter(full_text) {
        let block = cap.get(1).unwrap().as_str();
        let mut tts = TtsSection {
            character: String::new(),
            text: String::new(),
            emotion: String::new(),
        };
        for line in block.lines() {
            let line = line.trim();
            let line_lower = line.to_lowercase();
            if line_lower.starts_with("character:") || line_lower.starts_with("character :") {
                let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
                tts.character = val.to_string();
            } else if line_lower.starts_with("text:") || line_lower.starts_with("text :") {
                let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
                tts.text = val.to_string();
            } else if line_lower.starts_with("emotion:") || line_lower.starts_with("emotion :") {
                let val = line.split_once(':').map(|x| x.1).unwrap_or("").trim();
                tts.emotion = val.to_string();
            }
        }
        if !tts.text.is_empty() {
            sections.tts_list.push(tts);
        }
    }

    // Parse WORKFLOW blocks
    let wf_re = Regex::new(r"(?s)【WORKFLOW】\s*```(?:json)?\s*(.*?)```\s*【/WORKFLOW】").unwrap();
    if let Some(cap) = wf_re.captures(full_text) {
        let wf = cap.get(1).unwrap().as_str().trim().to_string();
        sections.workflow_json = Some(wf);
    }

    sections
}

async fn call_image_gen(
    api_base: &str,
    api_key: &str,
    image_model: &str,
    prompt: &str,
) -> Result<String, String> {
    let client = Client::new();
    let url = format!("{}/images/generations", api_base.trim_end_matches('/'));

    let resp = client
        .post(&url)
        .bearer_auth(api_key)
        .json(&json!({
            "model": image_model,
            "prompt": prompt,
            "n": 1,
            "size": "1024x1024"
        }))
        .send()
        .await
        .map_err(|e| format!("image request failed: {}", e))?;

    let status = resp.status();
    let data: Value = resp
        .json()
        .await
        .map_err(|e| format!("image response parse failed: {}", e))?;

    if !status.is_success() {
        return Err(format!("image upstream returned {}: {}", status, data));
    }

    Ok(data
        .get("data")
        .and_then(Value::as_array)
        .and_then(|items| items.first())
        .and_then(|item| item.get("url"))
        .and_then(Value::as_str)
        .unwrap_or("")
        .to_string())
}

async fn call_tts(
    api_base: &str,
    api_key: &str,
    text: &str,
    voice: &str,
) -> Result<Vec<u8>, String> {
    let client = Client::new();
    let url = format!("{}/audio/speech", api_base.trim_end_matches('/'));

    let resp = client
        .post(&url)
        .bearer_auth(api_key)
        .json(&json!({
            "model": "tts-1",
            "input": text,
            "voice": voice
        }))
        .send()
        .await
        .map_err(|e| format!("TTS request failed: {}", e))?;

    let status = resp.status();
    if !status.is_success() {
        let body = resp.text().await.unwrap_or_default();
        return Err(format!("TTS upstream returned {}: {}", status, body));
    }

    let bytes = resp
        .bytes()
        .await
        .map_err(|e| format!("TTS response read failed: {}", e))?;

    Ok(bytes.to_vec())
}

fn send_sse(tx: &UnboundedSender<Result<Event, Infallible>>, payload: Value) -> bool {
    tx.send(Ok(Event::default().data(payload.to_string())))
        .is_ok()
}

fn string_field(value: &Value, keys: &[&str]) -> Option<String> {
    keys.iter()
        .find_map(|key| value.get(*key).and_then(Value::as_str))
        .map(ToOwned::to_owned)
        .filter(|text| !text.is_empty())
}

fn extract_reasoning_delta(parsed: &Value) -> Option<String> {
    if let Some(delta) = parsed
        .get("choices")
        .and_then(Value::as_array)
        .and_then(|choices| choices.first())
        .and_then(|choice| choice.get("delta"))
    {
        if let Some(text) = string_field(
            delta,
            &[
                "reasoning_content",
                "reasoning",
                "thinking",
                "thinking_delta",
                "reasoning_delta",
            ],
        ) {
            return Some(text);
        }

        if let Some(text) = delta
            .get("reasoning")
            .and_then(|value| value.get("content").or_else(|| value.get("text")))
            .and_then(Value::as_str)
            .map(ToOwned::to_owned)
            .filter(|text| !text.is_empty())
        {
            return Some(text);
        }
    }

    let event_type = parsed
        .get("type")
        .and_then(Value::as_str)
        .unwrap_or_default();
    match event_type {
        "response.reasoning.delta"
        | "response.reasoning_text.delta"
        | "response.reasoning_summary_text.delta" => {
            string_field(parsed, &["delta", "text", "content"])
        }
        "content_block_delta" => {
            let delta = parsed.get("delta")?;
            match delta
                .get("type")
                .and_then(Value::as_str)
                .unwrap_or_default()
            {
                "thinking_delta" => string_field(delta, &["thinking", "text", "delta"]),
                "signature_delta" => None,
                other if other.contains("thinking") || other.contains("reasoning") => {
                    string_field(delta, &["thinking", "text", "content", "delta"])
                }
                _ => None,
            }
        }
        other if other.contains("reasoning") || other.contains("thinking") => string_field(
            parsed,
            &["delta", "text", "content", "thinking", "reasoning"],
        ),
        _ => None,
    }
}

fn attachment_summary(attachment: &AgentAttachment) -> String {
    let mime = attachment
        .mime_type
        .as_deref()
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .unwrap_or("application/octet-stream");
    let name = attachment.name.trim();
    let name = if name.is_empty() { "untitled" } else { name };
    let reference = attachment
        .reference_label
        .as_deref()
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .unwrap_or("@素材");
    let media_kind = attachment
        .media_kind
        .as_deref()
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .unwrap_or("file");
    let header =
        format!("\n\n[素材引用: {reference}, 名称: {name}, 类型: {mime}, 分类: {media_kind}]");

    if let Some(text) = attachment
        .text
        .as_deref()
        .map(str::trim)
        .filter(|value| !value.is_empty())
    {
        format!("{header}\n{text}")
    } else {
        header
    }
}

fn agent_message_to_upstream(message: AgentMessage) -> Value {
    let role = message.role;
    let mut text = message.content;
    let mut content_parts = Vec::new();

    for attachment in &message.attachments {
        let mime = attachment
            .mime_type
            .as_deref()
            .unwrap_or_default()
            .to_ascii_lowercase();
        if mime.starts_with("image/") {
            if let Some(data_url) = attachment
                .data_url
                .as_deref()
                .map(str::trim)
                .filter(|value| value.starts_with("data:image/"))
            {
                if !content_parts
                    .iter()
                    .any(|part: &Value| part["type"] == "text")
                {
                    content_parts.push(json!({
                        "type": "text",
                        "text": text,
                    }));
                    text = String::new();
                }
                content_parts.push(json!({
                    "type": "image_url",
                    "image_url": { "url": data_url },
                }));
                continue;
            }
        }

        text.push_str(&attachment_summary(attachment));
    }

    if !content_parts.is_empty() {
        if !text.trim().is_empty() {
            content_parts.push(json!({
                "type": "text",
                "text": text,
            }));
        }
        json!({
            "role": role,
            "content": content_parts,
        })
    } else {
        json!({
            "role": role,
            "content": text,
        })
    }
}

async fn chat_handler(
    State(state): State<AppState>,
    headers: HeaderMap,
    Json(body): Json<AgentChatRequest>,
) -> Result<Sse<impl Stream<Item = Result<Event, Infallible>>>, (axum::http::StatusCode, Json<Value>)>
{
    let auth = resolve_agent_auth(&headers, &state).await?;

    if body.messages.is_empty()
        || body.messages.iter().any(|message| {
            message.role.trim().is_empty()
                || (message.content.trim().is_empty() && message.attachments.is_empty())
        })
    {
        return Err((
            axum::http::StatusCode::BAD_REQUEST,
            Json(json!({ "error": "messages is required" })),
        ));
    }

    let requested_model = body
        .model
        .as_deref()
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(ToOwned::to_owned);
    let requested_image_model = body
        .image_model
        .as_deref()
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(ToOwned::to_owned);

    let messages: Vec<Value> = body
        .messages
        .into_iter()
        .map(agent_message_to_upstream)
        .collect();
    let (mode, api_base, api_key, key_name) = upstream_config(&state, &auth);
    info!(
        mode = mode.as_str(),
        upstream_base = %api_base,
        key_name = %key_name,
        user_id = ?auth.user_id,
        "comfyui-agent: resolved upstream",
    );
    let agent_model = requested_model.unwrap_or_else(|| state.config.comfyui.agent_model.clone());
    let image_model = requested_image_model.unwrap_or_else(|| "gpt-image-2".to_string());
    let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();

    tokio::spawn(async move {
        let upstream_url = format!("{}/chat/completions", api_base.trim_end_matches('/'));

        let mut upstream_messages = Vec::with_capacity(messages.len() + 1);
        upstream_messages.push(json!({
            "role": "system",
            "content": SYSTEM_PROMPT,
        }));
        upstream_messages.extend(messages);

        let client = Client::new();

        info!("comfyui-agent: sending chat request to upstream");

        let response = match client
            .post(&upstream_url)
            .bearer_auth(&api_key)
            .json(&json!({
                "model": agent_model,
                "messages": upstream_messages,
                "temperature": 0.7,
                "stream": true,
            }))
            .send()
            .await
        {
            Ok(response) => response,
            Err(e) => {
                send_sse(
                    &tx,
                    json!({
                        "type": "text",
                        "content": format!("服务暂时不可用，请稍后重试。({})", e),
                    }),
                );
                send_sse(&tx, json!({ "type": "done" }));
                return;
            }
        };

        if !response.status().is_success() {
            send_sse(
                &tx,
                json!({
                    "type": "text",
                    "content": "服务暂时不可用，请稍后重试",
                }),
            );
            send_sse(&tx, json!({ "type": "done" }));
            return;
        }

        let mut stream = response.bytes_stream();
        let mut buffer: Vec<u8> = Vec::new();
        let mut accumulated_text = String::new();

        // Phase 1: Stream the LLM text response
        loop {
            let chunk = tokio::select! {
                _ = tx.closed() => {
                    return;
                }
                chunk = stream.next() => chunk,
            };

            let Some(chunk) = chunk else {
                break;
            };

            let chunk = match chunk {
                Ok(chunk) => chunk,
                Err(error) => {
                    send_sse(
                        &tx,
                        json!({
                            "type": "text",
                            "content": format!("\n\n[错误: {}]", error),
                        }),
                    );
                    break;
                }
            };

            buffer.extend_from_slice(&chunk);

            while let Some(pos) = buffer.iter().position(|byte| *byte == b'\n') {
                let mut line_bytes: Vec<u8> = buffer.drain(..=pos).collect();

                if line_bytes.last() == Some(&b'\n') {
                    line_bytes.pop();
                }
                if line_bytes.last() == Some(&b'\r') {
                    line_bytes.pop();
                }

                let line = String::from_utf8_lossy(&line_bytes);
                let line = line.trim().to_string();

                if !line.starts_with("data: ") {
                    continue;
                }

                let data = line.trim_start_matches("data: ").trim().to_string();

                if data == "[DONE]" {
                    continue;
                }

                let Ok(parsed) = serde_json::from_str::<Value>(&data) else {
                    continue;
                };

                if let Some(reasoning) = extract_reasoning_delta(&parsed) {
                    if !send_sse(
                        &tx,
                        json!({
                            "type": "reasoning",
                            "content": reasoning,
                        }),
                    ) {
                        return;
                    }
                }

                let delta = parsed
                    .get("choices")
                    .and_then(Value::as_array)
                    .and_then(|choices| choices.first())
                    .and_then(|choice| choice.get("delta"))
                    .cloned()
                    .unwrap_or(Value::Null);

                let content = delta
                    .get("content")
                    .and_then(Value::as_str)
                    .unwrap_or_default()
                    .to_string();

                if !content.is_empty() {
                    accumulated_text.push_str(&content);
                    if !send_sse(
                        &tx,
                        json!({
                            "type": "text",
                            "content": content,
                        }),
                    ) {
                        return;
                    }
                }
            }
        }

        // Phase 2: Parse structured sections from accumulated text
        info!("comfyui-agent: text streaming complete, parsing sections");
        let sections = parse_sections(&accumulated_text);

        let has_structured_data = !sections.characters.is_empty()
            || !sections.scenes.is_empty()
            || !sections.tts_list.is_empty()
            || sections.workflow_json.is_some();

        if !has_structured_data {
            send_sse(&tx, json!({ "type": "done" }));
            return;
        }

        // Send a processing indicator
        send_sse(
            &tx,
            json!({
                "type": "text",
                "content": "\n\n🔄 正在处理结构化数据..."
            }),
        );

        // Generate images for each character (4 angles each)
        for (char_idx, character) in sections.characters.iter().enumerate() {
            // Send character_prompts event
            send_sse(
                &tx,
                json!({
                    "type": "character_prompts",
                    "character": character.name,
                }),
            );

            let angles = [
                ("正面", &character.prompt_front, "front"),
                ("背面", &character.prompt_back, "back"),
                ("左侧", &character.prompt_left, "left"),
                ("右侧", &character.prompt_right, "right"),
            ];

            for (angle_label, prompt, angle_key) in &angles {
                if prompt.is_empty() {
                    continue;
                }
                info!(
                    "comfyui-agent: generating image for {} - {}",
                    character.name, angle_label
                );

                // Add a short delay between image gen requests to avoid rate limiting
                if char_idx > 0 || *angle_key != "front" {
                    tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
                }

                match call_image_gen(&api_base, &api_key, &image_model, prompt).await {
                    Ok(url) if !url.is_empty() => {
                        send_sse(
                            &tx,
                            json!({
                                "type": "image",
                                "url": url,
                                "character": character.name,
                                "angle": angle_key,
                            }),
                        );
                    }
                    _ => {
                        send_sse(
                            &tx,
                            json!({
                                "type": "text",
                                "content": format!("\n⚠️ {} [{}] 图片生成失败，请重试", character.name, angle_label),
                            }),
                        );
                    }
                }
            }
        }

        // Generate images for each scene
        for scene in &sections.scenes {
            info!(
                "comfyui-agent: generating scene image for scene {}",
                scene.id
            );

            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;

            match call_image_gen(&api_base, &api_key, &image_model, &scene.prompt).await {
                Ok(url) if !url.is_empty() => {
                    send_sse(
                        &tx,
                        json!({
                            "type": "scene",
                            "id": scene.id,
                            "description": scene.description,
                            "prompt": scene.prompt,
                            "url": url,
                        }),
                    );
                }
                _ => {
                    // Still send scene info even without image
                    send_sse(
                        &tx,
                        json!({
                            "type": "scene",
                            "id": scene.id,
                            "description": scene.description,
                            "prompt": scene.prompt,
                        }),
                    );
                }
            }
        }

        // Generate TTS for each dialogue
        for tts in &sections.tts_list {
            info!("comfyui-agent: generating TTS for {}", tts.character);

            // Choose voice based on emotion
            let voice = match tts.emotion.as_str() {
                "激动" | "愤怒" => "nova",
                "悲伤" | "恐惧" => "shimmer",
                _ => "alloy",
            };

            match call_tts(&api_base, &api_key, &tts.text, voice).await {
                Ok(audio_bytes) => {
                    // For SSE, we'll send a data URL for small audio, or a placeholder
                    // Since SSE can't easily send binary, we store to file and send URL
                    let audio_path = format!(
                        "/tmp/comfyui_tts_{}_{}.mp3",
                        tts.character.replace(' ', "_"),
                        std::time::SystemTime::now()
                            .duration_since(std::time::UNIX_EPOCH)
                            .unwrap_or_default()
                            .as_millis()
                    );
                    if let Ok(()) = tokio::fs::write(&audio_path, &audio_bytes).await {
                        send_sse(
                            &tx,
                            json!({
                                "type": "tts",
                                "character": tts.character,
                                "text": tts.text,
                                "url": format!("/audio/{}", audio_path.trim_start_matches("/tmp/comfyui_tts_")),
                            }),
                        );
                    }
                }
                Err(e) => {
                    info!("comfyui-agent: TTS generation failed: {}", e);
                    send_sse(
                        &tx,
                        json!({
                            "type": "text",
                            "content": format!("\n⚠️ {} 的语音生成失败", tts.character),
                        }),
                    );
                }
            }

            tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
        }

        // Send ComfyUI workflow if present
        if let Some(wf_json) = &sections.workflow_json {
            info!("comfyui-agent: sending workflow JSON");
            send_sse(
                &tx,
                json!({
                    "type": "comfyui_workflow",
                    "content": wf_json,
                }),
            );
        }

        // Signal done with summary
        let mut summary = String::from("\n\n✅ **处理完成！**");
        if !sections.characters.is_empty() {
            summary.push_str(&format!("\n- 角色: {} 个", sections.characters.len()));
        }
        if !sections.scenes.is_empty() {
            summary.push_str(&format!("\n- 场景: {} 个", sections.scenes.len()));
        }
        if !sections.tts_list.is_empty() {
            summary.push_str(&format!("\n- TTS: {} 段", sections.tts_list.len()));
        }
        if sections.workflow_json.is_some() {
            summary.push_str("\n- ComfyUI 工作流: ✅ 已生成");
        }
        send_sse(&tx, json!({ "type": "text", "content": summary }));
        send_sse(&tx, json!({ "type": "done" }));
    });

    Ok(Sse::new(UnboundedReceiverStream::new(rx)))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig};
    use crate::services::builtin_keys::BuiltinKeys;
    use axum::http::HeaderValue;
    use sqlx::sqlite::SqlitePoolOptions;
    use tokio::sync::Mutex;
    use wiremock::matchers::{header, method, path};
    use wiremock::{Mock, MockServer, ResponseTemplate};

    /// `BOOSAPI_UPSTREAM_*` is process-global. Serialize tests that read or
    /// mutate it so cargo's parallel runner does not race their assertions.
    /// Uses `tokio::sync::Mutex` so the guard can safely span `.await` points
    /// (mock server setup, agent worker poll loop) without tripping
    /// `clippy::await_holding_lock`.
    static ENV_GUARD: Mutex<()> = Mutex::const_new(());

    async fn make_state(host: &str, port: u16) -> AppState {
        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect("sqlite::memory:")
            .await
            .unwrap();
        AppState::new(
            AppConfig {
                server: ServerConfig {
                    host: host.to_string(),
                    port,
                    log_level: "info".into(),
                    cors_origins: vec![],
                },
                db: DbConfig {
                    url: "sqlite::memory:".into(),
                    max_connections: 1,
                    ssl: false,
                    runtime_db_url: None,
                    runtime_db_dialect: None,
                    runtime_db_ssl: None,
                },
                auth: AuthConfig {
                    jwt_secret: "secret".into(),
                    jwt_expiry_hours: 1,
                    auth_token: "auth".into(),
                    proxy_token: "proxy".into(),
                },
                proxy: ProxyConfig {
                    upstream_timeout_secs: 30,
                    first_byte_timeout_secs: 30,
                    max_retries: 3,
                    proxy_log_retention_days: 7,
                    template_proxy_url: None,
                },
                comfyui: Default::default(),
            },
            pool,
            BuiltinKeys {
                comfyui_agent: (1, "sk-comfyui-test".into()),
            },
            crate::services::runtime_auth::RuntimeAuthSettings::from_token_for_test("proxy"),
        )
    }

    fn clear_upstream_env() {
        std::env::remove_var("BOOSAPI_UPSTREAM_BASE");
        std::env::remove_var("BOOSAPI_UPSTREAM_KEY");
    }

    fn auth_headers() -> HeaderMap {
        let mut headers = HeaderMap::new();
        headers.insert("Authorization", HeaderValue::from_static("Bearer auth"));
        headers
    }

    async fn call_status(state: AppState) -> Value {
        let result = status_handler(State(state), auth_headers()).await;
        result.expect("status_handler returned error").0
    }

    #[test]
    fn extracts_chat_reasoning_delta() {
        let parsed = json!({
            "choices": [{
                "delta": {
                    "reasoning_content": "正在拆解镜头..."
                }
            }]
        });
        assert_eq!(
            extract_reasoning_delta(&parsed).as_deref(),
            Some("正在拆解镜头...")
        );
    }

    #[test]
    fn extracts_responses_reasoning_delta() {
        let parsed = json!({
            "type": "response.reasoning_summary_text.delta",
            "delta": "先分析角色和场景。"
        });
        assert_eq!(
            extract_reasoning_delta(&parsed).as_deref(),
            Some("先分析角色和场景。")
        );
    }

    #[test]
    fn extracts_anthropic_thinking_delta() {
        let parsed = json!({
            "type": "content_block_delta",
            "index": 0,
            "delta": {
                "type": "thinking_delta",
                "thinking": "用户需要一个视频分镜。"
            }
        });
        assert_eq!(
            extract_reasoning_delta(&parsed).as_deref(),
            Some("用户需要一个视频分镜。")
        );
    }

    /// 3.5 — `/api/comfyui-agent/status` reflects the env state in one shot:
    /// no env → local; whitespace/empty env → still local (treated as unset);
    /// real env → fallback_remote with trimmed `upstreamBase`.
    #[tokio::test]
    async fn status_reflects_env_state_across_three_scenarios() {
        let _g = ENV_GUARD.lock().await;
        clear_upstream_env();

        // Scenario A: env unset → local; 0.0.0.0 host normalizes to 127.0.0.1.
        {
            let state = make_state("0.0.0.0", 7860).await;
            let body = call_status(state).await;
            assert_eq!(body["mode"], "local");
            assert_eq!(body["keyName"], "global-proxy");
            assert_eq!(body["upstreamBase"], "http://127.0.0.1:7860/v1");
        }

        // Scenario B: env set to whitespace/empty → treated as unset → local.
        std::env::set_var("BOOSAPI_UPSTREAM_BASE", "   ");
        std::env::set_var("BOOSAPI_UPSTREAM_KEY", "");
        {
            let state = make_state("127.0.0.1", 9000).await;
            let body = call_status(state).await;
            assert_eq!(body["mode"], "local");
            assert_eq!(body["keyName"], "global-proxy");
            assert_eq!(body["upstreamBase"], "http://127.0.0.1:9000/v1");
        }

        // Scenario C: env set → fallback_remote, base trimmed of trailing slash.
        std::env::set_var("BOOSAPI_UPSTREAM_BASE", "https://api.example.com/v1/");
        std::env::remove_var("BOOSAPI_UPSTREAM_KEY");
        {
            let state = make_state("0.0.0.0", 7860).await;
            let body = call_status(state).await;
            assert_eq!(body["mode"], "fallback_remote");
            assert_eq!(body["keyName"], "global-proxy");
            assert_eq!(body["upstreamBase"], "https://api.example.com/v1");
        }

        clear_upstream_env();
    }

    /// 3.6 — In Local mode the agent's chat call must hit the local
    /// `/v1/chat/completions` endpoint with the current caller's downstream
    /// key as the bearer token.
    #[tokio::test]
    async fn agent_local_mode_routes_chat_to_local_upstream() {
        let _g = ENV_GUARD.lock().await;
        clear_upstream_env();

        let mock = MockServer::start().await;
        let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n\n";
        Mock::given(method("POST"))
            .and(path("/v1/chat/completions"))
            .and(header("authorization", "Bearer proxy"))
            .respond_with(
                ResponseTemplate::new(200)
                    .insert_header("content-type", "text/event-stream")
                    .set_body_string(sse),
            )
            .mount(&mock)
            .await;

        let addr = mock.address();
        let host = addr.ip().to_string();
        let port = addr.port();

        let state = make_state(&host, port).await;

        let result = chat_handler(
            State(state),
            auth_headers(),
            Json(AgentChatRequest {
                messages: vec![AgentMessage {
                    role: "user".to_string(),
                    content: "hello".to_string(),
                    attachments: Vec::new(),
                }],
                model: None,
                image_model: None,
            }),
        )
        .await;
        assert!(result.is_ok(), "chat_handler should accept the request");
        // The Sse stream is dropped here; the worker still runs the upstream
        // POST before its streaming loop notices `tx.closed()`.

        // Wait up to ~2s for the spawned worker's POST to land at wiremock.
        let mut hit = false;
        for _ in 0..40 {
            tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
            let received = mock.received_requests().await.unwrap_or_default();
            if received
                .iter()
                .any(|r| r.url.path() == "/v1/chat/completions")
            {
                hit = true;
                break;
            }
        }
        assert!(
            hit,
            "expected /v1/chat/completions to be hit by the local-mode agent"
        );
    }
}