# DOCS_REF:
# - docs/user/06-exceptions-reliability.md
# - docs/user/reference/api-surface.md
# - examples/flow_exception_handler_demo.py
# DOCS_SYNC:
# - If handler API/behavior changes, update reliability docs and demo example.

from __future__ import annotations

from contextlib import contextmanager
from typing import Optional

from flownet.api.actor_method_ref import ActorMethodRef
from flownet.core.exceptions import FlowDefinitionError
from flownet.utils.context_vars import ContextSlot


_flowtask_ctx = ContextSlot[object]("flownet_flowtask", default=None)


def _set_current_flowtask(flowtask) -> None:
    _flowtask_ctx.set(flowtask)


def _clear_current_flowtask() -> None:
    _flowtask_ctx.set(None)


def _current_flowtask():
    return _flowtask_ctx.get()


def _resolve_handler(handler) -> Optional[ActorMethodRef]:
    if handler is None:
        return None
    if isinstance(handler, ActorMethodRef):
        return handler
    raise FlowDefinitionError(
        "flow_exception_handler expects ActorMethodRef. "
        "Create an actor first (create_actor/ActorRef) and pass actor.method."
    )


@contextmanager
def flow_exception_handler(handler: Optional[ActorMethodRef]):
    flowtask = _current_flowtask()
    if flowtask is None:
        raise RuntimeError("flow_exception_handler must be used inside flow compilation.")
    resolved = _resolve_handler(handler)
    flowtask._exception_handler_stack.append(resolved)
    try:
        yield
    finally:
        flowtask._exception_handler_stack.pop()


# Backwards-compatible alias
exception_handler = flow_exception_handler

__all__ = [
    "flow_exception_handler",
    "exception_handler",
    "_set_current_flowtask",
    "_clear_current_flowtask",
    "_current_flowtask",
]
