"""Tracing stage for recording step execution history."""

from __future__ import annotations

from typing import TYPE_CHECKING

from agent_flows.models.execution import TraceEntry
from agent_flows.utils.logging import get_logger

if TYPE_CHECKING:
    from agent_flows.core.flow_engine.pipeline import ExecutionState

logger = get_logger(__name__)


class TracingStage:
    """Adds trace entry for executed step.

    Runs: Post-execution only
    Modifies: Context trace_log

    When tracing is enabled, this stage records the step execution outcome
    so callers can inspect the flow history.
    """

    async def process(self, state: ExecutionState) -> ExecutionState:
        """Add trace entry for executed step.

        Args:
            state: Current execution state

        Returns:
            Same state (trace log is mutated in place)
        """
        # Only run after execution
        if state.is_pre_execution:
            return state

        # Skip if tracing not enabled
        if state.context.trace_log is None:
            return state

        # Create and append trace entry (mutates shared list)
        entry = TraceEntry(
            step_id=state.step.id or "",
            step_type=state.step.type,
            success=state.result.success,
            execution_time=state.result.execution_time,
            data=state.result.data,
            variables_updated=state.result.variables_updated or None,
        )
        state.context.trace_log.append(entry)

        logger.debug(
            "Appended trace entry",
            extra={
                "step_id": state.step.id,
                "step_type": state.step.type,
                "success": state.result.success,
            },
        )

        return state
