PicoBot/src/gateway/ws.rs

use super::GatewayState;
use crate::agent::{AgentError, CompositeSystemPromptProvider};
use crate::bus::{InboundMessage, MediaItem};
use crate::command::adapter::{InputAdapter, OutputAdapter};
use crate::command::adapters::websocket::{WebSocketInputAdapter, WebSocketOutputAdapter};
use crate::command::context::CommandContext;
use crate::command::handler::CommandRouter;
use crate::command::handlers::get_current::GetCurrentSessionCommandHandler;
use crate::command::handlers::help::HelpCommandHandler;
use crate::command::handlers::list_channels::ListChannelsCommandHandler;
use crate::command::handlers::list_sessions::ListSessionsCommandHandler;
use crate::command::handlers::list_sessions_by_channel::ListSessionsByChannelCommandHandler;
use crate::command::handlers::list_topics::ListTopicsCommandHandler;
use crate::command::handlers::load_task_messages::LoadTaskMessagesCommandHandler;
use crate::command::handlers::load_topic::LoadTopicCommandHandler;
use crate::command::handlers::save_session::SaveSessionCommandHandler;
use crate::command::handlers::session::SessionCommandHandler;
use crate::command::handlers::switch_topic::SwitchTopicCommandHandler;
use crate::gateway::agent_prompt_provider::AgentPromptProvider;
use crate::protocol::{WsInbound, WsOutbound, MediaSummary, parse_inbound, serialize_outbound};
use crate::skills::SkillPromptProvider;
use axum::extract::State;
use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade};
use axum::response::Response;
use base64::{Engine as _, engine::general_purpose::STANDARD};
use futures_util::{SinkExt, StreamExt};
use std::collections::HashMap;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::sync::mpsc;

const CLI_CHANNEL_NAME: &str = "cli";

/// Default media directory for WebSocket uploads
fn default_ws_media_dir() -> PathBuf {
    let home = dirs::home_dir().unwrap_or_else(|| PathBuf::from("."));
    home.join(".picobot").join("media").join("ws")
}

/// Build a unique filename for media upload
fn build_media_filename(media_type: &str, file_name: Option<&str>) -> String {
    if let Some(file_name) = file_name {
        let sanitized: String = file_name
            .chars()
            .map(|ch| match ch {
                '/' | '\\' | ':' | '*' | '?' | '"' | '<' | '>' | '|' => '_',
                _ => ch,
            })
            .collect();
        if !sanitized.trim().is_empty() {
            return format!("{}_{}", uuid::Uuid::new_v4(), sanitized);
        }
    }
    format!("{}_{}", media_type, uuid::Uuid::new_v4())
}

/// Process attachments with base64 content: save to local file and return MediaItem with correct path
/// Keeps content_base64 for frontend display/download
fn process_attachments_with_base64(attachments: Vec<MediaSummary>) -> Result<Vec<MediaItem>, AgentError> {
    if attachments.is_empty() {
        return Ok(Vec::new());
    }

    let media_dir = default_ws_media_dir();
    std::fs::create_dir_all(&media_dir)
        .map_err(|error| AgentError::Other(format!("Failed to create media dir: {}", error)))?;

    attachments
        .into_iter()
        .map(|att| {
            // If content_base64 exists, save to file and update path
            if let Some(base64_content) = &att.content_base64 {
                let decoded = STANDARD
                    .decode(base64_content)
                    .map_err(|error| AgentError::Other(format!("Failed to decode base64: {}", error)))?;

                let filename = build_media_filename(&att.media_type, att.file_name.as_deref());
                let file_path = media_dir.join(&filename);

                std::fs::write(&file_path, decoded)
                    .map_err(|error| AgentError::Other(format!("Failed to write media file: {}", error)))?;

                tracing::info!(
                    filename = %filename,
                    media_type = %att.media_type,
                    file_path = %file_path.to_string_lossy(),
                    "Saved WebSocket media to local file"
                );

                Ok(MediaItem {
                    path: file_path.to_string_lossy().to_string(),
                    media_type: att.media_type,
                    mime_type: att.mime_type,
                    original_key: None,
                    // Keep content_base64 for frontend display/download
                    content_base64: att.content_base64,
                    file_name: att.file_name,
                })
            } else {
                // No base64 content, keep original path (should already be valid)
                Ok(MediaItem {
                    path: att.path,
                    media_type: att.media_type,
                    mime_type: att.mime_type,
                    original_key: None,
                    content_base64: None,
                    file_name: att.file_name,
                })
            }
        })
        .collect()
}

