"""Saga Restate Workflow creation.

Restate Workflows are long-running processes identified by a unique workflow key.
The ``main`` handler executes exactly once per key, with automatic deduplication.
Additional handlers can be used for signals and queries while the workflow runs.
"""

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_workflow_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


# ---------------------------------------------------------------------------
# SagaRestateWorkflow
# ---------------------------------------------------------------------------


class SagaRestateWorkflow(restate.Workflow):
    """A Restate Workflow with saga support on the main handler.

    The ``main`` handler receives ``(saga, ctx, input)`` and gets automatic
    compensation on ``TerminalError``.  Additional handlers registered via
    ``@wf.handler()`` are shared handlers — ``(ctx, input)``, no saga.

    Can also be embedded in a parent saga via ``wf.run_as_step(parent_saga, input)``.

    Usage::

        wf = create_saga_restate_workflow("OrderWorkflow")

        @wf.main()
        async def run(saga: SagaContext, ctx: restate.WorkflowContext, 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"]}

        @wf.handler()
        async def get_status(ctx: restate.WorkflowSharedContext):
            return await ctx.get("status")
    """

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

    def __init__(self, name: str, **kwargs: Any) -> None:
        super().__init__(name, **kwargs)
        self._saga_handler = None
        self._saga_input_model = None

    def main(self, name: str | None = None, **kwargs: Any):  # type: ignore[override]
        """Register the main workflow handler with saga support.

        The decorated function receives ``(saga, ctx, input)`` and gets
        automatic compensation on ``TerminalError``.
        """
        original_decorator = super().main(name=name, **kwargs)

        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__

            self._saga_handler = fn
            self._saga_input_model = input_model

            async def wrapped(ctx: restate.WorkflowContext, 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

    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.
        """
        if self._saga_handler is None:
            raise RuntimeError(
                f"SagaRestateWorkflow({self.name!r}) has no @main handler registered"
            )
        input_val = validate_input(input_val, self._saga_input_model, self.name)
        return await self._saga_handler(parent_saga, parent_saga.ctx, input_val)

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


def create_saga_restate_workflow(
    name: str,
    options: SagaWorkflowOptions | None = None,
) -> SagaRestateWorkflow:
    """Create a Restate Workflow with saga support.

    Returns a :class:`SagaRestateWorkflow` — register the main handler with
    ``@wf.main()`` (with saga) and additional handlers with
    ``@wf.handler()`` (shared, no saga).

    Mount via ``restate.app(services=[wf])``.
    """
    workflow_kwargs = _build_workflow_kwargs(options)
    return SagaRestateWorkflow(name, **workflow_kwargs)
