from __future__ import annotations

import threading
import time
from typing import TYPE_CHECKING

import cloudpickle

from flownet.runtime.transport.transport_selector import TransportSelector

if TYPE_CHECKING:
    from flownet.core.request import FlownetPacket
    from flownet.runtime.transport.packet_transport import EndpointRef, PacketTransport


_LATENCY_BUCKETS_MS: tuple[tuple[float, str], ...] = (
    (0.1, "le_0_1ms"),
    (0.5, "le_0_5ms"),
    (1.0, "le_1ms"),
    (5.0, "le_5ms"),
    (10.0, "le_10ms"),
    (50.0, "le_50ms"),
)
_LATENCY_LAST_BUCKET = "gt_50ms"


class TransportHub:
    def __init__(
        self,
        *,
        selector: TransportSelector,
        tcp_transport: PacketTransport,
        uds_transport: PacketTransport | None,
        shm_transport: PacketTransport | None = None,
        backend_cooldown_seconds: float = 1.0,
        backend_cooldown_max_entries: int = 4096,
        backend_cooldown_cleanup_interval: int = 128,
    ):
        self.selector = selector
        self.tcp_transport = tcp_transport
        self.uds_transport = uds_transport
        self.shm_transport = shm_transport
        self.backend_cooldown_seconds = max(0.0, float(backend_cooldown_seconds))
        self.backend_cooldown_max_entries = max(1, int(backend_cooldown_max_entries))
        self.backend_cooldown_cleanup_interval = max(1, int(backend_cooldown_cleanup_interval))
        self._backend_cooldown_until: dict[tuple[str, str], float] = {}
        self._cooldown_op_seq = 0
        self._metrics_lock = threading.Lock()
        self._metrics: dict[str, object] = {
            "send_tcp_total": 0,
            "send_uds_total": 0,
            "send_shm_total": 0,
            "fallback_total": 0,
            "fallback_uds_to_tcp_total": 0,
            "fallback_shm_to_uds_total": 0,
            "fallback_shm_to_tcp_total": 0,
            "backend_cooldown_skips_total": 0,
            "backend_cooldown_evicted_total": 0,
            "backend_cooldown_expired_pruned_total": 0,
            "backend_cooldown_cleared_on_success_total": 0,
            "backend_cooldown_skip_counts": {name: 0 for name in ("tcp", "uds", "shm")},
            "backend_cooldown_cleared_on_success_counts": {name: 0 for name in ("tcp", "uds", "shm")},
            "backend_select_counts": {name: 0 for name in ("tcp", "uds", "shm")},
            "backend_error_counts": {name: 0 for name in ("tcp", "uds", "shm")},
            "backend_latency_ms": {
                name: {
                    "samples": 0,
                    "success_samples": 0,
                    "error_samples": 0,
                    "sum_ms": 0.0,
                    "max_ms": 0.0,
                    "bucket_counts": _empty_latency_buckets(),
                }
                for name in ("tcp", "uds", "shm")
            },
        }

    async def send_packet(self, address: str, packet: FlownetPacket) -> None:
        self._tick_and_maybe_prune_cooldowns()
        payload_size_bytes = None
        if self.selector.transport_mode == "hetero_auto":
            payload_size_bytes = self._estimate_payload_size(packet)
        route = self.selector.resolve(address, payload_size_bytes=payload_size_bytes)
        last_exc: Exception | None = None

        for index, backend in enumerate(route.backend_order):
            next_backend = route.backend_order[index + 1] if index + 1 < len(route.backend_order) else None
            if self._is_backend_in_cooldown(backend=backend, address=address):
                # Only skip a cooling-down backend when there is at least one
                # remaining backend that is currently usable. This avoids
                # dead-ending on single-backend routes.
                has_alternative = any(
                    self._is_backend_available(other_backend)
                    and not self._is_backend_in_cooldown(backend=other_backend, address=address)
                    for other_backend in route.backend_order[index + 1 :]
                )
                if has_alternative:
                    self._record_backend_cooldown_skip(backend=backend)
                    continue
            started_at = time.perf_counter()
            self._increment_backend_counter("backend_select_counts", backend)
            try:
                await self._send_with_backend(
                    backend=backend,
                    address=address,
                    packet=packet,
                    endpoint=route.endpoint,
                )
            except Exception as exc:
                last_exc = exc
                self._record_backend_latency_ms(
                    backend=backend,
                    latency_ms=(time.perf_counter() - started_at) * 1000.0,
                    succeeded=False,
                )
                self._increment_backend_counter("backend_error_counts", backend)
                self._activate_backend_cooldown(backend=backend, address=address)
                with self._metrics_lock:
                    self._metrics["fallback_total"] = int(self._metrics["fallback_total"]) + 1
                    if backend == "shm" and next_backend == "uds":
                        self._metrics["fallback_shm_to_uds_total"] = (
                            int(self._metrics["fallback_shm_to_uds_total"]) + 1
                        )
                    elif backend == "shm" and next_backend == "tcp":
                        self._metrics["fallback_shm_to_tcp_total"] = (
                            int(self._metrics["fallback_shm_to_tcp_total"]) + 1
                        )
                    elif backend == "uds" and next_backend == "tcp":
                        self._metrics["fallback_uds_to_tcp_total"] = (
                            int(self._metrics["fallback_uds_to_tcp_total"]) + 1
                        )
                continue

            self._record_backend_latency_ms(
                backend=backend,
                latency_ms=(time.perf_counter() - started_at) * 1000.0,
                succeeded=True,
            )
            self._clear_backend_cooldown_on_success(backend=backend, address=address)
            self._increment_send_counter(backend)
            return

        if last_exc is not None:
            raise last_exc
        raise RuntimeError(f"No transport backend available for {address}")

    async def _send_with_backend(
        self,
        *,
        backend: str,
        address: str,
        packet: FlownetPacket,
        endpoint: EndpointRef,
    ) -> None:
        if backend == "shm":
            if self.shm_transport is None:
                raise RuntimeError("SHM transport is unavailable.")
            await self.shm_transport.send(address, packet, endpoint=endpoint)
            return
        if backend == "uds":
            if self.uds_transport is None:
                raise RuntimeError("UDS transport is unavailable.")
            await self.uds_transport.send(address, packet, endpoint=endpoint)
            return
        if backend == "tcp":
            await self.tcp_transport.send(address, packet, endpoint=endpoint)
            return
        raise RuntimeError(f"Unknown backend {backend!r}")

    def _increment_send_counter(self, backend: str) -> None:
        key = f"send_{backend}_total"
        with self._metrics_lock:
            if key not in self._metrics:
                self._metrics[key] = 0
            self._metrics[key] = int(self._metrics[key]) + 1

    def _increment_backend_counter(self, metric: str, backend: str) -> None:
        with self._metrics_lock:
            counters = self._metrics.get(metric)
            if not isinstance(counters, dict):
                counters = {}
                self._metrics[metric] = counters
            counters[backend] = int(counters.get(backend, 0)) + 1

    def _record_backend_latency_ms(self, *, backend: str, latency_ms: float, succeeded: bool) -> None:
        with self._metrics_lock:
            latency = self._metrics.get("backend_latency_ms")
            if not isinstance(latency, dict):
                latency = {}
                self._metrics["backend_latency_ms"] = latency
            per_backend = latency.get(backend)
            if not isinstance(per_backend, dict):
                per_backend = {
                    "samples": 0,
                    "success_samples": 0,
                    "error_samples": 0,
                    "sum_ms": 0.0,
                    "max_ms": 0.0,
                    "bucket_counts": _empty_latency_buckets(),
                }
                latency[backend] = per_backend
            safe_latency_ms = max(0.0, float(latency_ms))
            per_backend["samples"] = int(per_backend.get("samples", 0)) + 1
            if succeeded:
                per_backend["success_samples"] = int(per_backend.get("success_samples", 0)) + 1
            else:
                per_backend["error_samples"] = int(per_backend.get("error_samples", 0)) + 1
            per_backend["sum_ms"] = float(per_backend.get("sum_ms", 0.0)) + safe_latency_ms
            per_backend["max_ms"] = max(float(per_backend.get("max_ms", 0.0)), safe_latency_ms)
            buckets = per_backend.get("bucket_counts")
            if not isinstance(buckets, dict):
                buckets = _empty_latency_buckets()
                per_backend["bucket_counts"] = buckets
            bucket_name = _latency_bucket_name(safe_latency_ms)
            buckets[bucket_name] = int(buckets.get(bucket_name, 0)) + 1

    def metrics_snapshot(self) -> dict[str, object]:
        with self._metrics_lock:
            send_tcp_total = int(self._metrics.get("send_tcp_total", 0))
            send_uds_total = int(self._metrics.get("send_uds_total", 0))
            send_shm_total = int(self._metrics.get("send_shm_total", 0))
            fallback_total = int(self._metrics.get("fallback_total", 0))
            fallback_uds_to_tcp_total = int(self._metrics.get("fallback_uds_to_tcp_total", 0))
            fallback_shm_to_uds_total = int(self._metrics.get("fallback_shm_to_uds_total", 0))
            fallback_shm_to_tcp_total = int(self._metrics.get("fallback_shm_to_tcp_total", 0))
            backend_cooldown_skips_total = int(self._metrics.get("backend_cooldown_skips_total", 0))
            backend_cooldown_evicted_total = int(self._metrics.get("backend_cooldown_evicted_total", 0))
            backend_cooldown_expired_pruned_total = int(
                self._metrics.get("backend_cooldown_expired_pruned_total", 0)
            )
            backend_cooldown_cleared_on_success_total = int(
                self._metrics.get("backend_cooldown_cleared_on_success_total", 0)
            )
            backend_cooldown_skip_counts = {
                str(name): int(value)
                for name, value in dict(self._metrics.get("backend_cooldown_skip_counts", {})).items()
            }
            backend_cooldown_cleared_on_success_counts = {
                str(name): int(value)
                for name, value in dict(
                    self._metrics.get("backend_cooldown_cleared_on_success_counts", {})
                ).items()
            }
            backend_cooldown_entries = len(self._backend_cooldown_until)
            backend_select_counts = {
                str(name): int(value)
                for name, value in dict(self._metrics.get("backend_select_counts", {})).items()
            }
            backend_error_counts = {
                str(name): int(value)
                for name, value in dict(self._metrics.get("backend_error_counts", {})).items()
            }
            raw_latency = self._metrics.get("backend_latency_ms")
            if not isinstance(raw_latency, dict):
                raw_latency = {}

            backend_latency_ms: dict[str, object] = {}
            for backend, raw_stats in raw_latency.items():
                if not isinstance(raw_stats, dict):
                    continue
                samples = int(raw_stats.get("samples", 0))
                total_sum_ms = float(raw_stats.get("sum_ms", 0.0))
                bucket_counts_raw = raw_stats.get("bucket_counts")
                bucket_counts = _empty_latency_buckets()
                if isinstance(bucket_counts_raw, dict):
                    for bucket, bucket_value in bucket_counts_raw.items():
                        bucket_counts[str(bucket)] = int(bucket_value)
                backend_latency_ms[str(backend)] = {
                    "samples": samples,
                    "success_samples": int(raw_stats.get("success_samples", 0)),
                    "error_samples": int(raw_stats.get("error_samples", 0)),
                    "avg": (total_sum_ms / float(samples)) if samples else 0.0,
                    "max": float(raw_stats.get("max_ms", 0.0)),
                    "bucket_counts": bucket_counts,
                }

        total_send = send_tcp_total + send_uds_total + send_shm_total
        total_select = sum(backend_select_counts.values())
        uds_select_total = int(backend_select_counts.get("uds", 0))
        shm_select_total = int(backend_select_counts.get("shm", 0))
        return {
            "send_tcp_total": send_tcp_total,
            "send_uds_total": send_uds_total,
            "send_shm_total": send_shm_total,
            "fallback_total": fallback_total,
            "fallback_uds_to_tcp_total": fallback_uds_to_tcp_total,
            "fallback_shm_to_uds_total": fallback_shm_to_uds_total,
            "fallback_shm_to_tcp_total": fallback_shm_to_tcp_total,
            "backend_cooldown_skips_total": backend_cooldown_skips_total,
            "backend_cooldown_entries": backend_cooldown_entries,
            "backend_cooldown_evicted_total": backend_cooldown_evicted_total,
            "backend_cooldown_expired_pruned_total": backend_cooldown_expired_pruned_total,
            "backend_cooldown_cleared_on_success_total": backend_cooldown_cleared_on_success_total,
            "backend_cooldown_skip_counts": backend_cooldown_skip_counts,
            "backend_cooldown_cleared_on_success_counts": backend_cooldown_cleared_on_success_counts,
            "backend_cooldown_config": {
                "seconds": self.backend_cooldown_seconds,
                "max_entries": self.backend_cooldown_max_entries,
                "cleanup_interval": self.backend_cooldown_cleanup_interval,
            },
            "backend_select_counts": backend_select_counts,
            "backend_error_counts": backend_error_counts,
            "fallback_rates": {
                "per_send": (float(fallback_total) / float(total_send)) if total_send > 0 else 0.0,
                "per_select": (float(fallback_total) / float(total_select)) if total_select > 0 else 0.0,
                "uds_to_tcp_per_uds_select": (
                    float(fallback_uds_to_tcp_total) / float(uds_select_total)
                    if uds_select_total > 0
                    else 0.0
                ),
                "shm_to_uds_per_shm_select": (
                    float(fallback_shm_to_uds_total) / float(shm_select_total)
                    if shm_select_total > 0
                    else 0.0
                ),
                "shm_to_tcp_per_shm_select": (
                    float(fallback_shm_to_tcp_total) / float(shm_select_total)
                    if shm_select_total > 0
                    else 0.0
                ),
            },
            "backend_latency_ms": backend_latency_ms,
        }

    @staticmethod
    def _estimate_payload_size(packet: FlownetPacket) -> int | None:
        try:
            return len(cloudpickle.dumps(packet))
        except Exception:
            return None

    def _is_backend_in_cooldown(self, *, backend: str, address: str) -> bool:
        if self.backend_cooldown_seconds <= 0:
            return False
        cooldown_key = (backend, address)
        with self._metrics_lock:
            until = self._backend_cooldown_until.get(cooldown_key)
        if until is None:
            return False
        now = time.monotonic()
        if now < until:
            return True
        with self._metrics_lock:
            current_until = self._backend_cooldown_until.get(cooldown_key)
            if current_until is not None and current_until <= now:
                self._backend_cooldown_until.pop(cooldown_key, None)
        return False

    def _activate_backend_cooldown(self, *, backend: str, address: str) -> None:
        if self.backend_cooldown_seconds <= 0:
            return
        cooldown_key = (backend, address)
        deadline = time.monotonic() + self.backend_cooldown_seconds
        with self._metrics_lock:
            self._backend_cooldown_until[cooldown_key] = deadline
            overflow = len(self._backend_cooldown_until) - self.backend_cooldown_max_entries
            if overflow > 0:
                sorted_items = sorted(
                    self._backend_cooldown_until.items(),
                    key=lambda item: item[1],
                )
                for key, _ in sorted_items[:overflow]:
                    self._backend_cooldown_until.pop(key, None)
                self._metrics["backend_cooldown_evicted_total"] = (
                    int(self._metrics.get("backend_cooldown_evicted_total", 0))
                    + overflow
                )

    def _record_backend_cooldown_skip(self, *, backend: str) -> None:
        with self._metrics_lock:
            self._metrics["backend_cooldown_skips_total"] = (
                int(self._metrics.get("backend_cooldown_skips_total", 0)) + 1
            )
            counters = self._metrics.get("backend_cooldown_skip_counts")
            if not isinstance(counters, dict):
                counters = {}
                self._metrics["backend_cooldown_skip_counts"] = counters
            counters[backend] = int(counters.get(backend, 0)) + 1

    def _clear_backend_cooldown_on_success(self, *, backend: str, address: str) -> None:
        cooldown_key = (backend, address)
        with self._metrics_lock:
            existed = self._backend_cooldown_until.pop(cooldown_key, None) is not None
            if not existed:
                return
            self._metrics["backend_cooldown_cleared_on_success_total"] = (
                int(self._metrics.get("backend_cooldown_cleared_on_success_total", 0)) + 1
            )
            counters = self._metrics.get("backend_cooldown_cleared_on_success_counts")
            if not isinstance(counters, dict):
                counters = {}
                self._metrics["backend_cooldown_cleared_on_success_counts"] = counters
            counters[backend] = int(counters.get(backend, 0)) + 1

    def _is_backend_available(self, backend: str) -> bool:
        if backend == "shm":
            return self.shm_transport is not None
        if backend == "uds":
            return self.uds_transport is not None
        if backend == "tcp":
            return self.tcp_transport is not None
        return False

    def _tick_and_maybe_prune_cooldowns(self) -> None:
        self._cooldown_op_seq += 1
        if (self._cooldown_op_seq % self.backend_cooldown_cleanup_interval) != 0:
            return
        self._prune_backend_cooldowns()

    def _prune_backend_cooldowns(self) -> None:
        now = time.monotonic()
        with self._metrics_lock:
            if not self._backend_cooldown_until:
                return
            expired_keys = [
                key
                for key, deadline in self._backend_cooldown_until.items()
                if deadline <= now
            ]
            for key in expired_keys:
                self._backend_cooldown_until.pop(key, None)
            if expired_keys:
                self._metrics["backend_cooldown_expired_pruned_total"] = (
                    int(self._metrics.get("backend_cooldown_expired_pruned_total", 0))
                    + len(expired_keys)
                )
            overflow = len(self._backend_cooldown_until) - self.backend_cooldown_max_entries
            if overflow > 0:
                # Evict nearest-expiry entries first to keep map bounded.
                sorted_items = sorted(
                    self._backend_cooldown_until.items(),
                    key=lambda item: item[1],
                )
                for key, _ in sorted_items[:overflow]:
                    self._backend_cooldown_until.pop(key, None)
                self._metrics["backend_cooldown_evicted_total"] = (
                    int(self._metrics.get("backend_cooldown_evicted_total", 0))
                    + overflow
                )


def _empty_latency_buckets() -> dict[str, int]:
    buckets = {bucket_name: 0 for _, bucket_name in _LATENCY_BUCKETS_MS}
    buckets[_LATENCY_LAST_BUCKET] = 0
    return buckets


def _latency_bucket_name(latency_ms: float) -> str:
    for upper_bound_ms, bucket_name in _LATENCY_BUCKETS_MS:
        if latency_ms <= upper_bound_ms:
            return bucket_name
    return _LATENCY_LAST_BUCKET


__all__ = [
    "TransportHub",
]
