from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Awaitable, Callable

from flownet.api.actor_method_ref import ActorMethodRef
from flownet.core.flow_program import FlowProgramRef
from flownet.compiler.transformation import Transformation
from flownet.core.request import FlowRequest
from flownet.core.types import PayloadData


@dataclass(frozen=True)
class OperatorContext:
    local_address: Callable[[], str]
    forward_flowrequest: Callable[[str, FlowRequest], Awaitable[Any]]
    call_operator: Callable[
        [ActorMethodRef, PayloadData, FlowRequest, Transformation | None],
        Awaitable[tuple[bool, Any]],
    ]
    call_stateless: Callable[[Any, PayloadData, FlowRequest, Transformation | None], Awaitable[tuple[bool, Any]]]
    post_process_result: Callable[[FlowRequest, Any, PayloadData], Awaitable[None]]
    request_add_delta: Callable[[FlowRequest, int], Awaitable[None]]
    submit_runtime_coro: Callable[[Any], Any]
    emit_flow_request: Callable[[str, FlowRequest], Awaitable[Any]]
    resolve_flow_program: Callable[[FlowProgramRef], Awaitable[Any]] | None = None
    frame_meta: dict | None = None
    state_services: Any | None = None

    @property
    def current_frame_meta(self) -> dict | None:
        return self.frame_meta

    def clone_request_with_meta(
        self, request: FlowRequest, meta: dict | None
    ) -> FlowRequest:
        from flownet.engine.utils.flow_stack import clone_request_with_frame_meta

        return clone_request_with_frame_meta(request, meta)


FlowEngineContext = OperatorContext


__all__ = [
    "OperatorContext",
    "FlowEngineContext",
]
