import asyncio, time
import heapq
import os
import warnings
from typing import Callable

from flownet.config import FlownetConfig
from flownet.core.request import FlownetPacket
from flownet.runtime import loops
from flownet.runtime.contexts.actor import system_calls_allowed
from flownet.runtime.transport.tcp_transport import TcpPacketTransport
from flownet.runtime.transport.transport_hub import TransportHub
from flownet.runtime.transport.transport_selector import TransportSelector
from flownet.runtime.transport.shm_transport import ShmPacketTransport
from flownet.runtime.transport.uds_transport import UdsPacketTransport


_REMOTE_ERROR_KEY = "__flownet_remote_error__"


def _decode_remote_error(payload: object) -> str | None:
    if not isinstance(payload, dict):
        return None
    marker = payload.get(_REMOTE_ERROR_KEY)
    if not isinstance(marker, dict):
        return None
    message = marker.get("message")
    if not isinstance(message, str) or not message.strip():
        return "Remote actor call failed"
    error_type = marker.get("type")
    if isinstance(error_type, str) and error_type.strip():
        return f"Remote actor call failed [{error_type}]: {message}"
    return f"Remote actor call failed: {message}"


def _resolve_default_timeout_seconds(config: FlownetConfig) -> float:
    explicit = getattr(config, "rpc_call_timeout_seconds", None)
    if explicit is not None:
        try:
            return max(0.1, float(explicit))
        except Exception:
            pass

    raw_env = os.environ.get("FLOWNET_RPC_TIMEOUT_SECONDS", "").strip()
    if raw_env:
        try:
            return max(0.1, float(raw_env))
        except Exception:
            pass
    return 30.0


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