pub async fn ws_handler(ws: WebSocketUpgrade, State(state): State<Arc<GatewayState>>) -> Response {
    ws.on_upgrade(|socket| async {
        handle_socket(socket, state).await;
    })
}

async fn handle_socket(ws: WebSocket, state: Arc<GatewayState>) {
    let (sender, receiver) = mpsc::channel::<WsOutbound>(100);

    let cli_sessions = state.session_manager.cli_sessions();
    let store = state.session_manager.store();

    // 1. 先查询 websocket 通道的 Sessions
    let websocket_sessions = store.list_sessions("websocket", false)
        .unwrap_or_default();

    // 2. 如果没有，自动创建一个默认 Session
    let initial_record = if websocket_sessions.is_empty() {
        match cli_sessions.create_with_channel("websocket", Some("默认会话")) {
            Ok(record) => record,
            Err(e) => {
                tracing::error!(error = %e, "Failed to create initial WebSocket session");
                return;
            }
        }
    } else {
        // 使用最新的 Session
        websocket_sessions[0].clone()
    };

    let runtime_session_id = uuid::Uuid::new_v4().to_string();
    let mut current_session_id = initial_record.id.clone();
    let mut current_topic_id: Option<String> = None;
    state
        .channel_manager
        .cli_channel()
        .register_connection(
            current_session_id.clone(),
            runtime_session_id.clone(),
            sender.clone(),
        )
        .await;
    tracing::info!(runtime_session_id = %runtime_session_id, session_id = %current_session_id, "WebSocket session established");

    let _ = sender
        .send(WsOutbound::SessionEstablished {
            session_id: current_session_id.clone(),
        })
        .await;

    // 连接建立后立即发送通道列表
    let channels = ListChannelsCommandHandler::get_default_channels();
    let _ = sender
        .send(WsOutbound::ChannelList { channels })
        .await;

    // 3. 重新查询 websocket 通道的 Session 列表（包含刚创建的）
    let final_sessions = store.list_sessions("websocket", false)
        .unwrap_or_default();

    tracing::info!("Sending {} websocket sessions to client", final_sessions.len());
    for s in &final_sessions {
        tracing::info!("  - {}: {} (channel: {})", s.id, s.title, s.channel_name);
    }

    let session_summaries: Vec<crate::protocol::SessionSummary> = final_sessions
        .into_iter()
        .map(|s| crate::protocol::SessionSummary {
            session_id: s.id,
            title: s.title,
            channel_name: s.channel_name,
            chat_id: s.chat_id,
            message_count: s.message_count,
            last_active_at: s.last_active_at,
            archived_at: s.archived_at,
        })
        .collect();

    let _ = sender
        .send(WsOutbound::SessionList {
            sessions: session_summaries,
            current_session_id: Some(current_session_id.clone()),
            channel_name: Some("websocket".to_string()),
        })
        .await;

    let (mut ws_sender, mut ws_receiver) = ws.split();

    let mut receiver = receiver;
    let session_id_for_sender = runtime_session_id.clone();
    tokio::spawn(async move {
        while let Some(msg) = receiver.recv().await {
            if let Ok(text) = serialize_outbound(&msg) {
                if ws_sender.send(WsMessage::Text(text.into())).await.is_err() {
                    #[cfg(debug_assertions)]
                    tracing::debug!(session_id = %session_id_for_sender, "WebSocket send error");
                    break;
                }
            }
        }
    });

    while let Some(msg) = ws_receiver.next().await {
        match msg {
            Ok(WsMessage::Text(text)) => {
                let text = text.to_string();
                match parse_inbound(&text) {
                    Ok(inbound) => {
                        if let Err(e) = handle_inbound(
                            &state,
                            &sender,
                            &runtime_session_id,
                            &mut current_session_id,
                            &mut current_topic_id,
                            inbound,
                        )
                        .await
                        {
                            tracing::warn!(error = %e, session_id = %current_session_id, "Failed to handle inbound message");
                            let _ = sender
                                .send(WsOutbound::Error {
                                    code: "SESSION_ERROR".to_string(),
                                    message: e.to_string(),
                                })
                                .await;
                        }
                    }
                    Err(e) => {
                        tracing::warn!(error = %e, "Failed to parse inbound message");
                        let _ = sender
                            .send(WsOutbound::Error {
                                code: "PARSE_ERROR".to_string(),
                                message: e.to_string(),
                            })
                            .await;
                    }
                }
            }
            Ok(WsMessage::Close(_)) | Err(_) => {
                #[cfg(debug_assertions)]
                tracing::debug!(session_id = %runtime_session_id, "WebSocket closed");
                break;
            }
            _ => {}
        }
    }

    state
        .channel_manager
        .cli_channel()
        .unregister_connection(&runtime_session_id)
        .await;
    tracing::info!(session_id = %runtime_session_id, current_session_id = %current_session_id, "CLI session ended");
}


