//! Memory step types for different agent actions.
//!
//! This module defines all the step types that can be stored in agent memory:
//! - `SystemPromptStep` - The initial system prompt
//! - `TaskStep` - A new task submitted to the agent
//! - `ActionStep` - An action taken by the agent
//! - `PlanningStep` - A planning step
//! - `FinalAnswerStep` - The final answer

use std::any::Any;

use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::message::{ChatMessage, MessageContent, MessageRole};
use crate::multimodal::AgentImage;
use crate::providers::common::TokenUsage;

use super::timing::Timing;
use super::types::{MemoryStep, ToolCall};

/// System prompt step.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemPromptStep {
    /// The system prompt.
    pub system_prompt: String,
}

impl MemoryStep for SystemPromptStep {
    fn to_messages(&self, summary_mode: bool) -> Vec<ChatMessage> {
        if summary_mode {
            return vec![];
        }
        vec![ChatMessage::system(&self.system_prompt)]
    }

    fn to_value(&self) -> Value {
        serde_json::to_value(self).unwrap_or_default()
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

/// Task step representing a new task.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskStep {
    /// The task description.
    pub task: String,
    /// Optional images associated with the task (for vision models).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub task_images: Option<Vec<AgentImage>>,
}

impl TaskStep {
    /// Create a new task step with just a task description.
    #[must_use]
    pub fn new(task: impl Into<String>) -> Self {
        Self {
            task: task.into(),
            task_images: None,
        }
    }

    /// Create a new task step with images.
    #[must_use]
    pub fn with_images(task: impl Into<String>, images: Vec<AgentImage>) -> Self {
        Self {
            task: task.into(),
            task_images: if images.is_empty() {
                None
            } else {
                Some(images)
            },
        }
    }

    /// Add images to the task step.
    #[must_use]
    pub fn add_images(mut self, images: Vec<AgentImage>) -> Self {
        self.task_images = if images.is_empty() {
            None
        } else {
            Some(images)
        };
        self
    }

    /// Check if this task has images.
    #[must_use]
    pub fn has_images(&self) -> bool {
        self.task_images
            .as_ref()
            .is_some_and(|imgs| !imgs.is_empty())
    }
}

impl MemoryStep for TaskStep {
    fn to_messages(&self, _summary_mode: bool) -> Vec<ChatMessage> {
        let mut content = vec![MessageContent::text(format!("New task:\n{}", self.task))];

        // Add images if present (for vision models)
        if let Some(images) = &self.task_images {
            for image in images {
                if let Some(img_content) = MessageContent::from_agent_image(image) {
                    content.push(img_content);
                }
            }
        }

        vec![ChatMessage::with_contents(MessageRole::User, content)]
    }

    fn to_value(&self) -> Value {
        serde_json::json!({
            "task": self.task,
            "has_images": self.has_images(),
            "image_count": self.task_images.as_ref().map_or(0, Vec::len)
        })
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

/// Action step representing an agent action.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ActionStep {
    /// Step number.
    pub step_number: usize,
    /// Timing information.
    pub timing: Timing,
    /// Input messages to the model.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub model_input_messages: Option<Vec<ChatMessage>>,
    /// Tool calls made in this step.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,
    /// Error that occurred (if any).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
    /// Model output message.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub model_output_message: Option<ChatMessage>,
    /// Raw model output.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub model_output: Option<String>,
    /// Code action (for code agents).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub code_action: Option<String>,
    /// Observations from tool execution.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub observations: Option<String>,
    /// Images from tool observations (e.g., screenshots, generated images).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub observations_images: Option<Vec<AgentImage>>,
    /// Output of the action.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub action_output: Option<Value>,
    /// Token usage for this step.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub token_usage: Option<TokenUsage>,
    /// Whether this is the final answer step.
    #[serde(default)]
    pub is_final_answer: bool,
}

impl ActionStep {
    /// Check if this step has observation images.
    #[must_use]
    pub fn has_observation_images(&self) -> bool {
        self.observations_images
            .as_ref()
            .is_some_and(|imgs| !imgs.is_empty())
    }

    /// Add observation images to this step.
    pub fn add_observation_images(&mut self, images: Vec<AgentImage>) {
        if images.is_empty() {
            return;
        }
        if let Some(ref mut existing) = self.observations_images {
            existing.extend(images);
        } else {
            self.observations_images = Some(images);
        }
    }

    /// Clear observation images (useful for memory management in long runs).
    pub fn clear_observation_images(&mut self) {
        self.observations_images = None;
    }
}

impl MemoryStep for ActionStep {
    fn to_messages(&self, summary_mode: bool) -> Vec<ChatMessage> {
        let mut messages = vec![];

        // Add model output
        if let Some(output) = &self.model_output
            && !summary_mode
        {
            messages.push(ChatMessage::assistant(output.trim()));
        }

        // Add tool calls
        if let Some(tool_calls) = &self.tool_calls {
            let calls_str = serde_json::to_string(tool_calls).unwrap_or_default();
            messages.push(ChatMessage {
                role: MessageRole::ToolCall,
                content: Some(vec![MessageContent::text(format!(
                    "Calling tools:\n{calls_str}"
                ))]),
                tool_calls: None,
                tool_call_id: None,
            });
        }

        // Add observation images (for vision models)
        if let Some(images) = &self.observations_images {
            let content: Vec<MessageContent> = images
                .iter()
                .filter_map(MessageContent::from_agent_image)
                .collect();
            if !content.is_empty() {
                messages.push(ChatMessage::with_contents(MessageRole::User, content));
            }
        }

        // Add text observations
        if let Some(obs) = &self.observations {
            messages.push(ChatMessage {
                role: MessageRole::ToolResponse,
                content: Some(vec![MessageContent::text(format!("Observation:\n{obs}"))]),
                tool_calls: None,
                tool_call_id: None,
            });
        }

        // Add error
        if let Some(err) = &self.error {
            let error_msg =
                format!("Error:\n{err}\nNow let's retry: take care not to repeat previous errors!");
            messages.push(ChatMessage {
                role: MessageRole::ToolResponse,
                content: Some(vec![MessageContent::text(error_msg)]),
                tool_calls: None,
                tool_call_id: None,
            });
        }

        messages
    }

    fn to_value(&self) -> Value {
        serde_json::to_value(self).unwrap_or_default()
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

/// Planning step for agent planning.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PlanningStep {
    /// Input messages to the model.
    pub model_input_messages: Vec<ChatMessage>,
    /// Model output message.
    pub model_output_message: ChatMessage,
    /// The plan text.
    pub plan: String,
    /// Timing information.
    pub timing: Timing,
    /// Token usage.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub token_usage: Option<TokenUsage>,
}

impl MemoryStep for PlanningStep {
    fn to_messages(&self, summary_mode: bool) -> Vec<ChatMessage> {
        if summary_mode {
            return vec![];
        }
        vec![
            ChatMessage::assistant(self.plan.trim()),
            ChatMessage::user("Now proceed and carry out this plan."),
        ]
    }

    fn to_value(&self) -> Value {
        serde_json::to_value(self).unwrap_or_default()
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

/// Final answer step.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FinalAnswerStep {
    /// The final output.
    pub output: Value,
}

impl MemoryStep for FinalAnswerStep {
    fn to_messages(&self, _summary_mode: bool) -> Vec<ChatMessage> {
        vec![]
    }

    fn to_value(&self) -> Value {
        serde_json::json!({ "output": self.output })
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}
