from __future__ import annotations

import hashlib
import os
import socket
from dataclasses import dataclass
from typing import Callable, Iterable

from flownet.runtime.transport.packet_transport import EndpointRef, split_host_port


_LOCAL_HOST_ALIASES = frozenset({"localhost", "127.0.0.1", "::1", "0.0.0.0"})
_UDS_MODES = frozenset({"auto", "uds_prefer", "hetero_auto"})
_SHM_MODES = frozenset({"hetero_auto"})
_VALID_CAPS = frozenset({"tcp", "uds", "shm"})


@dataclass(frozen=True)
class TransportRoute:
    endpoint: EndpointRef
    backend_order: tuple[str, ...]


class TransportSelector:
    def __init__(
        self,
        *,
        local_host: str,
        local_port: int,
        transport_mode: str = "auto",
        local_ipc_dir: str = "/tmp/flownet",
        local_host_id: str | None = None,
        local_uds_path: str | None = None,
        local_shm_path: str | None = None,
        local_transport_caps: Iterable[str] | None = None,
        shm_threshold_bytes: int = 262144,
    ):
        self.local_host = local_host
        self.local_port = int(local_port)
        self.transport_mode = (transport_mode or "auto").strip().lower()
        self.local_ipc_dir = local_ipc_dir or "/tmp/flownet"
        self.local_host_id = local_host_id or self._build_default_host_id()
        self.local_transport_caps = self._normalize_transport_caps(
            local_transport_caps,
            transport_mode=self.transport_mode,
        )
        self.shm_threshold_bytes = max(1, int(shm_threshold_bytes))
        self.local_address = f"{self.local_host}:{self.local_port}"
        self.local_uds_path = local_uds_path or self.uds_path_for_local_port(self.local_port)
        self.local_shm_path = local_shm_path or self.shm_path_for_local_port(self.local_port)
        self._local_host_aliases = self._build_local_host_aliases(local_host)
        self._endpoint_lookup: Callable[[str], dict[str, object] | None] | None = None
        self._metrics = {
            "endpoint_lookup_hits": 0,
            "endpoint_lookup_misses": 0,
            "endpoint_lookup_errors": 0,
        }

    def set_endpoint_lookup(
        self,
        endpoint_lookup: Callable[[str], dict[str, object] | None] | None,
    ) -> None:
        self._endpoint_lookup = endpoint_lookup

    def resolve(self, address: str, payload_size_bytes: int | None = None) -> TransportRoute:
        host, port = split_host_port(address)
        local_host = self.is_local_host(host)
        endpoint_meta = self._resolve_endpoint_metadata(address)
        endpoint_host_id = endpoint_meta.get("host_id")
        if not isinstance(endpoint_host_id, str):
            endpoint_host_id = None
        endpoint_caps = self._normalize_transport_caps(
            endpoint_meta.get("transport_caps"),
            transport_mode=self.transport_mode,
            default_local=False,
            same_machine_hint=local_host,
        )
        endpoint_uds_path = endpoint_meta.get("uds_path")
        if not isinstance(endpoint_uds_path, str):
            endpoint_uds_path = self.uds_path_for_local_port(port) if local_host else self.uds_path_for_address(address)
        endpoint_shm_path = endpoint_meta.get("shm_path")
        if not isinstance(endpoint_shm_path, str):
            endpoint_shm_path = self.shm_path_for_local_port(port) if local_host else self.shm_path_for_address(address)
        same_machine = local_host or (
            isinstance(endpoint_host_id, str)
            and self.local_host_id is not None
            and endpoint_host_id == self.local_host_id
        )
        endpoint = EndpointRef(
            address=address,
            host=host,
            port=port,
            uds_path=endpoint_uds_path,
            host_id=endpoint_host_id,
            transport_caps=endpoint_caps,
            shm_path=endpoint_shm_path,
        )
        if self.transport_mode == "tcp":
            return TransportRoute(endpoint=endpoint, backend_order=("tcp",))
        uds_allowed = (
            same_machine
            and "uds" in self.local_transport_caps
            and "uds" in endpoint_caps
        )
        shm_allowed = (
            same_machine
            and "shm" in self.local_transport_caps
            and "shm" in endpoint_caps
            and payload_size_bytes is not None
            and payload_size_bytes >= self.shm_threshold_bytes
        )
        if (
            self.transport_mode in _SHM_MODES
            and shm_allowed
        ):
            backends: list[str] = ["shm"]
            if uds_allowed:
                backends.append("uds")
            backends.append("tcp")
            return TransportRoute(endpoint=endpoint, backend_order=tuple(backends))
        if self.transport_mode in _UDS_MODES and uds_allowed:
            return TransportRoute(endpoint=endpoint, backend_order=("uds", "tcp"))
        return TransportRoute(endpoint=endpoint, backend_order=("tcp",))

    def is_local_host(self, host: str) -> bool:
        return host.strip().lower() in self._local_host_aliases

    def uds_path_for_address(self, address: str) -> str:
        safe = address.replace("/", "_").replace(":", "_")
        digest = hashlib.sha1(address.encode("utf-8")).hexdigest()[:12]
        filename = f"flownet-{safe[:32]}-{digest}.sock"
        return os.path.join(self.local_ipc_dir, filename)

    def uds_path_for_local_port(self, port: int) -> str:
        # Canonical local IPC path should not depend on host spelling
        # (127.0.0.1 / localhost / 0.0.0.0 / hostname).
        return self.uds_path_for_address(f"local:{int(port)}")

    def shm_path_for_address(self, address: str) -> str:
        safe = address.replace("/", "_").replace(":", "_")
        digest = hashlib.sha1(address.encode("utf-8")).hexdigest()[:12]
        filename = f"flownet-shm-{safe[:32]}-{digest}.sock"
        return os.path.join(self.local_ipc_dir, filename)

    def shm_path_for_local_port(self, port: int) -> str:
        return self.shm_path_for_address(f"local:{int(port)}")

    def metrics_snapshot(self) -> dict[str, int]:
        return dict(self._metrics)

    def _resolve_endpoint_metadata(self, address: str) -> dict[str, object]:
        if self._endpoint_lookup is None:
            return {}
        try:
            payload = self._endpoint_lookup(address)
        except Exception:
            self._metrics["endpoint_lookup_errors"] += 1
            return {}
        if isinstance(payload, dict):
            self._metrics["endpoint_lookup_hits"] += 1
            return payload
        self._metrics["endpoint_lookup_misses"] += 1
        return {}

    @staticmethod
    def _build_default_host_id() -> str:
        try:
            return socket.gethostname().strip().lower()
        except Exception:
            return "unknown-host"

    @staticmethod
    def _normalize_transport_caps(
        transport_caps: Iterable[str] | None,
        *,
        transport_mode: str,
        default_local: bool = True,
        same_machine_hint: bool = False,
    ) -> tuple[str, ...]:
        normalized: list[str] = []
        if transport_caps is not None:
            source = (
                [item.strip() for item in transport_caps.split(",")]
                if isinstance(transport_caps, str)
                else transport_caps
            )
            for cap in source:
                value = str(cap).strip().lower()
                if value in _VALID_CAPS and value not in normalized:
                    normalized.append(value)
        if normalized:
            if "tcp" not in normalized:
                normalized.append("tcp")
            return tuple(normalized)
        if not default_local:
            if same_machine_hint and transport_mode != "tcp":
                if transport_mode == "hetero_auto":
                    return ("uds", "shm", "tcp")
                return ("uds", "tcp")
            return ("tcp",)
        if transport_mode == "tcp":
            return ("tcp",)
        if transport_mode == "hetero_auto":
            return ("uds", "shm", "tcp")
        return ("uds", "tcp")

    @staticmethod
    def _build_local_host_aliases(local_host: str) -> set[str]:
        aliases: set[str] = {local_host.strip().lower()}
        aliases.update(_LOCAL_HOST_ALIASES)
        try:
            infos = socket.getaddrinfo(socket.gethostname(), None)
            aliases.update(info[4][0].strip().lower() for info in infos if info[4])
        except Exception:
            pass
        try:
            infos = socket.getaddrinfo(local_host, None)
            aliases.update(info[4][0].strip().lower() for info in infos if info[4])
        except Exception:
            pass
        return aliases


__all__ = [
    "TransportRoute",
    "TransportSelector",
]
