"""Progress tracking and callbacks for verbose agent execution."""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Callable, Awaitable
from datetime import datetime


class ProgressStage(Enum):
    """Stages of agent execution."""
    # Session stages
    SESSION_START = "session_start"
    SESSION_END = "session_end"

    # Retrieval stages
    RETRIEVING_SKILLS = "retrieving_skills"
    SKILLS_RETRIEVED = "skills_retrieved"
    RETRIEVING_TOOLS = "retrieving_tools"
    TOOLS_RETRIEVED = "tools_retrieved"

    # Context gathering stages
    GATHERING_CONTEXT = "gathering_context"
    CONTEXT_GATHERED = "context_gathered"

    # Planning stages
    PLANNING = "planning"
    PLAN_CREATED = "plan_created"

    # Tool building stages
    BUILDING_TOOLS = "building_tools"
    TOOL_BUILD_FAILED = "tool_build_failed"
    TOOLS_BUILT = "tools_built"

    # Execution stages
    EXECUTING_PLAN = "executing_plan"
    EXECUTING_STEP = "executing_step"
    STEP_COMPLETED = "step_completed"
    STEP_FAILED = "step_failed"

    # Self-healing stages
    SELF_HEALING = "self_healing"
    HEAL_ATTEMPT = "heal_attempt"

    # Replanning stages
    REPLANNING = "replanning"
    PLAN_UPDATED = "plan_updated"

    # User feedback
    USER_FEEDBACK = "user_feedback"

    # Tool execution stages
    TOOL_STARTING = "tool_starting"
    TOOL_COMPLETED = "tool_completed"
    TOOL_FAILED = "tool_failed"

    # Agentic loop stages
    LOOP_ITERATION = "loop_iteration"
    LOOP_TOOL_CALL = "loop_tool_call"

    # Learning stages
    LEARNING = "learning"
    SKILL_LEARNED = "skill_learned"

    # Final stages
    SYNTHESIZING = "synthesizing"
    TASK_COMPLETED = "task_completed"
    TASK_FAILED = "task_failed"


@dataclass
class ProgressEvent:
    """A progress event with details."""
    stage: ProgressStage
    message: str
    timestamp: datetime = field(default_factory=datetime.now)
    details: dict[str, Any] = field(default_factory=dict)

    # For nested progress (e.g., step 2 of 5)
    current: int | None = None
    total: int | None = None


# Type for progress callbacks
ProgressCallback = Callable[[ProgressEvent], Awaitable[None] | None]


class ProgressTracker:
    """Tracks and reports progress during agent execution."""

    def __init__(self, callback: ProgressCallback | None = None):
        """
        Initialize the progress tracker.

        Args:
            callback: Optional callback function to receive progress events.
                      Can be sync or async.
        """
        self._callback = callback
        self._events: list[ProgressEvent] = []

    async def emit(
        self,
        stage: ProgressStage,
        message: str,
        current: int | None = None,
        total: int | None = None,
        **details: Any,
    ) -> None:
        """
        Emit a progress event.

        Args:
            stage: The current stage
            message: Human-readable message
            current: Current item number (for progress bars)
            total: Total items (for progress bars)
            **details: Additional details to include
        """
        event = ProgressEvent(
            stage=stage,
            message=message,
            current=current,
            total=total,
            details=details,
        )
        self._events.append(event)

        if self._callback:
            result = self._callback(event)
            # Handle async callbacks
            if hasattr(result, '__await__'):
                await result

    @property
    def events(self) -> list[ProgressEvent]:
        """Get all recorded events."""
        return self._events.copy()

    def clear(self) -> None:
        """Clear recorded events."""
        self._events.clear()


class NullProgressTracker(ProgressTracker):
    """A no-op progress tracker for when verbose mode is disabled."""

    def __init__(self):
        super().__init__(callback=None)

    async def emit(self, *args, **kwargs) -> None:
        """No-op emit."""
        pass