async fn handle_inbound(
    state: &Arc<GatewayState>,
    sender: &mpsc::Sender<WsOutbound>,
    runtime_session_id: &str,
    current_session_id: &mut String,
    current_topic_id: &mut Option<String>,
    inbound: WsInbound,
) -> Result<(), AgentError> {
    match inbound {
        WsInbound::Message {
            content,
            attachments,
            chat_id,
            sender_id,
            ..
        } => {
            let chat_id = chat_id.unwrap_or_else(|| current_session_id.clone());
            let sender_id = resolve_ws_sender_id(sender_id.as_deref(), runtime_session_id);

            state
                .channel_manager
                .cli_channel()
                .register_connection(
                    chat_id.clone(),
                    runtime_session_id.to_string(),
                    sender.clone(),
                )
                .await;

            // Process attachments: save base64 content to local files and build MediaItems with correct paths
            let media = process_attachments_with_base64(attachments)?;

            state
                .bus
                .publish_inbound(InboundMessage {
                    channel: CLI_CHANNEL_NAME.to_string(),
                    sender_id,
                    chat_id,
                    content,
                    timestamp: current_timestamp(),
                    media,
                    metadata: HashMap::new(),
                    forwarded_metadata: HashMap::new(),
                })
                .await
                .map_err(|error| AgentError::Other(error.to_string()))?;

            Ok(())
        }
        WsInbound::Command { payload } => {
            // 使用 Command 系统处理命令
            let input_adapter = WebSocketInputAdapter::new();
            let output_adapter = WebSocketOutputAdapter::new();

            // 解析命令
            let adapter_ctx = crate::command::context::AdapterContext::new("websocket")
                .with_session_id(current_session_id.as_str());

            let cmd = match input_adapter.try_parse(&payload, adapter_ctx) {
                Ok(Some(cmd)) => cmd,
                Ok(None) => {
                    // 不是命令，返回错误
                    let _ = sender
                        .send(WsOutbound::Error {
                            code: "INVALID_COMMAND".to_string(),
                            message: "Invalid command payload".to_string(),
                        })
                        .await;
                    return Ok(());
                }
                Err(e) => {
                    let _ = sender
                        .send(WsOutbound::Error {
                            code: "PARSE_ERROR".to_string(),
                            message: e.to_string(),
                        })
                        .await;
                    return Ok(());
                }
            };

            // 创建命令路由器
            let _cli_sessions = state.session_manager.cli_sessions();
            let store = state.session_manager.store();
            let skills = state.session_manager.skills();
            let provider_config = state.config.get_provider_config("default")
                .map_err(|e| AgentError::Other(e.to_string()))?;
            let prompt_repository = state.session_manager.store().clone();

            let system_prompt_provider: Arc<dyn crate::agent::SystemPromptProvider> = Arc::new(CompositeSystemPromptProvider::new(vec![
                Box::new(AgentPromptProvider::new(
                    0,
                    provider_config.clone(),
                    prompt_repository.clone(),
                )),
                Box::new(SkillPromptProvider::new(skills)),
            ]));

            let mut router = CommandRouter::new();
            // 注册 Session 处理器
            let session_handler = SessionCommandHandler::new(store.clone())
                .with_session_manager(state.session_manager.clone());
            router.register(Box::new(session_handler));
            // 注册 list_sessions 处理器
            router.register(Box::new(ListSessionsCommandHandler::new(store.clone())));
            // 注册 list_sessions_by_channel 处理器
            router.register(Box::new(ListSessionsByChannelCommandHandler::new(store.clone())));
            // 注册 list_channels 处理器
            router.register(Box::new(ListChannelsCommandHandler::new()));
            // 注册 list_topics 处理器
            router.register(Box::new(ListTopicsCommandHandler::new(store.clone())));
            // 注册 switch_topic 处理器
            let switch_handler = SwitchTopicCommandHandler::new(store.clone())
                .with_session_manager(state.session_manager.clone());
            router.register(Box::new(switch_handler));
            // 注册 get_current 处理器
            router.register(Box::new(GetCurrentSessionCommandHandler::new(store.clone())));
            // 注册 load_topic 处理器
            router.register(Box::new(LoadTopicCommandHandler::new(store.clone())));
            // 注册 load_task_messages 处理器
            router.register(Box::new(LoadTaskMessagesCommandHandler::new(
                state.task_repository.clone(),
                store.clone(),
            )));
            router.register(Box::new(SaveSessionCommandHandler::new(
                store.clone(),
                state.task_repository.clone(),
                system_prompt_provider.clone(),
            )));
            // 注册 help 处理器
            let metadata = router.metadata_arc();
            router.register(Box::new(HelpCommandHandler::new(metadata)));

            // 构建命令上下文
            tracing::debug!(
                current_session_id = %current_session_id,
                current_topic_id = ?current_topic_id,
                "Building CommandContext for WebSocket command"
            );
            let mut cmd_ctx = CommandContext::new("websocket", "cli")
                .with_session_id(current_session_id.as_str())
                .with_chat_id(current_session_id.as_str());
            // 只在有 topic_id 时才设置
            if let Some(ref topic_id) = *current_topic_id {
                cmd_ctx = cmd_ctx.with_topic_id(topic_id.as_str());
            }

            // 执行命令
            let response = router.dispatch_with_response(cmd, cmd_ctx).await;

            // 处理响应
            if response.success {
                // 更新当前会话 ID（如果是创建会话）
                if let Some(session_id) = response.metadata.get("session_id") {
                    tracing::info!(
                        old_session_id = %current_session_id,
                        new_session_id = %session_id,
                        "Updating current_session_id"
                    );
                    *current_session_id = session_id.clone();
                    state
                        .channel_manager
                        .cli_channel()
                        .register_connection(
                            session_id.clone(),
                            runtime_session_id.to_string(),
                            sender.clone(),
                        )
                        .await;
                }
                // 更新当前话题 ID（如果是创建话题或切换话题）
                if let Some(topic_id) = response.metadata.get("topic_id") {
                    tracing::info!(
                        old_topic_id = ?current_topic_id,
                        new_topic_id = %topic_id,
                        "Updating current_topic_id"
                    );
                    *current_topic_id = Some(topic_id.clone());

                    // 加载并发送该话题的历史消息
                    if let Err(e) = send_topic_history(&store, current_session_id, topic_id, sender).await {
                        tracing::warn!(error = %e, topic_id = %topic_id, "Failed to send topic history");
                    }
                }
                // 加载子智能体任务消息
                if let Some(task_session_id) = response.metadata.get("task_session_id") {
                    if let Err(e) = send_task_messages(&store, task_session_id, sender).await {
                        tracing::warn!(error = %e, task_session_id = %task_session_id, "Failed to send task messages");
                    }

                    // 发送 TaskMessagesLoaded 元数据
                    let task_id = response.metadata.get("task_id").cloned().unwrap_or_default();
                    let description = response.metadata.get("task_description").cloned().unwrap_or_default();
                    let subagent_type = response.metadata.get("task_subagent_type").cloned().unwrap_or_default();
                    let status = response.metadata.get("task_status").cloned().unwrap_or_default();
                    let summary = response.metadata.get("task_summary").cloned();

                    let _ = sender.send(WsOutbound::TaskMessagesLoaded {
                        task_id,
                        description,
                        subagent_type,
                        status,
                        summary,
                    }).await;
                }

                if current_topic_id.is_none() {
                    if let Some(topics_json) = response.metadata.get("topics") {
                        match serde_json::from_str::<Vec<crate::protocol::TopicSummary>>(topics_json) {
                            Ok(topics) => {
                                if let Some(first_topic) = topics.first() {
                                    let topic_id = first_topic.topic_id.clone();
                                    *current_topic_id = Some(topic_id.clone());
                                    if let Err(e) = send_topic_history(&store, current_session_id, &topic_id, sender).await {
                                        tracing::warn!(error = %e, topic_id = %topic_id, "Failed to send initial topic history");
                                    }
                                }
                            }
                            Err(e) => {
                                tracing::warn!(error = %e, "Failed to parse topics metadata for initial history");
                            }
                        }
                    }
                }
            } else if let Some(ref error) = response.error {
                tracing::warn!(
                    error_code = %error.code,
                    error_message = %error.message,
                    "Command failed"
                );
            }

            // 适配并发送响应
            let outbounds = output_adapter.adapt(response);
            for msg in outbounds {
                let _ = sender.send(msg).await;
            }

            Ok(())
        }
        WsInbound::Ping => {
            let _ = sender.send(WsOutbound::Pong).await;
            Ok(())
        }
    }
}

