//! WebSocket 消息处理器
//! 专门处理 ComfyUI WebSocket 消息的解析和事件转换

use anyhow::{Result, anyhow};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{RwLock, broadcast};
use tracing::{info, warn, error, debug};
use serde_json::Value;

use crate::business::services::realtime_monitor_v2::RealtimeEventV2;
use crate::data::repositories::comfyui_repository::ComfyUIRepository;

/// WebSocket 消息处理器
pub struct WebSocketHandler {
    /// 事件发送器
    event_sender: broadcast::Sender<RealtimeEventV2>,
    /// 提示 ID 到执行 ID 的映射
    prompt_execution_map: Arc<RwLock<HashMap<String, String>>>,
    /// 数据仓库
    repository: Arc<ComfyUIRepository>,
}

impl WebSocketHandler {
    /// 创建新的消息处理器
    pub fn new(
        event_sender: broadcast::Sender<RealtimeEventV2>,
        prompt_execution_map: Arc<RwLock<HashMap<String, String>>>,
        repository: Arc<ComfyUIRepository>,
    ) -> Self {
        Self {
            event_sender,
            prompt_execution_map,
            repository,
        }
    }

    /// 处理 WebSocket 消息
    pub async fn handle_message(&self, message: Value) -> Result<()> {
        debug!("处理 WebSocket 消息: {:?}", message);

        // 解析消息类型
        let message_type = message.get("type")
            .and_then(|t| t.as_str())
            .unwrap_or("unknown");

        let timestamp = chrono::Utc::now().to_rfc3339();

        match message_type {
            "executing" => self.handle_executing_message(message, timestamp).await?,
            "progress" => self.handle_progress_message(message, timestamp).await?,
            "executed" => self.handle_executed_message(message, timestamp).await?,
            "execution_error" => self.handle_execution_error_message(message, timestamp).await?,
            "status" => self.handle_status_message(message, timestamp).await?,
            _ => {
                debug!("未处理的消息类型: {}", message_type);
            }
        }

        Ok(())
    }

    /// 处理执行消息
    async fn handle_executing_message(&self, message: Value, timestamp: String) -> Result<()> {
        if let Some(data) = message.get("data") {
            let prompt_id = data.get("prompt_id")
                .and_then(|p| p.as_str())
                .unwrap_or("")
                .to_string();
            
            let node_id = data.get("node")
                .and_then(|n| n.as_str())
                .map(|s| s.to_string());

            let execution_id = {
                let map = self.prompt_execution_map.read().await;
                map.get(&prompt_id).cloned()
            };

            if let Some(node_id) = node_id {
                let _ = self.event_sender.send(RealtimeEventV2::NodeExecuting {
                    prompt_id: prompt_id.clone(),
                    execution_id: execution_id.clone(),
                    node_id: node_id.clone(),
                    node_title: None,
                    timestamp: timestamp.clone(),
                });
            } else {
                let _ = self.event_sender.send(RealtimeEventV2::ExecutionStarted {
                    prompt_id,
                    execution_id,
                    node_id: None,
                    timestamp,
                });
            }
        }
        Ok(())
    }

    /// 处理进度消息
    async fn handle_progress_message(&self, message: Value, timestamp: String) -> Result<()> {
        if let Some(data) = message.get("data") {
            let prompt_id = data.get("prompt_id")
                .and_then(|p| p.as_str())
                .unwrap_or("")
                .to_string();
            
            let node_id = data.get("node")
                .and_then(|n| n.as_str())
                .unwrap_or("")
                .to_string();

            let progress = data.get("value")
                .and_then(|v| v.as_u64())
                .unwrap_or(0) as u32;

            let max_progress = data.get("max")
                .and_then(|m| m.as_u64())
                .unwrap_or(100) as u32;

            let percentage = if max_progress > 0 {
                (progress as f32 / max_progress as f32) * 100.0
            } else {
                0.0
            };

            let execution_id = {
                let map = self.prompt_execution_map.read().await;
                map.get(&prompt_id).cloned()
            };

            let _ = self.event_sender.send(RealtimeEventV2::ExecutionProgress {
                prompt_id,
                execution_id,
                node_id,
                progress,
                max_progress,
                percentage,
                timestamp,
            });
        }
        Ok(())
    }

