"""Base step abstraction for workflows."""

from __future__ import annotations

import asyncio
from abc import ABC, abstractmethod
from datetime import datetime
from typing import Any, Generic, Type

from agentbyte.workflow.core.models import (
    InputType,
    OutputType,
    StepMetadata,
    StepStatus,
    WorkflowContext,
    WorkflowPendingRequest,
    _WorkflowRequestSuspended,
)


class BaseStep(Generic[InputType, OutputType], ABC):
    """Abstract workflow step.

    Subclasses implement `execute`, which receives validated input and a shared
    WorkflowContext. The public `run` method handles validation, retry, timeout, and
    status bookkeeping.
    """

    def __init__(
        self,
        step_id: str,
        metadata: StepMetadata,
        input_type: Type[InputType],
        output_type: Type[OutputType],
    ) -> None:
        self.step_id = step_id
        self.metadata = metadata
        self.input_type = input_type
        self.output_type = output_type

        self._status: StepStatus = StepStatus.PENDING
        self._start_time: datetime | None = None
        self._end_time: datetime | None = None
        self._error: str | None = None

    # Public attributes -------------------------------------------------
    @property
    def status(self) -> StepStatus:
        return self._status

    @property
    def start_time(self) -> datetime | None:
        return self._start_time

    @property
    def end_time(self) -> datetime | None:
        return self._end_time

    @property
    def error(self) -> str | None:
        return self._error

    @property
    def duration(self) -> float | None:
        if self._start_time and self._end_time:
            return (self._end_time - self._start_time).total_seconds()
        return None

    # Lifecycle ---------------------------------------------------------
    async def run(
        self,
        input_data: dict[str, Any],
        context: dict[str, Any] | WorkflowContext,
        *,
        pending_request: WorkflowPendingRequest | None = None,
        response: Any | None = None,
    ) -> dict[str, Any]:
        """Validate input/output, apply retry + timeout, and return output dict."""

        attempts = self.metadata.max_retries + 1
        last_error: Exception | None = None

        for attempt in range(attempts):
            self._start_time = datetime.now()
            self._status = StepStatus.RUNNING
            try:
                validated_input = self.input_type(**input_data)
                context_obj = (
                    context
                    if isinstance(context, WorkflowContext)
                    else WorkflowContext.from_state_ref(context, step_id=self.step_id)
                )

                if pending_request is not None:
                    coro = self.resume_from_request(
                        validated_input,
                        context_obj,
                        pending_request,
                        response,
                    )
                else:
                    coro = self.execute(validated_input, context_obj)
                if self.metadata.timeout_seconds:
                    output_model = await asyncio.wait_for(
                        coro, timeout=self.metadata.timeout_seconds
                    )
                else:
                    output_model = await coro

                output = output_model.model_dump()
                self._status = StepStatus.COMPLETED
                self._error = None
                self._end_time = datetime.now()
                return output
            except _WorkflowRequestSuspended:
                self._status = StepStatus.SUSPENDED
                self._error = None
                self._end_time = datetime.now()
                raise
            except Exception as exc:  # noqa: PERF203 - intentionally broad to capture timeouts
                last_error = exc
                self._status = StepStatus.FAILED
                self._error = str(exc)
                self._end_time = datetime.now()
                if attempt < attempts - 1:
                    await asyncio.sleep(1)
                    continue
                raise

        # If loop exits without return/raise, propagate last error for safety
        if last_error:
            raise last_error
        raise RuntimeError("Step failed without an exception")

    @abstractmethod
    async def execute(self, input_data: InputType, context: WorkflowContext) -> OutputType:
        """Run the step's business logic. Must be implemented by subclasses."""

    async def resume_from_request(
        self,
        input_data: InputType,
        context: WorkflowContext,
        request: WorkflowPendingRequest,
        response: Any,
    ) -> OutputType:
        """Resume a previously suspended step with a typed external response."""
        msg = (
            f"Step '{self.step_id}' does not implement request/resume handling "
            f"for request '{request.key}'"
        )
        raise RuntimeError(msg)

    def on_checkpoint_save(self) -> dict[str, Any]:
        """Return step-local state to persist in a workflow checkpoint."""
        return {}

    def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
        """Restore step-local state from a workflow checkpoint."""
        del state

    # Validation helpers -----------------------------------------------
    def validate_input(self, data: dict[str, Any]) -> bool:
        self.input_type(**data)
        return True

    def validate_output(self, data: dict[str, Any]) -> bool:
        self.output_type(**data)
        return True

    # Schema helpers ----------------------------------------------------
    def get_schema(self) -> dict[str, Any]:
        return {
            "step_id": self.step_id,
            "metadata": self.metadata.model_dump(exclude_none=True),
            "input_type": self.input_type.__name__,
            "output_type": self.output_type.__name__,
            "input_schema": self.input_type.model_json_schema(),
            "output_schema": self.output_type.model_json_schema(),
        }


__all__ = ["BaseStep"]