fn current_timestamp() -> i64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_millis() as i64
}

fn resolve_ws_sender_id(sender_id: Option<&str>, runtime_session_id: &str) -> String {
    sender_id
        .map(str::trim)
        .filter(|sender_id| !sender_id.is_empty())
        .map(ToOwned::to_owned)
        .unwrap_or_else(|| runtime_session_id.to_string())
}

/// 加载并发送话题历史消息
async fn send_topic_history(
    store: &Arc<crate::storage::SessionStore>,
    _session_id: &str,
    topic_id: &str,
    sender: &mpsc::Sender<WsOutbound>,
) -> Result<(), Box<dyn std::error::Error>> {
    // 加载话题消息
    let messages = store.load_messages_for_topic(topic_id)?;

    tracing::info!(topic_id = %topic_id, message_count = messages.len(), "Sending topic history");

    // 将消息转换为 WsOutbound 并发送
    for msg in messages {
        let outbound = chat_message_to_ws_outbound(&msg);
        if let Some(outbound) = outbound {
            let _ = sender.send(outbound).await;
        }
    }

    Ok(())
}

/// 加载并发送子智能体任务的历史消息
async fn send_task_messages(
    store: &Arc<crate::storage::SessionStore>,
    session_id: &str,
    sender: &mpsc::Sender<WsOutbound>,
) -> Result<(), Box<dyn std::error::Error>> {
    let messages = store.load_messages(session_id)?;

    tracing::info!(session_id = %session_id, message_count = messages.len(), "Sending task messages");

    for msg in messages {
        let outbound = chat_message_to_ws_outbound(&msg);
        if let Some(outbound) = outbound {
            let _ = sender.send(outbound).await;
        }
    }

    Ok(())
}

