from __future__ import annotations

import traceback
from typing import Any, Awaitable, Callable

from flownet.api.actor_method_ref import ActorMethodRef
from flownet.compiler.transformation import Transformation
from flownet.core.exceptions import (
    ExceptionContext,
    ExceptionDecision,
    ExceptionEvent,
    FlowException,
)
from flownet.core.request import FlowRequest
from flownet.core.types import PayloadData


class FlowExceptionResolver:
    """Resolve exceptions raised during flow execution."""

    def __init__(
        self,
        *,
        send_flow_error: Callable[..., Awaitable[None]],
        handle_actor_call: Callable[..., Awaitable[Any]],
        post_process_result: Callable[[FlowRequest, Any, PayloadData], Awaitable[None]],
        resolve_ack_target: Callable[[FlowRequest], str | None] | None = None,
    ) -> None:
        self._send_flow_error_fn = send_flow_error
        self._handle_actor_call = handle_actor_call
        self._post_process_result = post_process_result
        self._resolve_ack_target = resolve_ack_target

    async def send_flow_error(
        self,
        request: FlowRequest,
        *,
        error_type: str,
        message: str,
        traceback_text: str,
        event: ExceptionEvent | None,
    ) -> None:
        resolve_ack_target = self._resolve_ack_target
        if resolve_ack_target is None:
            return
        coordinator_target = resolve_ack_target(request)
        if coordinator_target is None:
            return
        await self._send_flow_error_fn(
            request.origin_address,
            coordinator_target,
            request.request_id,
            error_type=error_type,
            message=message,
            traceback_text=traceback_text,
            event=event,
        )

    def build_exception_event(
        self,
        request: FlowRequest,
        rmf: ActorMethodRef,
        payload: PayloadData,
        transformation: Transformation | None,
        exc: BaseException,
    ) -> ExceptionEvent:
        context = ExceptionContext(
            request_id=request.request_id,
            phase="actor_call",
            actor_id=rmf.actor_id,
            method=rmf.method,
            address=rmf.address,
            transformation_id=transformation.trans_id if transformation else None,
            payload=payload,
            stack_depth=len(request.flow_stack.stack),
            metadata={
                "transformation_type": transformation.type if transformation else None,
            },
        )
        return ExceptionEvent(
            error_type=type(exc).__name__,
            message=str(exc),
            traceback=traceback.format_exc(),
            context=context,
            error=exc,
        )

    async def resolve_exception_decision(
        self,
        event: ExceptionEvent,
        transformation: Transformation | None,
    ) -> ExceptionDecision | None:
        handlers: list[ActorMethodRef] = []
        if transformation is not None:
            stack = getattr(transformation, "exception_handler_stack", [])
            handlers = [handler for handler in stack if handler is not None]

        for handler in reversed(handlers):
            decision = await self._handle_actor_call(handler, event)
            if decision is None:
                continue
            if not isinstance(decision, ExceptionDecision):
                raise TypeError(
                    f"Exception handler must return ExceptionDecision, got {type(decision)}"
                )
            if decision.action == "propagate":
                continue
            return decision
        return None

    async def handle_flow_exception(
        self,
        request: FlowRequest,
        rmf: ActorMethodRef,
        payload: PayloadData,
        transformation: Transformation | None,
        exc: BaseException,
    ) -> None:
        event = self.build_exception_event(request, rmf, payload, transformation, exc)
        decision = await self.resolve_exception_decision(event, transformation)

        if decision is None:
            await self.send_flow_error(
                request,
                error_type=event.error_type,
                message=event.message,
                traceback_text=event.traceback,
                event=event,
            )
            await self._post_process_result(request, None, payload)
            raise FlowException(event)

        if decision.action == "abort":
            await self._post_process_result(request, None, payload)
            return

        if decision.action == "fallback":
            payloads = decision.payloads if decision.payloads is not None else []
            await self._post_process_result(request, payloads, payload)
            return

        if decision.action == "propagate":
            await self.send_flow_error(
                request,
                error_type=event.error_type,
                message=event.message,
                traceback_text=event.traceback,
                event=event,
            )
            await self._post_process_result(request, None, payload)
            raise FlowException(event)

        raise ValueError(f"Unknown ExceptionDecision action: {decision.action}")


__all__ = ["FlowExceptionResolver"]
