"""Saga workflow creation."""

from __future__ import annotations

from typing import Any, Awaitable, Callable

import restate

from restate_saga._validation import get_input_model, validate_input
from restate_saga.types import SagaContext, SagaWorkflowOptions, WorkflowRetryPolicy


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


def _to_invocation_retry_policy(
    policy: WorkflowRetryPolicy | None,
) -> restate.InvocationRetryPolicy | None:
    """Convert WorkflowRetryPolicy to Restate InvocationRetryPolicy."""
    if policy is None:
        return None
    return restate.InvocationRetryPolicy(
        max_attempts=policy.max_attempts,
        on_max_attempts=policy.on_max_attempts,
        initial_interval=policy.initial_interval,
        exponentiation_factor=policy.exponentiation_factor,
        max_interval=policy.max_interval,
    )


def _build_service_kwargs(options: SagaWorkflowOptions | None) -> dict[str, Any]:
    """Build keyword arguments for ``restate.Service(...)`` from workflow options."""
    if options is None:
        return {}
    kwargs: dict[str, Any] = {}
    retry = _to_invocation_retry_policy(options.retry_policy)
    if retry is not None:
        kwargs["invocation_retry_policy"] = retry
    if options.idempotency_retention is not None:
        kwargs["idempotency_retention"] = options.idempotency_retention
    if options.journal_retention is not None:
        kwargs["journal_retention"] = options.journal_retention
    if options.inactivity_timeout is not None:
        kwargs["inactivity_timeout"] = options.inactivity_timeout
    if options.abort_timeout is not None:
        kwargs["abort_timeout"] = options.abort_timeout
    if options.ingress_private is not None:
        kwargs["ingress_private"] = options.ingress_private
    return kwargs


# ---------------------------------------------------------------------------
# SagaWorkflowService
# ---------------------------------------------------------------------------


class SagaWorkflowService(restate.Service):
    """A Restate service with saga support and a typed ``run_as_step`` method.

    Can be:

    1. Mounted via ``restate.app(services=[workflow])``
    2. Embedded in another workflow via ``workflow.run_as_step(parent_saga, input)``
    """

    _saga_handler: Callable[[SagaContext, Any], Awaitable[Any]]
    _saga_input_model: type | None

    async def run_as_step(self, parent_saga: SagaContext, input_val: Any) -> Any:
        """Run this workflow as a step in a parent saga.

        Compensations are registered in the parent's compensation stack so
        they execute if the parent saga fails later.
        """
        input_val = validate_input(input_val, self._saga_input_model, self.name)
        return await self._saga_handler(parent_saga, input_val)

    def __repr__(self) -> str:
        return f"SagaWorkflowService({self.name!r})"


# ---------------------------------------------------------------------------
# create_saga_workflow
# ---------------------------------------------------------------------------


def create_saga_workflow(
    name: str,
    handler: Callable[[SagaContext, Any], Awaitable[Any]],
    options: SagaWorkflowOptions | None = None,
) -> SagaWorkflowService:
    """Create a saga workflow that orchestrates steps with automatic compensation.

    The returned object is a :class:`SagaWorkflowService` that can be:

    1. Mounted via ``restate.app(services=[workflow])``
    2. Embedded in another workflow via ``workflow.run_as_step(parent_saga, input)``

    On ``TerminalError``, all registered compensations run in reverse order
    before the error propagates.
    """
    service_kwargs = _build_service_kwargs(options)
    service = SagaWorkflowService(name, **service_kwargs)
    service._saga_handler = handler
    service._saga_input_model = get_input_model(handler)

    @service.handler()
    async def run(ctx: restate.Context, input_val: Any) -> Any:
        input_val = validate_input(input_val, service._saga_input_model, name)
        saga = SagaContext(ctx=ctx, compensations=[])
        try:
            return await handler(saga, input_val)
        except restate.TerminalError:
            for compensate in reversed(saga.compensations):
                await compensate()
            raise

    return service


def saga_workflow(
    name: str,
    options: SagaWorkflowOptions | None = None,
) -> Callable[..., SagaWorkflowService]:
    """Decorator version of :func:`create_saga_workflow`.

    Usage::

        @saga_workflow("CheckoutWorkflow")
        async def checkout_workflow(saga, input_val):
            order = await create_order(saga, input_val)
            payment = await process_payment(saga, {"order_id": order["order_id"]})
            return {"order_id": order["order_id"], "payment_id": payment["payment_id"]}
    """

    def decorator(
        fn: Callable[[SagaContext, Any], Awaitable[Any]],
    ) -> SagaWorkflowService:
        return create_saga_workflow(name, fn, options)

    return decorator
