from __future__ import annotations

import asyncio
import inspect
import os
import traceback
from typing import Any

from flownet.api.actor_method_ref import ActorMethodRef
from flownet.compiler.transformation import Transformation
from flownet.core.flow_program import FlowProgramRef, ResolvedFlowProgram
from flownet.core.exceptions import FlowException, FlowDefinitionError
from flownet.core.request import FlowRequest, TransformationFrame
from flownet.core.stateless import StatelessOpRef
from flownet.core.types import PayloadData
from flownet.engine.advancer import FlowAdvancer
from flownet.engine.context import OperatorContext
from flownet.engine.exception_resolver import FlowExceptionResolver
from flownet.engine.operators.base import OperatorAction, OperatorHandler
from flownet.engine.operator_registry import OperatorRegistry
from flownet.engine.port import EnginePort
from flownet.runtime.kernel import try_get_runtime_kernel


def _rpc_debug_logs_enabled() -> bool:
    value = (os.environ.get("FLOWNET_RPC_DEBUG_LOGS") or "").strip().lower()
    return value in {"1", "true", "yes", "on"}


class FlowEngine:
    def __init__(
        self,
        port: EnginePort,
        operator_registry: OperatorRegistry | None = None,
    ):
        self._io = port
        self._submit_runtime_coro = port.submit_runtime_coro
        self._operator_registry = operator_registry or OperatorRegistry.with_defaults()
        self._default_operator_handler = self._operator_registry.default_handler
        self._advancer = FlowAdvancer(
            send_ack=self._io.send_ack,
            submit_runtime_coro=self._submit_runtime_coro,
            handle_flow_request=self.handle_flow_request,
        )
        self._exception_resolver = FlowExceptionResolver(
            send_flow_error=self._io.send_flow_error,
            handle_actor_call=self._io.handle_actor_call,
            post_process_result=self._advancer.post_process_result,
            resolve_ack_target=self._advancer.resolve_ack_coordinator_target,
        )
        self._flow_program_cache: dict[tuple[str, str], ResolvedFlowProgram] = {}

    # 我们是不是可以做一个operator解析器？
    def _get_operator_handler(
        self, transformation_type: str | None
    ) -> OperatorHandler:
        return self._operator_registry.resolve(transformation_type)

    def operator_stats_snapshot(self) -> dict[str, dict[str, int]]:
        return self._operator_registry.stats_snapshot()

    # 这是主要的方法入口
    async def handle_flow_request(self, request: FlowRequest):
        while True:
            top_frame = request.flow_stack.top()
            if top_frame is None:
                return
            if not isinstance(top_frame, TransformationFrame):
                raise TypeError(f"Unexpected flow frame: {type(top_frame)}")

            transformation = top_frame.transformation
            operator_key = transformation.operator_key or transformation.type
            handler = self._get_operator_handler(operator_key)
            frame_meta = self._merge_operator_meta(
                transformation.operator_config, getattr(top_frame, "meta", None)
            )
            handler_context = self._build_operator_context(frame_meta)
            try:
                action = await handler.handle_request(handler_context, request, top_frame)
            except asyncio.CancelledError:
                raise
            except (KeyboardInterrupt, SystemExit):
                raise
            except Exception as exc:
                await self._handle_request_exception(request, top_frame, exc)
                return
            if action == OperatorAction.CONTINUE:
                continue
            return

    # 这个便捷方法可以去掉
    def local_address(self) -> str:
        return self._io.local_address()

    # 这个便捷wiring方法也可以被吸收掉
    async def forward_flowrequest(self, address: str, request: FlowRequest):
        local_address = self._io.local_address()
        if _rpc_debug_logs_enabled():
            print("[runtime] routing request to remote actor:", address, "local", local_address)
        return await self._io.forward_flowrequest(address, request)

    async def emit_flow_request(self, address: str, request: FlowRequest):
        return await self._io.forward_flowrequest(address, request)

    # 这一个方法和actor dispatcher的功能重复了，我们应该带一个context，然后调用actor dispatcher
    async def call_local(
        self,
        rmf: ActorMethodRef,
        payload: PayloadData,
        request: FlowRequest,
        transformation: Transformation | None,
    ) -> tuple[bool, Any]:
        try:
            result = self._io.invoke_local(rmf, payload, request)
        except asyncio.CancelledError:
            raise
        except (KeyboardInterrupt, SystemExit):
            raise
        except Exception as exc:
            await self._exception_resolver.handle_flow_exception(
                request, rmf, payload, transformation, exc
            )
            return False, None
        return True, result

    async def call_stateless(
        self,
        op_ref: StatelessOpRef,
        payload: PayloadData,
        request: FlowRequest,
        transformation: Transformation | None,
    ) -> tuple[bool, Any]:
        try:
            fn = await self._io.resolve_stateless_op(op_ref)
            from flownet.runtime.contexts.flow import flow_context_scope
            with flow_context_scope(request):
                result = fn(*payload.args, **payload.kwargs)
                if inspect.isawaitable(result):
                    result = await result
        except asyncio.CancelledError:
            raise
        except (KeyboardInterrupt, SystemExit):
            raise
        except Exception as exc:
            rmf = ActorMethodRef(
                address=op_ref.origin or self.local_address(),
                actor_id="__stateless__",
                method=op_ref.name or op_ref.qualname or op_ref.op_id,
            )
            await self._exception_resolver.handle_flow_exception(
                request, rmf, payload, transformation, exc
            )
            return False, None
        return True, result

    async def resolve_flow_program(
        self,
        program_ref: FlowProgramRef,
    ) -> Any:
        cache_key = (program_ref.flow_uri, program_ref.version)
        cached = self._flow_program_cache.get(cache_key)
        if cached is not None:
            return cached

        kernel = try_get_runtime_kernel()
        if kernel is None:
            raise RuntimeError("Runtime kernel not available for FlowProgramRef resolution.")
        flow_services = kernel.require_flows()
        try:
            payload = flow_services.resolve_flow_program_payload(
                program_ref.flow_uri,
                version=program_ref.version,
            )
            catalog_record = flow_services.get_flow_program(
                program_ref.flow_uri,
                version=program_ref.version,
            )
            resolved = self._build_resolved_flow_program(
                program_ref=program_ref,
                payload=payload,
                catalog_record=catalog_record,
            )
            self._flow_program_cache[cache_key] = resolved
            return resolved
        except Exception as exc:
            if not self._is_flow_program_not_found(exc):
                raise

            resolved = await self._resolve_flow_program_from_attached_node(
                kernel=kernel,
                program_ref=program_ref,
            )
            if resolved is None:
                raise
            self._flow_program_cache[cache_key] = resolved
            return resolved

    @staticmethod
    def _is_flow_program_not_found(exc: Exception) -> bool:
        text = str(exc)
        return "FlowProgram not found" in text

    @staticmethod
    def _build_resolved_flow_program(
        *,
        program_ref: FlowProgramRef,
        payload: Any,
        catalog_record: Any,
    ) -> ResolvedFlowProgram:
        program_id: str | None = None
        program_hash: str | None = None
        if isinstance(catalog_record, dict):
            maybe_program_id = catalog_record.get("program_id")
            maybe_program_hash = catalog_record.get("program_hash")
            if isinstance(maybe_program_id, str) and maybe_program_id.strip():
                program_id = maybe_program_id
            if isinstance(maybe_program_hash, str) and maybe_program_hash.strip():
                program_hash = maybe_program_hash
        return ResolvedFlowProgram(
            payload=payload,
            flow_uri=program_ref.flow_uri,
            version=program_ref.version,
            program_id=program_id,
            program_hash=program_hash,
        )

    async def _resolve_flow_program_from_attached_node(
        self,
        *,
        kernel: Any,
        program_ref: FlowProgramRef,
    ) -> ResolvedFlowProgram | None:
        target_address = str(os.environ.get("FLOWNET_NODE_ADDR") or "").strip()
        if not target_address:
            return None

        local_address = str(kernel.local_address() or "").strip()
        if local_address and target_address == local_address:
            return None

        actors = kernel.require_actors()
        payload = await actors.call(
            ActorMethodRef(
                address=target_address,
                actor_id="__container__",
                method="resolve_flow_program_payload",
            ),
            program_ref.flow_uri,
            version=program_ref.version,
        )
        catalog_record = await actors.call(
            ActorMethodRef(
                address=target_address,
                actor_id="__container__",
                method="get_flow_program",
            ),
            program_ref.flow_uri,
            version=program_ref.version,
        )
        return self._build_resolved_flow_program(
            program_ref=program_ref,
            payload=payload,
            catalog_record=catalog_record,
        )

    async def _handle_request_exception(
        self,
        request: FlowRequest,
        frame: TransformationFrame,
        exc: Exception,
    ) -> None:
        if isinstance(exc, FlowException):
            return
        if isinstance(exc, FlowDefinitionError):
            await self._exception_resolver.send_flow_error(
                request,
                error_type=type(exc).__name__,
                message=str(exc),
                traceback_text=traceback.format_exc(),
                event=None,
            )
            return
        try:
            await self._exception_resolver.handle_flow_exception(
                request, frame.transformation.target, frame.payload, frame.transformation, exc
            )
        except FlowException as flow_exc:
            await self._exception_resolver.send_flow_error(
                request,
                error_type=flow_exc.event.error_type,
                message=flow_exc.event.message,
                traceback_text=flow_exc.event.traceback,
                event=flow_exc.event,
            )
            return
    def _build_operator_context(self, frame_meta: dict | None) -> OperatorContext:
        state_services = None
        kernel = try_get_runtime_kernel()
        if kernel is not None and getattr(kernel, "state", None) is not None:
            state_services = kernel.state
        return OperatorContext(
            local_address=self.local_address,
            forward_flowrequest=self.forward_flowrequest,
            call_operator=self.call_local,
            call_stateless=self.call_stateless,
            post_process_result=self._advancer.post_process_result,
            request_add_delta=self._advancer.apply_request_delta,
            submit_runtime_coro=self._submit_runtime_coro,
            emit_flow_request=self.emit_flow_request,
            resolve_flow_program=self.resolve_flow_program,
            frame_meta=frame_meta,
            state_services=state_services,
        )

    @staticmethod
    def _merge_operator_meta(base: dict | None, override: dict | None) -> dict | None:
        if base is None and override is None:
            return None
        if base is None:
            return FlowEngine._clone_meta(override)
        if override is None:
            return FlowEngine._clone_meta(base)
        if not isinstance(base, dict) or not isinstance(override, dict):
            return override
        merged = FlowEngine._clone_meta(base)
        for key, value in override.items():
            if isinstance(value, dict) and isinstance(merged.get(key), dict):
                merged[key] = FlowEngine._merge_operator_meta(merged.get(key), value)
            else:
                merged[key] = value
        return merged

    @staticmethod
    def _clone_meta(meta: dict | None) -> dict | None:
        if meta is None:
            return None
        if not isinstance(meta, dict):
            return meta
        cloned: dict = {}
        for key, value in meta.items():
            if isinstance(value, dict):
                cloned[key] = FlowEngine._clone_meta(value)
            else:
                cloned[key] = value
        return cloned
