//! Event types for ComfyUI SDK

use serde::{Deserialize, Serialize};
use chrono::{DateTime, Utc};
use std::collections::HashMap;

/// Execution progress information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionProgress {
    pub node_id: String,
    pub progress: u32,
    pub max: u32,
    pub timestamp: DateTime<Utc>,
}

/// Execution result information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionResult {
    pub prompt_id: String,
    pub outputs: HashMap<String, serde_json::Value>,
    pub execution_time: u64,
    pub timestamp: DateTime<Utc>,
}

/// Execution error information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionError {
    pub node_id: Option<String>,
    pub message: String,
    pub details: Option<serde_json::Value>,
    pub timestamp: DateTime<Utc>,
}

/// Callback trait for execution events
pub trait ExecutionCallbacks: Send + Sync {
    /// Called when progress is updated
    fn on_progress(&self, _progress: ExecutionProgress) {}

    /// Called when a node starts executing
    fn on_executing(&self, _node_id: String) {}

    /// Called when execution completes successfully
    fn on_executed(&self, _result: ExecutionResult) {}

    /// Called when an error occurs
    fn on_error(&self, _error: ExecutionError) {}
}

/// WebSocket message types
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum WSMessage {
    #[serde(rename = "progress")]
    Progress {
        data: WSProgressData,
    },
    #[serde(rename = "executing")]
    Executing {
        data: WSExecutingData,
    },
    #[serde(rename = "executed")]
    Executed {
        data: WSExecutedData,
    },
    #[serde(rename = "execution_error")]
    ExecutionError {
        data: WSErrorData,
    },
    #[serde(other)]
    Unknown,
}

/// WebSocket progress message data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WSProgressData {
    pub value: u32,
    pub max: u32,
    pub node: Option<String>,
}

/// WebSocket executing message data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WSExecutingData {
    pub node: Option<String>,
    pub prompt_id: String,
}

/// WebSocket executed message data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WSExecutedData {
    pub node: String,
    pub output: HashMap<String, serde_json::Value>,
    pub prompt_id: String,
}

/// WebSocket error message data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WSErrorData {
    pub node_id: Option<String>,
    pub node: Option<String>,
    pub message: Option<String>,
    pub details: Option<serde_json::Value>,
    pub error: Option<String>,
    pub exception_message: Option<String>,
    pub exception_type: Option<String>,
    pub traceback: Option<Vec<String>>,
    pub prompt_id: Option<String>,
    // 使用 flatten 来捕获任何其他字段
    #[serde(flatten)]
    pub extra: HashMap<String, serde_json::Value>,
}

/// Execution options
#[derive(Debug, Clone)]
pub struct ExecutionOptions {
    pub timeout: Option<std::time::Duration>,
    pub priority: Option<i32>,
}

impl Default for ExecutionOptions {
    fn default() -> Self {
        Self {
            timeout: Some(std::time::Duration::from_secs(300)), // 5 minutes default
            priority: None,
        }
    }
}

/// Template execution result
#[derive(Debug, Clone)]
pub struct TemplateExecutionResult {
    pub prompt_id: String,
    pub success: bool,
    pub outputs: Option<HashMap<String, serde_json::Value>>,
    pub error: Option<ExecutionError>,
    pub execution_time: u64,
}

impl TemplateExecutionResult {
    /// Create a successful execution result
    pub fn success(
        prompt_id: String,
        outputs: HashMap<String, serde_json::Value>,
        execution_time: u64,
    ) -> Self {
        Self {
            prompt_id,
            success: true,
            outputs: Some(outputs),
            error: None,
            execution_time,
        }
    }

    /// Create a failed execution result
    pub fn failure(
        prompt_id: String,
        error: ExecutionError,
        execution_time: u64,
    ) -> Self {
        Self {
            prompt_id,
            success: false,
            outputs: None,
            error: Some(error),
            execution_time,
        }
    }
}