# 注意，只能在kernel runtime loop中使用
class RemoteCallManager:
    def __init__(
        self,
        config: FlownetConfig,
        handler: callable,
        *,
        log_fn: Callable[[str], None] | None = None,
    ):
        self.pending: dict[str, asyncio.Future] = {}
        self.timestamps: dict[str, float] = {}
        self._deadline_heap: list[tuple[float, str]] = []
        self._deadlines: dict[str, float] = {}
        if log_fn is not None:
            self._log_fn = log_fn
        elif _rpc_debug_logs_enabled():
            self._log_fn = print
        else:
            self._log_fn = None
        self._stats: dict[str, int] = {
            "received_packets": 0,
            "resolved_packets": 0,
            "legacy_reply_packets": 0,
            "missing_future_packets": 0,
            "timeout_rejections": 0,
        }
        self.default_timeout = _resolve_default_timeout_seconds(config)
        local_caps = getattr(config, "transport_caps", None)
        selector = TransportSelector(
            local_host=config.host,
            local_port=config.port,
            transport_mode=getattr(config, "transport_mode", "auto"),
            local_ipc_dir=getattr(config, "local_ipc_dir", "/tmp/flownet"),
            local_host_id=getattr(config, "host_id", None),
            local_uds_path=getattr(config, "uds_path", None),
            local_shm_path=getattr(config, "shm_path", None),
            local_transport_caps=local_caps,
            shm_threshold_bytes=getattr(config, "transport_shm_threshold_bytes", 262144),
        )
        max_inflight_handlers = getattr(config, "transport_max_inflight_handlers", 256)
        tcp_transport = TcpPacketTransport(
            host=config.host,
            port=config.port,
            handler=self.receive_packet,
            max_inflight_handlers=max_inflight_handlers,
        )
        mode = (getattr(config, "transport_mode", "auto") or "auto").strip().lower()
        uds_transport = None
        shm_transport = None
        if mode != "tcp" and "uds" in selector.local_transport_caps:
            uds_transport = UdsPacketTransport(
                uds_path=selector.local_uds_path,
                handler=self.receive_packet,
                max_inflight_handlers=max_inflight_handlers,
            )
        if mode == "hetero_auto" and "shm" in selector.local_transport_caps:
            shm_transport = ShmPacketTransport(
                shm_path=selector.local_shm_path,
                handler=self.receive_packet,
                max_payload_bytes=getattr(config, "transport_shm_max_payload_bytes", 67108864),
                max_inflight_handlers=max_inflight_handlers,
            )
        self.transport_hub = TransportHub(
            selector=selector,
            tcp_transport=tcp_transport,
            uds_transport=uds_transport,
            shm_transport=shm_transport,
            backend_cooldown_seconds=getattr(config, "transport_backend_cooldown_seconds", 1.0),
            backend_cooldown_max_entries=getattr(
                config, "transport_backend_cooldown_max_entries", 4096
            ),
            backend_cooldown_cleanup_interval=getattr(
                config, "transport_backend_cooldown_cleanup_interval", 128
            ),
        )
        # Keep compatibility for callers that still access network_interface.host/port.
        self.network_interface = tcp_transport
        self.runtime_handler = handler
        loops.submit_runtime_coro(self.timeout_watcher())

    def set_endpoint_lookup(self, endpoint_lookup=None) -> None:
        self.transport_hub.selector.set_endpoint_lookup(endpoint_lookup)

    def local_endpoint_metadata(self) -> dict[str, object]:
        selector = self.transport_hub.selector
        return {
            "address": selector.local_address,
            "host_id": selector.local_host_id,
            "uds_path": selector.local_uds_path,
            "shm_path": selector.local_shm_path,
            "transport_caps": list(selector.local_transport_caps),
        }
        
    async def send_packet(self, address:str, packet:FlownetPacket):
        if not system_calls_allowed():
            warnings.warn(
                "Direct packet send detected outside system context. "
                "Prefer ActorContext/service/source emit APIs.",
                RuntimeWarning,
                stacklevel=2,
            )
        if packet.reply == False:
            await self.transport_hub.send_packet(address, packet)
            return None
        future = self.register(packet.uuid)
        try:
            await self.transport_hub.send_packet(address, packet)
        except Exception as exc:
            # Transport send failed before any response path can resolve this
            # registered future. Cancel it to avoid an unobserved exception
            # on a detached asyncio Future.
            self._cancel_pending(packet.uuid)
            raise
        return await future

    def transport_metrics_snapshot(self) -> dict[str, object]:
        selector = self.transport_hub.selector
        snapshot: dict[str, object] = {
            "transport_mode": selector.transport_mode,
            "local_address": selector.local_address,
            "local_uds_path": selector.local_uds_path,
            "local_shm_path": selector.local_shm_path,
            "local_host_id": selector.local_host_id,
            "local_transport_caps": list(selector.local_transport_caps),
        }
        snapshot.update(selector.metrics_snapshot())
        snapshot.update(self.transport_hub.metrics_snapshot())
        snapshot.update(self.call_stats_snapshot())
        return snapshot

    def call_stats_snapshot(self) -> dict[str, int]:
        return dict(getattr(self, "_stats", {}))

    def _log(self, message: str) -> None:
        log_fn = getattr(self, "_log_fn", None)
        if log_fn is not None:
            log_fn(message)

    def _bump(self, key: str, delta: int = 1) -> None:
        stats = getattr(self, "_stats", None)
        if stats is None:
            stats = {}
            self._stats = stats
        stats[key] = int(stats.get(key, 0)) + int(delta)

    async def receive_packet(self, packet:FlownetPacket):
        self._bump("received_packets")
        # self._log(f"[CallManager] Received packet type={packet.type} uuid={packet.uuid}")
        if packet.type in {"response", "reply"}:
            if packet.type == "reply":
                warnings.warn(
                    "Packet type 'reply' is deprecated; use 'response'.",
                    DeprecationWarning,
                    stacklevel=2,
                )
                self._bump("legacy_reply_packets")
            self._log(f"[CallManager] Resolving future uuid={packet.uuid}")

            future = self.pending.get(packet.uuid, None)
            self._log(f"[CallManager] Future found={future is not None}")
            if future:
                remote_error = _decode_remote_error(packet.packet_content)
                if remote_error is not None:
                    self._bump("resolved_packets")
                    self.reject(packet.uuid, RuntimeError(remote_error))
                else:
                    self._bump("resolved_packets")
                    self.resolve(packet.uuid, packet.packet_content)
            else:
                self._bump("missing_future_packets")
        else:
            return await self.runtime_handler(packet)

    def register(self, uuid: str) -> asyncio.Future:
        self._log(f"[CallManager] Registering call uuid={uuid}")
        loop = loops.get_runtime_loop()
        fut = loop.create_future()
        now = time.time()
        deadline = now + self.default_timeout
        self.pending[uuid] = fut
        self.timestamps[uuid] = now
        self._deadlines[uuid] = deadline
        heapq.heappush(self._deadline_heap, (deadline, uuid))
        return fut
 
    def resolve(self, uuid: str, result):
        fut = self.pending.pop(uuid, None)
        self.timestamps.pop(uuid, None)
        self._deadlines.pop(uuid, None)
        if fut and not fut.done():
            fut.set_result(result)

    def reject(self, uuid: str, exc: Exception):
        fut = self.pending.pop(uuid, None)
        self.timestamps.pop(uuid, None)
        self._deadlines.pop(uuid, None)
        if fut and not fut.done():
            fut.set_exception(exc)

    def _cancel_pending(self, uuid: str) -> None:
        fut = self.pending.pop(uuid, None)
        self.timestamps.pop(uuid, None)
        self._deadlines.pop(uuid, None)
        if fut and not fut.done():
            fut.cancel()

    def _collect_expired(self, now: float) -> list[str]:
        expired: list[str] = []
        while self._deadline_heap and self._deadline_heap[0][0] <= now:
            deadline, uuid = heapq.heappop(self._deadline_heap)
            current_deadline = self._deadlines.get(uuid)
            if current_deadline is None:
                continue
            if current_deadline != deadline:
                continue
            expired.append(uuid)
        return expired

    def _next_timeout_sleep(self, now: float) -> float:
        while self._deadline_heap:
            deadline, uuid = self._deadline_heap[0]
            current_deadline = self._deadlines.get(uuid)
            if current_deadline is None or current_deadline != deadline:
                heapq.heappop(self._deadline_heap)
                continue
            wait_s = max(0.05, deadline - now)
            return min(wait_s, 5.0)
        return 5.0

    async def timeout_watcher(self):
        while True:
            now = time.time()
            for uuid in self._collect_expired(now):
                self._bump("timeout_rejections")
                self.reject(uuid, TimeoutError(f"[CallManager] Call {uuid} timed out"))
            await asyncio.sleep(self._next_timeout_sleep(now))
