"""Workflow step decorator and tracking."""

from __future__ import annotations

import functools
from datetime import datetime, timezone
from typing import Any, Callable, Dict, Optional

from databridge_core.sdk.core.exceptions import NotImplementedStubError as _NotImplementedStubError
from databridge_core.sdk.core.types import Status, StepOutcome


def workflow_step(name: str, description: str = ""):
    """Decorator that wraps a workflow method as a tracked step.

    Records timing, status, and audit entry for each step execution.

    Usage::

        class MyWorkflow(BaseWorkflow):
            @workflow_step("load_data", "Load and parse input files")
            def _load(self, params, ctx):
                ...
                return {"rows": 100}

    Args:
        name: Step identifier (used in audit trail and StepOutcome).
        description: Human-readable description of the step.
    """

    def decorator(fn: Callable) -> Callable:
        @functools.wraps(fn)
        def wrapper(self, params: Dict[str, Any], ctx=None) -> StepOutcome:
            started = datetime.now(timezone.utc)

            if ctx:
                ctx.audit.log(
                    action="step_start",
                    workflow=getattr(self, "name", ""),
                    step_name=name,
                )

            try:
                data = fn(self, params, ctx)
                outcome = StepOutcome(
                    step_name=name,
                    status=Status.SUCCESS,
                    data=data if isinstance(data, dict) else {},
                    started_at=started.isoformat(),
                    finished_at=datetime.now(timezone.utc).isoformat(),
                )
            except (NotImplementedError, _NotImplementedStubError) as exc:
                outcome = StepOutcome(
                    step_name=name,
                    status=Status.NOT_IMPLEMENTED,
                    started_at=started.isoformat(),
                    finished_at=datetime.now(timezone.utc).isoformat(),
                    error=str(exc),
                )
            except Exception as exc:
                outcome = StepOutcome(
                    step_name=name,
                    status=Status.FAILED,
                    started_at=started.isoformat(),
                    finished_at=datetime.now(timezone.utc).isoformat(),
                    error=str(exc),
                )

            if ctx:
                ctx.audit.log_step(
                    step_name=name,
                    workflow=getattr(self, "name", ""),
                    status=outcome.status.value,
                    data=outcome.data,
                )

            return outcome

        wrapper._step_name = name
        wrapper._step_description = description
        return wrapper

    return decorator
