from __future__ import annotations

from enum import Enum
from typing import TYPE_CHECKING, Any

from flownet.api.actor_method_ref import ActorMethodRef
from flownet.compiler.flow_def import BoundFlow, FlowDef
from flownet.compiler.flowtask import FlowTask
from flownet.core.request import FlowRequest, TransformationFrame
from flownet.core.stateless import StatelessOpRef
from flownet.core.types import PayloadData
from flownet.engine.context import OperatorContext


if TYPE_CHECKING:
    from flownet.compiler.transformation import Transformation


class OperatorAction(Enum):
    CONTINUE = "continue"
    RETURN = "return"



class OperatorHandler:
    """Default operator behavior for map/flatmap/tap."""

    def resolve_target(self, transformation_frame: TransformationFrame) -> ActorMethodRef:
        return transformation_frame.transformation.target

    async def handle_request(
        self,
        ctx: OperatorContext,
        request: FlowRequest,
        transformation_frame: TransformationFrame,
    ) -> OperatorAction:
        transformation = transformation_frame.transformation
        target = self.resolve_target(transformation_frame)
        if isinstance(target, StatelessOpRef):
            if (
                target.scope == "local"
                and target.origin is not None
                and target.origin != ctx.local_address()
            ):
                await ctx.forward_flowrequest(target.origin, request)
                return OperatorAction.RETURN
            has_result, result = await ctx.call_stateless(
                target, transformation_frame.payload, request, transformation
            )
        else:
            rmf = target
            if rmf.address != ctx.local_address():
                await ctx.forward_flowrequest(rmf.address, request)
                return OperatorAction.RETURN
            has_result, result = await ctx.call_operator(
                rmf, transformation_frame.payload, request, transformation
            )
        if not has_result:
            return OperatorAction.RETURN
        return await self.handle_result(
            ctx, request, transformation, transformation_frame.payload, result
        )

    async def handle_result(
        self,
        ctx: OperatorContext,
        request: FlowRequest,
        transformation: Transformation,
        payload: PayloadData,
        result: Any,
    ) -> OperatorAction:
        if isinstance(result, (FlowTask, FlowDef, BoundFlow)):
            raise TypeError(
                "Subgraph results are only allowed for pipe(); use pipe() instead of map/flatmap."
            )
        await ctx.post_process_result(request, result, payload)
        return OperatorAction.RETURN

