"""
Module for tracking execution steps in Agentle agents.

This module provides the Step class, which represents a single execution step
in an agent's reasoning process. Steps track the specific actions taken by an
agent during its execution, particularly focusing on tool calls.

Steps are stored in the Context object's 'steps' sequence and provide a
detailed record of the agent's decision-making and actions during execution.
This is useful for debugging, logging, and understanding the agent's reasoning path.

Example:
```python
from agentle.agents.step import Step
from agentle.generations.models.message_parts.tool_execution_suggestion import ToolExecutionSuggestion

# Create a tool execution suggestion
tool_suggestion = ToolExecutionSuggestion(
    id="tool-call-1",
    tool_name="get_weather",
    args={"location": "London"}
)

# Create a step recording that tool call
step = Step(tool_execution_suggestions=[tool_suggestion])

# Steps are typically stored in a Context object
context.steps = list(context.steps) + [step]
```
"""

from collections.abc import MutableMapping, Sequence
from datetime import datetime
from typing import Any, Literal

from rsb.models.base_model import BaseModel
from rsb.models.field import Field

from agentle.generations.models.message_parts.tool_execution_suggestion import (
    ToolExecutionSuggestion,
)
from agentle.generations.models.generation.usage import Usage


type StepType = Literal[
    "tool_execution", "generation", "reasoning", "error", "user_input"
]


class ToolExecutionResult(BaseModel):
    """
    Represents the result of a tool execution within a step.

    This class captures both the suggestion that was executed and the result
    that was returned, providing a complete record of the tool interaction.
    """

    suggestion: ToolExecutionSuggestion
    """The tool execution suggestion that was executed."""

    result: Any
    """The result returned by the tool execution."""

    execution_time_ms: float | None = Field(default=None)
    """Time taken to execute the tool in milliseconds."""

    success: bool = Field(default=True)
    """Whether the tool execution was successful."""

    error_message: str | None = Field(default=None)
    """Error message if the tool execution failed."""


class Step(BaseModel):
    """
    Represents a single execution step in an agent's reasoning process.

    The Step class tracks the actions and decisions made by an agent during execution,
    with a primary focus on tool calls. Each step captures a discrete point in the
    agent's execution where it interacted with external tools or made significant
    progress in its reasoning.

    Steps serve several important purposes:
    1. Tracking the sequence of tool calls made by the agent
    2. Providing a chronological record of agent actions for debugging
    3. Supporting audit trails of what external resources were accessed
    4. Facilitating analysis of agent decision-making patterns
    5. Enabling context resumption for asynchronous operations

    Steps are stored in the Context object's 'steps' sequence. They are primarily
    used internally by the agent execution logic, but are exposed in the
    AgentRunOutput for inspection and logging.

    Attributes:
        step_id: Unique identifier for this step
        step_type: The type of step (tool_execution, generation, reasoning, etc.)
        timestamp: When this step was created
        iteration: Which iteration of the agent loop this step belongs to
        tool_execution_suggestions: Tool execution suggestions made during this step
        tool_execution_results: Results from tool executions in this step
        generation_text: Text generated by the model in this step
        reasoning: Any reasoning or thought process captured in this step
        token_usage: Token usage statistics for this step
        metadata: Additional metadata for this step
        duration_ms: How long this step took to complete in milliseconds
        error_message: Error message if this step failed

    Example:
        ```python
        # Creating a step to record a weather API call
        from agentle.generations.models.message_parts.tool_execution_suggestion import ToolExecutionSuggestion

        weather_tool_call = ToolExecutionSuggestion(
            id="call-123",
            tool_name="get_weather",
            args={"location": "New York", "units": "celsius"}
        )

        step = Step(
            step_type="tool_execution",
            iteration=1,
            tool_execution_suggestions=[weather_tool_call]
        )

        # After tool execution, update with results
        step.tool_execution_results = [
            ToolExecutionResult(
                suggestion=weather_tool_call,
                result={"temperature": 22, "condition": "sunny"},
                execution_time_ms=150.5,
                success=True
            )
        ]

        # Steps are typically added to a context during agent execution
        context.steps = list(context.steps) + [step]

        # Examining steps after execution
        for step in result.steps:
            print(f"Step {step.step_id}: {step.step_type} at {step.timestamp}")
            for result in step.tool_execution_results:
                print(f"  Tool: {result.suggestion.tool_name} -> {result.result}")
        ```
    """

    step_id: str = Field(default_factory=lambda: f"step_{datetime.now().timestamp()}")
    """Unique identifier for this step."""

    step_type: StepType = Field(default="tool_execution")
    """The type of step being performed."""

    timestamp: datetime = Field(default_factory=datetime.now)
    """When this step was created."""

    iteration: int = Field(default=0)
    """Which iteration of the agent loop this step belongs to."""

    tool_execution_suggestions: Sequence[ToolExecutionSuggestion] = Field(
        default_factory=list
    )
    """
    A sequence of tool execution suggestions that were made during this step.
    
    Each ToolExecutionSuggestion contains information about a specific tool call,
    including the tool name, arguments passed to the tool, and a unique identifier
    for the call.
    """

    tool_execution_results: Sequence[ToolExecutionResult] = Field(default_factory=list)
    """Results from tool executions performed in this step."""

    generation_text: str | None = Field(default=None)
    """Text generated by the model in this step."""

    reasoning: str | None = Field(default=None)
    """Any reasoning or thought process captured in this step."""

    token_usage: Usage | None = Field(default=None)
    """Token usage statistics for this step."""

    metadata: MutableMapping[str, Any] = Field(default_factory=dict)
    """Additional metadata for this step."""

    duration_ms: float | None = Field(default=None)
    """How long this step took to complete in milliseconds."""

    error_message: str | None = Field(default=None)
    """Error message if this step failed."""

    def add_tool_execution_result(
        self,
        suggestion: ToolExecutionSuggestion,
        result: Any,
        execution_time_ms: float | None = None,
        success: bool = True,
        error_message: str | None = None,
    ) -> None:
        """
        Add a tool execution result to this step.

        Args:
            suggestion: The tool execution suggestion that was executed
            result: The result returned by the tool
            execution_time_ms: Time taken to execute the tool in milliseconds
            success: Whether the execution was successful
            error_message: Error message if execution failed
        """
        execution_result = ToolExecutionResult(
            suggestion=suggestion,
            result=result,
            execution_time_ms=execution_time_ms,
            success=success,
            error_message=error_message,
        )

        # Convert to mutable list if needed
        if not isinstance(self.tool_execution_results, list):
            self.tool_execution_results = list(self.tool_execution_results)

        self.tool_execution_results.append(execution_result)

    def mark_completed(self, duration_ms: float | None = None) -> None:
        """
        Mark this step as completed.

        Args:
            duration_ms: How long the step took to complete in milliseconds
        """
        if duration_ms is not None:
            self.duration_ms = duration_ms

    def mark_failed(self, error_message: str, duration_ms: float | None = None) -> None:
        """
        Mark this step as failed.

        Args:
            error_message: The error message describing the failure
            duration_ms: How long the step took before failing in milliseconds
        """
        self.error_message = error_message
        if duration_ms is not None:
            self.duration_ms = duration_ms

    @property
    def is_successful(self) -> bool:
        """Check if this step completed successfully."""
        return self.error_message is None

    @property
    def has_tool_executions(self) -> bool:
        """Check if this step contains any tool executions."""
        return (
            len(self.tool_execution_suggestions) > 0
            or len(self.tool_execution_results) > 0
        )