    /// 处理执行完成消息
    async fn handle_executed_message(&self, message: Value, timestamp: String) -> Result<()> {
        if let Some(data) = message.get("data") {
            let prompt_id = data.get("prompt_id")
                .and_then(|p| p.as_str())
                .unwrap_or("")
                .to_string();

            let outputs = data.get("output")
                .and_then(|o| o.as_object())
                .map(|obj| {
                    obj.iter()
                        .map(|(k, v)| (k.clone(), v.clone()))
                        .collect::<HashMap<String, serde_json::Value>>()
                })
                .unwrap_or_default();

            let output_urls = self.extract_output_urls(&outputs);

            let execution_id = {
                let map = self.prompt_execution_map.read().await;
                map.get(&prompt_id).cloned()
            };

            // 更新数据库中的执行状态
            if let Some(exec_id) = &execution_id {
                if let Ok(Some(mut execution)) = self.repository.get_execution(exec_id).await {
                    execution.set_results(outputs.clone(), output_urls.clone());
                    let _ = self.repository.update_execution(&execution).await;
                }
            }

            let _ = self.event_sender.send(RealtimeEventV2::ExecutionCompleted {
                prompt_id: prompt_id.clone(),
                execution_id: execution_id.clone(),
                outputs,
                output_urls,
                execution_time: 0, // TODO: 计算实际执行时间
                timestamp,
            });

            // 清理映射
            {
                let mut map = self.prompt_execution_map.write().await;
                map.remove(&prompt_id);
            }
        }
        Ok(())
    }

    /// 处理执行错误消息
    async fn handle_execution_error_message(&self, message: Value, timestamp: String) -> Result<()> {
        if let Some(data) = message.get("data") {
            let prompt_id = data.get("prompt_id")
                .and_then(|p| p.as_str())
                .unwrap_or("")
                .to_string();

            let node_id = data.get("node_id")
                .and_then(|n| n.as_str())
                .map(|s| s.to_string());

            let error_msg = data.get("exception_message")
                .and_then(|e| e.as_str())
                .unwrap_or("未知错误")
                .to_string();

            let details = data.get("traceback").cloned();

            let execution_id = {
                let map = self.prompt_execution_map.read().await;
                map.get(&prompt_id).cloned()
            };

            // 更新数据库中的执行状态
            if let Some(exec_id) = &execution_id {
                if let Ok(Some(mut execution)) = self.repository.get_execution(exec_id).await {
                    execution.set_error(error_msg.clone());
                    let _ = self.repository.update_execution(&execution).await;
                }
            }

            let _ = self.event_sender.send(RealtimeEventV2::ExecutionFailed {
                prompt_id: prompt_id.clone(),
                execution_id: execution_id.clone(),
                node_id,
                error: error_msg,
                details,
                timestamp,
            });

            // 清理映射
            {
                let mut map = self.prompt_execution_map.write().await;
                map.remove(&prompt_id);
            }
        }
        Ok(())
    }

    /// 处理状态消息
    async fn handle_status_message(&self, message: Value, timestamp: String) -> Result<()> {
        if let Some(data) = message.get("data") {
            let queue_running = data.get("status")
                .and_then(|s| s.get("exec_info"))
                .and_then(|e| e.get("queue_remaining"))
                .and_then(|q| q.as_u64())
                .unwrap_or(0) as u32;

            let _ = self.event_sender.send(RealtimeEventV2::QueueUpdated {
                running_count: if queue_running > 0 { 1 } else { 0 },
                pending_count: queue_running,
                timestamp,
            });
        }
        Ok(())
    }

    /// 提取输出 URLs
    fn extract_output_urls(&self, outputs: &HashMap<String, serde_json::Value>) -> Vec<String> {
        let mut urls = Vec::new();

        for (_, output) in outputs {
            if let Some(images) = output.get("images").and_then(|v| v.as_array()) {
                for image in images {
                    if let Some(filename) = image.get("filename").and_then(|v| v.as_str()) {
                        let url = format!("/view?filename={}", filename);
                        urls.push(url);
                    }
                }
            }
        }

        urls
    }
}