/// 将 ChatMessage 转换为 WsOutbound
fn chat_message_to_ws_outbound(msg: &crate::bus::ChatMessage) -> Option<WsOutbound> {
    use crate::bus::message::ToolMessageState;

    // Helper function to strip media_refs_json from content
    fn strip_media_refs_json(content: &str) -> String {
        // Remove the media_refs_json suffix if present
        if let Some(pos) = content.find("\n\nmedia_refs_json:") {
            content[..pos].to_string()
        } else {
            content.to_string()
        }
    }

    // Build attachments from media_refs, reading file content for base64
    let attachments: Vec<MediaSummary> = msg
        .media_refs
        .iter()
        .filter_map(|path| {
            // Try to read file and encode as base64
            let file_content = std::fs::read(path).ok()?;
            let base64_content = STANDARD.encode(&file_content);

            // Guess mime type from path
            let mime_type = mime_guess::from_path(path)
                .first_raw()
                .map(ToOwned::to_owned);

            // Determine media type from mime type
            let media_type = mime_type
                .as_ref()
                .map(|m| {
                    if m.starts_with("image/") { "image" }
                    else if m.starts_with("audio/") { "audio" }
                    else if m.starts_with("video/") { "video" }
                    else { "file" }
                })
                .unwrap_or("file");

            // Get file name from path
            let file_name = std::path::Path::new(path)
                .file_name()
                .and_then(|name| name.to_str())
                .map(ToOwned::to_owned);

            Some(MediaSummary {
                path: path.clone(),
                media_type: media_type.to_string(),
                mime_type,
                content_base64: Some(base64_content),
                file_name,
            })
        })
        .collect();

    match msg.role.as_str() {
        "assistant" => {
            if let Some(tool_calls) = &msg.tool_calls {
                // 如果有 tool_calls，发送 ToolCall 消息
                if let Some(tool_call) = tool_calls.first() {
                    return Some(WsOutbound::ToolCall {
                        id: msg.id.clone(),
                        tool_call_id: tool_call.id.clone(),
                        tool_name: tool_call.name.clone(),
                        arguments: tool_call.arguments.clone(),
                        content: format!("{}\nargs: {}", tool_call.name, tool_call.arguments),
                        role: msg.role.clone(),
                        subagent_task_id: None,
                    });
                }
            }
            // 普通助手消息
            Some(WsOutbound::AssistantResponse {
                id: msg.id.clone(),
                content: msg.content.clone(),
                role: msg.role.clone(),
                attachments: Vec::new(),
                subagent_task_id: None,
            })
        }
        "tool" => {
            let tool_state = msg.tool_state.as_ref().unwrap_or(&ToolMessageState::Completed);
            match tool_state {
                ToolMessageState::Completed => Some(WsOutbound::ToolResult {
                    id: msg.id.clone(),
                    tool_call_id: msg.tool_call_id.clone().unwrap_or_default(),
                    tool_name: msg.tool_name.clone().unwrap_or_default(),
                    content: msg.content.clone(),
                    role: msg.role.clone(),
                    subagent_task_id: None,
                }),
                ToolMessageState::PendingUserAction => Some(WsOutbound::ToolPending {
                    id: msg.id.clone(),
                    tool_call_id: msg.tool_call_id.clone().unwrap_or_default(),
                    tool_name: msg.tool_name.clone().unwrap_or_default(),
                    content: msg.content.clone(),
                    role: msg.role.clone(),
                    resume_hint: "完成外部操作后，直接发一条继续消息即可。".to_string(),
                    subagent_task_id: None,
                }),
            }
        }
        "user" => Some(WsOutbound::AssistantResponse {
            id: msg.id.clone(),
            content: strip_media_refs_json(&msg.content),
            role: msg.role.clone(),
            attachments,
            subagent_task_id: None,
        }),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::{resolve_ws_sender_id, build_media_filename, process_attachments_with_base64};
    use crate::protocol::MediaSummary;
    use base64::{Engine as _, engine::general_purpose::STANDARD};

    #[test]
    fn test_resolve_ws_sender_id_prefers_inbound_sender() {
        assert_eq!(
            resolve_ws_sender_id(Some("user-42"), "runtime-1"),
            "user-42"
        );
        assert_eq!(
            resolve_ws_sender_id(Some("  user-42  "), "runtime-1"),
            "user-42"
        );
    }

    #[test]
    fn test_resolve_ws_sender_id_falls_back_to_runtime_session_id() {
        assert_eq!(resolve_ws_sender_id(None, "runtime-1"), "runtime-1");
        assert_eq!(resolve_ws_sender_id(Some("   "), "runtime-1"), "runtime-1");
    }

    #[test]
    fn test_build_media_filename_preserves_original_name() {
        let filename = build_media_filename("image", Some("photo.png"));
        assert!(filename.ends_with("_photo.png"));
        // UUID is 36 chars, plus underscore and original name
        assert!(filename.len() >= 36 + 1 + "photo.png".len());
    }

    #[test]
    fn test_build_media_filename_generates_default_when_no_name() {
        let filename = build_media_filename("image", None);
        assert!(filename.starts_with("image_"));
    }

    #[test]
    fn test_process_attachments_with_base64_saves_to_file() {
        let test_content = "test image content";
        let base64_content = STANDARD.encode(test_content.as_bytes());

        let attachments = vec![MediaSummary {
            path: "test_image.png".to_string(),
            media_type: "image".to_string(),
            mime_type: Some("image/png".to_string()),
            content_base64: Some(base64_content.clone()),
            file_name: Some("test_image.png".to_string()),
        }];

        let result = process_attachments_with_base64(attachments).unwrap();

        // Verify path is now a full path
        assert!(result[0].path.contains(".picobot"));
        assert!(result[0].path.contains("media"));
        assert!(result[0].path.contains("ws"));

        // Verify content_base64 is kept for frontend display
        assert!(result[0].content_base64.is_some());
        assert_eq!(result[0].content_base64.as_ref().unwrap(), &base64_content);

        // Verify file was actually written
        let file_content = std::fs::read(&result[0].path).unwrap();
        assert_eq!(file_content, test_content.as_bytes());

        // Cleanup
        std::fs::remove_file(&result[0].path).ok();
    }

    #[test]
    fn test_process_attachments_without_base64_keeps_original_path() {
        let attachments = vec![MediaSummary {
            path: "/existing/path/image.jpg".to_string(),
            media_type: "image".to_string(),
            mime_type: Some("image/jpeg".to_string()),
            content_base64: None,
            file_name: Some("image.jpg".to_string()),
        }];

        let result = process_attachments_with_base64(attachments).unwrap();

        // Path should remain unchanged
        assert_eq!(result[0].path, "/existing/path/image.jpg");
    }

    #[test]
    fn test_process_empty_attachments_returns_empty_vec() {
        let result = process_attachments_with_base64(Vec::new()).unwrap();
        assert!(result.is_empty());
    }
}