"""Saga step creation: create_saga_step (hybrid) and create_saga_step_strict."""

from __future__ import annotations

from typing import Any, Awaitable, Callable

import restate
from restate import RunOptions

from restate_saga._validation import get_input_model, validate_input
from restate_saga.error_registry import resolve_terminal_error
from restate_saga.step_response import StepResponse
from restate_saga.types import SagaContext, SagaStepOptions, StepRetryPolicy

# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _to_run_options(policy: StepRetryPolicy | None) -> RunOptions | None:
    """Convert a StepRetryPolicy to Restate RunOptions."""
    if policy is None:
        return None
    return RunOptions(
        max_attempts=policy.max_retry_attempts,
        max_duration=policy.max_retry_duration,
        initial_retry_interval=policy.initial_retry_interval,
        max_retry_interval=policy.max_retry_interval,
        retry_interval_factor=policy.retry_interval_factor,
    )


async def _run_with_options(
    ctx: restate.Context,
    name: str,
    fn: Callable[[], Any],
    options: RunOptions | None = None,
) -> Any:
    """Call ctx.run / ctx.run_typed depending on whether options are provided."""
    if options is not None:
        return await ctx.run_typed(name, fn, options)
    return await ctx.run(name, fn)


# ---------------------------------------------------------------------------
# create_saga_step  (hybrid compensation)
# ---------------------------------------------------------------------------


def create_saga_step(
    name: str,
    run: Callable[[restate.Context, Any], Awaitable[StepResponse]],
    compensate: Callable[..., Awaitable[None]] | None = None,
    options: SagaStepOptions | None = None,
) -> Callable[[SagaContext, Any], Awaitable[Any]]:
    """Create a saga step with hybrid compensation (Restate + Medusa pattern).

    - Compensation is registered **before** execution so it runs even if the
      step fails.
    - If the step succeeds, compensation receives the ``compensation_data``
      returned by the step.
    - If the step fails or throws, compensation falls back to the original
      *input* as the data argument.
    - ``compensate(data, failed=bool)`` — *failed* indicates whether the
      forward action completed successfully.
    """
    run_options = _to_run_options(options.retry if options else None)
    comp_options = _to_run_options(
        (options.compensation_retry or options.retry) if options else None
    )
    error_mapper = options.as_terminal_error if options else None
    input_model = get_input_model(run)

    async def step(saga: SagaContext, input_val: Any) -> Any:
        input_val = validate_input(input_val, input_model, name)
        ctx = saga.ctx
        step_failed = True
        compensation_data: Any = None

        # 1. Register compensation BEFORE execution (hybrid pattern)
        if compensate is not None:

            async def _compensate() -> None:
                data = compensation_data if compensation_data is not None else input_val

                async def _do_compensate() -> None:
                    await compensate(data, failed=step_failed)

                await _run_with_options(
                    ctx, f"compensate:{name}", _do_compensate, comp_options
                )

            saga.compensations.append(_compensate)

        # 2. Execute forward action
        async def _action() -> dict:
            try:
                response: StepResponse = await run(ctx, input_val)
                return {
                    "output": response.output,
                    "compensation_data": response.compensation_data,
                    "failed": response.failed,
                    "error_message": response.error_message,
                }
            except restate.TerminalError:
                raise
            except Exception as err:
                terminal = resolve_terminal_error(err, error_mapper)
                if terminal:
                    raise terminal from err
                raise

        result = await _run_with_options(ctx, name, _action, run_options)

        # 3. Capture compensation data
        compensation_data = result["compensation_data"]

        # 4. Check for permanent failure
        if result["failed"]:
            raise restate.TerminalError(
                result.get("error_message") or "Step failed permanently"
            )

        # 5. Mark succeeded
        step_failed = False
        return result["output"]

    return step


def saga_step(
    name: str,
    compensate: Callable[..., Awaitable[None]] | None = None,
    options: SagaStepOptions | None = None,
) -> Callable:
    """Decorator version of :func:`create_saga_step` (hybrid compensation).

    Usage::

        async def cancel_order(data, failed=False):
            await api.cancel(data["order_id"])

        @saga_step("CreateOrder", compensate=cancel_order)
        async def create_order(ctx, input_val):
            order_id = await api.create(input_val)
            return StepResponse(output={"order_id": order_id})
    """

    def decorator(
        fn: Callable[[restate.Context, Any], Awaitable[StepResponse]],
    ) -> Callable[[SagaContext, Any], Awaitable[Any]]:
        return create_saga_step(name, fn, compensate, options)

    return decorator


# ---------------------------------------------------------------------------
# create_saga_step_strict  (strict / Medusa-style compensation)
# ---------------------------------------------------------------------------


def create_saga_step_strict(
    name: str,
    run: Callable[[restate.Context, Any], Awaitable[StepResponse]],
    compensate: Callable[..., Awaitable[None]] | None = None,
    options: SagaStepOptions | None = None,
) -> Callable[[SagaContext, Any], Awaitable[Any]]:
    """Create a saga step with strict Medusa-style compensation.

    - Forward action executes **first**.
    - Compensation is registered **after** success — only completed steps get
      compensated.
    - Compensation receives the exact ``compensation_data`` returned by the
      step.
    - Does **not** support ``permanent_failure`` (use ``create_saga_step``
      for that).
    """
    run_options = _to_run_options(options.retry if options else None)
    comp_options = _to_run_options(
        (options.compensation_retry or options.retry) if options else None
    )
    error_mapper = options.as_terminal_error if options else None
    input_model = get_input_model(run)

    async def step(saga: SagaContext, input_val: Any) -> Any:
        input_val = validate_input(input_val, input_model, name)
        ctx = saga.ctx

        # 1. Execute forward action FIRST
        async def _action() -> dict:
            try:
                response: StepResponse = await run(ctx, input_val)
                return {
                    "output": response.output,
                    "compensation_data": response.compensation_data,
                    "failed": response.failed,
                    "error_message": response.error_message,
                }
            except restate.TerminalError:
                raise
            except Exception as err:
                terminal = resolve_terminal_error(err, error_mapper)
                if terminal:
                    raise terminal from err
                raise

        result = await _run_with_options(ctx, name, _action, run_options)

        # 2. If permanent failure, throw without registering compensation
        if result["failed"]:
            raise restate.TerminalError(
                result.get("error_message") or "Step failed permanently"
            )

        # 3. Register compensation AFTER success
        if compensate is not None:
            comp_data = result["compensation_data"]

            async def _compensate() -> None:
                async def _do_compensate() -> None:
                    await compensate(comp_data)

                await _run_with_options(
                    ctx, f"compensate:{name}", _do_compensate, comp_options
                )

            saga.compensations.append(_compensate)

        return result["output"]

    return step


def saga_step_strict(
    name: str,
    compensate: Callable[..., Awaitable[None]] | None = None,
    options: SagaStepOptions | None = None,
) -> Callable:
    """Decorator version of :func:`create_saga_step_strict` (strict compensation).

    Usage::

        async def cancel_order(data):
            await api.cancel(data["order_id"])

        @saga_step_strict("CreateOrder", compensate=cancel_order)
        async def create_order(ctx, input_val):
            order_id = await api.create(input_val)
            return StepResponse(output={"order_id": order_id})
    """

    def decorator(
        fn: Callable[[restate.Context, Any], Awaitable[StepResponse]],
    ) -> Callable[[SagaContext, Any], Awaitable[Any]]:
        return create_saga_step_strict(name, fn, compensate, options)

    return decorator
