"""Saga virtual object 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:
    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_object_kwargs(options: SagaWorkflowOptions | None) -> dict[str, Any]:
    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


# ---------------------------------------------------------------------------
# SagaVirtualObject
# ---------------------------------------------------------------------------


class SagaVirtualObject(restate.VirtualObject):
    """A Restate Virtual Object with saga support on exclusive handlers.

    Exclusive handlers receive ``(saga, ctx, input)`` and get automatic
    compensation on ``TerminalError``.  Shared handlers work like normal
    Restate shared handlers — ``(ctx, input)``, no saga.

    Usage::

        cart = create_saga_virtual_object("ShoppingCart")

        @cart.handler()
        async def checkout(saga: SagaContext, ctx: restate.ObjectContext, input_val):
            payment = await charge_payment(saga, {...})
            ctx.clear("items")
            return {"order_id": payment["order_id"]}

        @cart.handler(kind="shared")
        async def get_items(ctx: restate.ObjectSharedContext):
            return await ctx.get("items") or []
    """

    def handler(self, name=None, kind="exclusive", **kwargs):  # type: ignore[override]
        """Register a handler on this virtual object.

        - ``kind="exclusive"`` (default): receives ``(saga, ctx, input)``,
          automatic compensation on TerminalError.
        - ``kind="shared"``: receives ``(ctx, input)``, no saga support,
          concurrent access.
        """
        original_decorator = super().handler(name=name, kind=kind, **kwargs)

        if kind == "shared":
            return original_decorator

        def saga_decorator(
            fn: Callable[..., Awaitable[Any]],
        ) -> Callable[..., Awaitable[Any]]:
            input_model = get_input_model(fn, param_index=2)
            handler_name = name or fn.__name__

            async def wrapped(ctx: restate.ObjectContext, input_val: Any) -> Any:
                input_val = validate_input(input_val, input_model, handler_name)
                saga = SagaContext(ctx=ctx, compensations=[])
                try:
                    return await fn(saga, ctx, input_val)
                except restate.TerminalError:
                    for compensate in reversed(saga.compensations):
                        await compensate()
                    raise

            wrapped.__name__ = fn.__name__
            wrapped.__qualname__ = fn.__qualname__
            return original_decorator(wrapped)

        return saga_decorator


def create_saga_virtual_object(
    name: str,
    options: SagaWorkflowOptions | None = None,
) -> SagaVirtualObject:
    """Create a Restate Virtual Object with saga support.

    Returns a :class:`SagaVirtualObject` — register handlers with
    ``@obj.handler()`` (exclusive, with saga) or
    ``@obj.handler(kind="shared")`` (concurrent, no saga).

    Mount via ``restate.app(services=[obj])``.
    """
    object_kwargs = _build_object_kwargs(options)
    return SagaVirtualObject(name, **object_kwargs)
