import functools
from typing import Any, Awaitable, Callable

import pgqrs
from ._pgqrs import Run


class WorkflowDef:
    """Stores workflow metadata generated by the decorator."""

    __slots__ = ("name", "run", "_func")

    def __init__(
        self,
        name: str,
        run: Callable[..., Awaitable[Any]],
        func: Callable[..., Awaitable[Any]],
    ) -> None:
        self.name = name
        self.run = run
        self._func = func
        # Keep `__doc__` / `__name__` on the underlying function; avoid slot conflicts.

    def __call__(self, *args: Any, **kwargs: Any) -> Awaitable[Any]:
        return self.run(*args, **kwargs)


def workflow(*, name: str):
    """Decorator to mark a function as a durable workflow entry point."""

    if not isinstance(name, str):
        raise TypeError("Workflow name must be a string")

    if not name:
        raise ValueError("Workflow name cannot be empty")

    def decorator(func: Callable[..., Awaitable[Any]]) -> WorkflowDef:
        @functools.wraps(func)
        async def runner(ctx: Run, *args: Any, **kwargs: Any) -> Any:
            if not isinstance(ctx, Run):
                raise TypeError(
                    f"First argument to a workflow must be a Run instance, got {type(ctx)}"
                )

            try:
                await ctx.start()
            except pgqrs.ValidationError as e:
                if "terminal SUCCESS" not in str(e):
                    raise

            try:
                # If invoked by the DequeueBuilder wrapper, args/kwargs might be empty
                # We can extract the input payload from the context's record
                if not args and not kwargs and hasattr(ctx, "input") and ctx.input:
                    payload = ctx.input
                    if isinstance(payload, dict) and "input" in payload:
                        args = (payload["input"],)
                    else:
                        args = (payload,)

                result = await func(ctx, *args, **kwargs)
                await ctx.success(result)
                return result
            except (
                pgqrs.TransientStepError,
                pgqrs.StepNotReadyError,
                pgqrs.RetriesExhaustedError,
                pgqrs.PausedError,
            ):
                raise
            except Exception as exc:
                await ctx.fail(str(exc))
                raise

        return WorkflowDef(name=name, run=runner, func=func)

    return decorator


def step(func):
    """Decorator to mark a function as a durable workflow step."""

    @functools.wraps(func)
    async def wrapper(ctx, *args, **kwargs):
        if not isinstance(ctx, Run):
            raise TypeError(
                f"First argument to a step must be a Run instance, got {type(ctx)}"
            )

        step_name = func.__name__

        current_time = getattr(ctx, "current_time", None)
        step_result = await ctx.acquire_step(step_name, current_time=current_time)

        if step_result.status == pgqrs.StepResultStatus.Skipped:
            return step_result.value

        elif step_result.status == pgqrs.StepResultStatus.Execute:
            guard = step_result.guard
            try:
                result = await func(ctx, *args, **kwargs)
                await guard.success(result)
                return result
            except pgqrs.TransientStepError as e:
                code = getattr(e, "code", "TRANSIENT_ERROR")
                retry_after = getattr(e, "retry_after", None)
                await guard.fail_transient(code, str(e), retry_after)
                raise
            except (
                pgqrs.StepNotReadyError,
                pgqrs.RetriesExhaustedError,
                pgqrs.PausedError,
            ):
                raise
            except Exception as e:
                await guard.fail(str(e))
                raise
        else:
            raise RuntimeError(f"Unknown step status: {step_result.status}")

    return wrapper
