"""Execution-related types.

Types for execution budgets, step outcomes, loop results, and execution modes.
"""

from enum import StrEnum
from typing import Annotated, Literal

from pydantic import BaseModel, Field

# =============================================================================
# EXECUTION MODE
# =============================================================================


class ExecutionMode(StrEnum):
    """Execution mode that influences runtime defaults.

    The mode is a configuration convenience — it sets sensible defaults
    for streaming, error handling, checkpointing, and history management
    without changing the pathway graph itself.

    Modes
    -----
    CONVERSATIONAL
        Chat agent pattern.  Streaming enabled, conversation history
        maintained, errors reported to the user.
    WORKFLOW
        Autonomous, goal-directed execution.  Checkpoints created
        automatically, errors retried, no streaming needed.
    PIPELINE
        Batch iteration.  Per-record error isolation (skip on failure),
        no conversation history, no streaming.
    """

    CONVERSATIONAL = "conversational"
    WORKFLOW = "workflow"
    PIPELINE = "pipeline"


# JSON-serializable output value
StepOutputValue = str | int | float | bool | None | list[object] | dict[str, object]


class ExecutionBudget(BaseModel):
    """Budget constraints for execution."""

    max_tokens: int = 100000  # Total tokens per turn
    max_steps: int = 10  # Max tool calls per turn
    max_duration_seconds: int = 300  # 5 minutes
    max_cost_usd: float = 1.0  # Dollar limit per turn

    # Context limits
    max_context_messages: int = 20  # History truncation


# =============================================================================
# STEP OUTCOME DISCRIMINATED UNION
# =============================================================================


class StepSuccess(BaseModel):
    """Step completed successfully."""

    type: Literal["success"] = "success"
    output: StepOutputValue = "Completed"
    tool_name: str | None = None
    tool_call_id: str | None = None
    duration_ms: float | None = None


class StepFailure(BaseModel):
    """Step failed with an error."""

    type: Literal["failure"] = "failure"
    error: str
    can_retry: bool = False
    retry_after_seconds: float | None = None
    tool_name: str | None = None
    exception_type: str | None = None


class StepSkipped(BaseModel):
    """Step was skipped (e.g., conditional not met, dependency failed)."""

    type: Literal["skipped"] = "skipped"
    reason: str
    skipped_tool: str | None = None


# Discriminated union with type field
StepOutcome = Annotated[
    StepSuccess | StepFailure | StepSkipped,
    Field(discriminator="type"),
]


# =============================================================================
# LOOP RESULT
# =============================================================================


class LoopResult(BaseModel):
    """Base result for all loops.

    All loop-specific results should extend this with additional fields.
    """

    success: bool
    response: str
    steps_taken: int
    tool_calls: list[object] = Field(default_factory=list)  # List[ToolCallRecord]
    tokens_used: int = 0
    cost_usd: float = 0.0
    error: str | None = None


# =============================================================================
# HELPER FUNCTIONS
# =============================================================================


def is_step_successful(outcome: StepSuccess | StepFailure | StepSkipped) -> bool:
    """Check if a step outcome represents success."""
    return isinstance(outcome, StepSuccess)


def count_failures(outcomes: list[StepSuccess | StepFailure | StepSkipped]) -> int:
    """Count the number of failed steps."""
    return sum(1 for o in outcomes if isinstance(o, StepFailure))


def get_last_error(outcomes: list[StepSuccess | StepFailure | StepSkipped]) -> str | None:
    """Get the error message from the last failed step, if any."""
    for outcome in reversed(outcomes):
        if isinstance(outcome, StepFailure):
            return outcome.error
    return None


class ExecutionResult(BaseModel):
    """Result of an agent turn."""

    success: bool
    response: str
    agent_id: str = ""
    thread_id: str = ""
    pathway: str = ""
    steps_taken: int = 0
    tokens_used: int = 0
    cost_usd: float = 0.0
    tool_calls: list[object] = Field(default_factory=list)
    error: str | None = None


__all__ = [
    "ExecutionBudget",
    "ExecutionMode",
    "ExecutionResult",
    "LoopResult",
    "StepFailure",
    "StepOutcome",
    "StepOutputValue",
    "StepSkipped",
    "StepSuccess",
    "count_failures",
    "get_last_error",
    "is_step_successful",
]
