"""Cluster client for topology-aware routing without cluster membership.

``ClusterClient`` connects to a Casty cluster as an external observer,
receives topology updates (membership + shard allocations), and routes
``ShardEnvelope`` messages directly to the owning node — zero hops,
no proxy overhead, no cluster participation.
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, cast
from uuid import uuid4

from casty.actor import Behavior, Behaviors
from casty.address import ActorAddress
from casty.cluster_state import NodeAddress, ServiceEntry
from casty.receptionist import Listing, ServiceInstance, ServiceKey
from casty.ref import ActorRef
from casty.remote_transport import (
    GetPort,
    InboundMessageHandler,
    RemoteTransport,
    TcpTransportConfig,
    TcpTransportMsg,
    tcp_transport,
)
from casty.scheduler import ScheduleOnce, scheduler
from casty.serialization import PickleSerializer
from casty.shard_coordinator_actor import GetShardLocation, ShardLocation
from casty.sharding import ShardEnvelope, entity_shard
from casty.system import ActorSystem
from casty.topology import SubscribeTopology, TopologySnapshot
from casty.transport import LocalTransport, MessageTransport

if TYPE_CHECKING:
    from casty.context import ActorContext

SUBSCRIPTION_TIMEOUT = 10.0


@dataclass(frozen=True)
class NodeFailed:
    """TCP send failure detected — evict allocations for this node."""

    host: str
    port: int


@dataclass(frozen=True)
class RetryShardQuery:
    """Delayed retry for a shard location query after a node failure."""

    shard_id: int


@dataclass(frozen=True)
class SubscriptionTimeout:
    """No topology update received within the liveness window."""


type ClientProxyMsg = (
    ShardEnvelope[Any]
    | ShardLocation
    | TopologySnapshot
    | NodeFailed
    | RetryShardQuery
)

type TopologySubscriberMsg = TopologySnapshot | RegisterProxy | SubscriptionTimeout


def client_proxy_behavior(
    *,
    shard_name: str,
    num_shards: int,
    remote_transport: RemoteTransport,
    system_name: str,
    logger: logging.Logger | None = None,
) -> Behavior[ClientProxyMsg]:
    """Shard proxy that always routes remotely.

    Uses allocation table from ``TopologySnapshot`` for known shards.
    On cache miss, queries the coordinator on the leader node via
    ``GetShardLocation`` (same protocol as the cluster's shard proxy).
    """

    log = logger or logging.getLogger(f"casty.client_proxy.{shard_name}")

    def _route(envelope: ShardEnvelope[Any], node: NodeAddress) -> None:
        remote_addr = ActorAddress(
            system=system_name,
            path=f"/_region-{shard_name}",
            host=node.host,
            port=node.port,
        )
        remote_transport.make_ref(remote_addr).tell(envelope)

    def _query_coordinator(
        shard_id: int,
        leader: NodeAddress,
        reply_to: ActorRef[ClientProxyMsg],
    ) -> None:
        coord_addr = ActorAddress(
            system=system_name,
            path=f"/_coord-{shard_name}",
            host=leader.host,
            port=leader.port,
        )
        coord_ref: ActorRef[GetShardLocation] = remote_transport.make_ref(coord_addr)
        coord_ref.tell(GetShardLocation(shard_id=shard_id, reply_to=reply_to))
        log.debug("Shard %d: requesting location from coordinator", shard_id)

    async def setup(ctx: ActorContext[ClientProxyMsg]) -> Behavior[ClientProxyMsg]:
        scheduler_ref = ctx.spawn(scheduler(), "_proxy_scheduler")

        def active(
            allocations: dict[int, NodeAddress],
            leader: NodeAddress | None,
            buffer: dict[int, tuple[ShardEnvelope[Any], ...]],
            failed_nodes: frozenset[NodeAddress],
        ) -> Behavior[ClientProxyMsg]:
            async def receive(
                ctx: ActorContext[ClientProxyMsg],
                msg: ClientProxyMsg,
            ) -> Behavior[ClientProxyMsg]:
                match msg:
                    case ShardEnvelope() as envelope:
                        shard_id = entity_shard(envelope.entity_id, num_shards)

                        if shard_id in allocations:
                            _route(envelope, allocations[shard_id])
                            return Behaviors.same()

                        existing = buffer.get(shard_id, ())
                        new_buf = {**buffer, shard_id: (*existing, envelope)}
                        if not existing and leader is not None:
                            _query_coordinator(shard_id, leader, ctx.self)
                        return active(allocations, leader, new_buf, failed_nodes)

                    case ShardLocation(shard_id=sid, node=node):
                        if node in failed_nodes:
                            log.debug(
                                "Shard %d -> %s:%d REJECTED (failed node)",
                                sid,
                                node.host,
                                node.port,
                            )
                            scheduler_ref.tell(
                                ScheduleOnce(
                                    key=f"retry-shard-{sid}",
                                    target=ctx.self,
                                    message=RetryShardQuery(shard_id=sid),
                                    delay=0.5,
                                )
                            )
                            return Behaviors.same()

                        new_allocs = {**allocations, sid: node}
                        buffered = buffer.get(sid, ())
                        log.debug(
                            "Shard %d -> %s:%d (flushed %d)",
                            sid,
                            node.host,
                            node.port,
                            len(buffered),
                        )
                        new_buf = {k: v for k, v in buffer.items() if k != sid}
                        for envelope in buffered:
                            _route(envelope, node)
                        return active(new_allocs, leader, new_buf, failed_nodes)

                    case TopologySnapshot() as snapshot:
                        allocs = snapshot.shard_allocations.get(shard_name, {})
                        merged = {
                            **allocations,
                            **{sid: alloc.primary for sid, alloc in allocs.items()},
                        }
                        unreachable = snapshot.unreachable
                        new_allocations = {
                            sid: node
                            for sid, node in merged.items()
                            if node not in unreachable
                        }
                        evicted = len(merged) - len(new_allocations)
                        if evicted:
                            log.info(
                                "Evicted %d shards routed to unreachable nodes",
                                evicted,
                            )
                        new_leader = snapshot.leader
                        still_buffered: dict[int, tuple[ShardEnvelope[Any], ...]] = {}
                        for sid, envelopes in buffer.items():
                            if sid in new_allocations:
                                for envelope in envelopes:
                                    _route(envelope, new_allocations[sid])
                            else:
                                still_buffered[sid] = envelopes
                        if new_leader is not None:
                            for sid in still_buffered:
                                if sid not in allocations:
                                    _query_coordinator(sid, new_leader, ctx.self)
                        healthy_addrs = (
                            {m.address for m in snapshot.members} - snapshot.unreachable
                        )
                        new_failed = failed_nodes - healthy_addrs
                        log.debug(
                            "Topology: %d shards, leader=%s (epoch=%d)",
                            len(new_allocations),
                            f"{new_leader.host}:{new_leader.port}"
                            if new_leader
                            else "none",
                            snapshot.allocation_epoch,
                        )
                        return active(
                            new_allocations, new_leader, still_buffered, new_failed
                        )

                    case NodeFailed(host=failed_host, port=failed_port):
                        failed = NodeAddress(host=failed_host, port=failed_port)
                        new_allocs = {
                            sid: node
                            for sid, node in allocations.items()
                            if node != failed
                        }
                        evicted_sids = [
                            sid for sid in allocations if sid not in new_allocs
                        ]
                        new_failed = failed_nodes | {failed}
                        if evicted_sids:
                            log.info(
                                "TCP failure %s:%d — evicted %d shards",
                                failed_host,
                                failed_port,
                                len(evicted_sids),
                            )
                        for sid in evicted_sids:
                            scheduler_ref.tell(
                                ScheduleOnce(
                                    key=f"retry-shard-{sid}",
                                    target=ctx.self,
                                    message=RetryShardQuery(shard_id=sid),
                                    delay=0.5,
                                )
                            )
                        new_leader = None if leader == failed else leader
                        return active(new_allocs, new_leader, buffer, new_failed)

                    case RetryShardQuery(shard_id=sid):
                        if sid in buffer and leader is not None:
                            _query_coordinator(sid, leader, ctx.self)
                        return Behaviors.same()

                    case _:
                        return Behaviors.same()

            return Behaviors.receive(receive)

        return active({}, None, {}, frozenset())

    return Behaviors.setup(setup)


@dataclass(frozen=True)
class RegisterProxy:
    """Register an additional proxy to receive topology updates."""

    proxy_ref: ActorRef[ClientProxyMsg]


def topology_subscriber(
    *,
    contact_points: list[tuple[str, int]],
    remote_transport: RemoteTransport,
    system_name: str,
    proxies: list[ActorRef[ClientProxyMsg]],
    logger: logging.Logger | None = None,
    on_snapshot: Callable[[TopologySnapshot], None] | None = None,
) -> Behavior[TopologySubscriberMsg]:
    """Actor that subscribes to topology updates from a cluster contact point.

    On startup, sends ``SubscribeTopology`` to the first reachable contact
    point's ``/_topology`` actor. Forwards received ``TopologySnapshot``
    to all registered proxies. Reconnects to the next contact point if
    no topology update is received within ``SUBSCRIPTION_TIMEOUT``.
    """

    if not contact_points:
        msg = "No contact points provided"
        raise ValueError(msg)

    log = logger or logging.getLogger("casty.topology_subscriber")

    static_contacts = tuple(contact_points)

    def _build_contacts(
        snapshot: TopologySnapshot | None,
    ) -> tuple[tuple[str, int], ...]:
        if snapshot is None:
            return static_contacts
        member_addrs = tuple(
            (m.address.host, m.address.port)
            for m in snapshot.members
            if (m.address.host, m.address.port) not in static_contacts
        )
        return (*static_contacts, *member_addrs)

    def _subscribe_to(
        idx: int,
        known: tuple[tuple[str, int], ...],
        ctx: ActorContext[TopologySubscriberMsg],
        sched: ActorRef[Any],
    ) -> None:
        host, port = known[idx % len(known)]
        log.info("Subscribing to topology from %s:%d", host, port)
        addr = ActorAddress(
            system=system_name,
            path="/_cluster/_topology",
            host=host,
            port=port,
        )
        ref: ActorRef[SubscribeTopology] = remote_transport.make_ref(addr)
        ref.tell(SubscribeTopology(reply_to=ctx.self))
        sched.tell(
            ScheduleOnce(
                key="liveness",
                target=ctx.self,
                message=SubscriptionTimeout(),
                delay=SUBSCRIPTION_TIMEOUT,
            )
        )

    def connected(
        registered_proxies: tuple[ActorRef[ClientProxyMsg], ...],
        last_snapshot: TopologySnapshot | None,
        contact_index: int,
        scheduler_ref: ActorRef[Any],
    ) -> Behavior[TopologySubscriberMsg]:
        async def receive(
            ctx: ActorContext[TopologySubscriberMsg],
            msg: TopologySubscriberMsg,
        ) -> Behavior[TopologySubscriberMsg]:
            match msg:
                case TopologySnapshot() as snapshot:
                    log.info(
                        "Topology: %d members, %d shard types, epoch=%d",
                        len(snapshot.members),
                        len(snapshot.shard_allocations),
                        snapshot.allocation_epoch,
                    )
                    if on_snapshot is not None:
                        on_snapshot(snapshot)
                    for proxy in registered_proxies:
                        proxy.tell(snapshot)
                    scheduler_ref.tell(
                        ScheduleOnce(
                            key="liveness",
                            target=ctx.self,
                            message=SubscriptionTimeout(),
                            delay=SUBSCRIPTION_TIMEOUT,
                        )
                    )
                    return connected(
                        registered_proxies, snapshot, contact_index, scheduler_ref
                    )

                case RegisterProxy(proxy_ref=proxy_ref):
                    if last_snapshot is not None:
                        proxy_ref.tell(last_snapshot)
                    return connected(
                        (*registered_proxies, proxy_ref),
                        last_snapshot,
                        contact_index,
                        scheduler_ref,
                    )

                case SubscriptionTimeout():
                    known = _build_contacts(last_snapshot)
                    next_idx = (contact_index + 1) % len(known)
                    host, port = known[next_idx]
                    log.warning(
                        "No topology in %.0fs, reconnecting to %s:%d (index=%d/%d)",
                        SUBSCRIPTION_TIMEOUT,
                        host,
                        port,
                        next_idx,
                        len(known),
                    )
                    _subscribe_to(next_idx, known, ctx, scheduler_ref)
                    return connected(
                        registered_proxies, last_snapshot, next_idx, scheduler_ref
                    )

                case _:
                    return Behaviors.same()

        return Behaviors.receive(receive)

    async def setup(
        ctx: ActorContext[TopologySubscriberMsg],
    ) -> Behavior[TopologySubscriberMsg]:
        sched = ctx.spawn(scheduler(), "_sub_scheduler")
        _subscribe_to(0, static_contacts, ctx, sched)
        return connected(tuple(proxies), None, 0, sched)

    return Behaviors.setup(setup)


class _RemoteActorSystem(ActorSystem):
    """ActorSystem that provides remote addressing for spawned actors."""

    def __init__(
        self,
        name: str,
        remote_transport: RemoteTransport | None,
        host: str,
        port: int,
    ) -> None:
        super().__init__(name=name)
        self._remote = remote_transport
        self._host = host
        self._port = port

    @property
    def local_transport(self) -> LocalTransport:
        return self._local_transport

    def _get_ref_transport(self) -> MessageTransport | None:
        return self._remote

    def _get_ref_host(self) -> str | None:
        return self._host

    def _get_ref_port(self) -> int | None:
        return self._port

    def set_remote(self, transport: RemoteTransport) -> None:
        self._remote = transport

    async def shutdown(self) -> None:
        tcp_cell = self._root_cells.pop("_tcp_transport", None)
        await super().shutdown()
        if tcp_cell is not None:
            await tcp_cell.stop()

    def set_port(self, port: int) -> None:
        self._port = port

    def set_host(self, host: str) -> None:
        self._host = host


class ClusterClient:
    """External client that routes messages to a Casty cluster.

    Connects to cluster contact points via TCP, subscribes to topology
    updates, and routes ``ShardEnvelope`` messages directly to the node
    owning each shard — zero hops, no cluster membership.

    Parameters
    ----------
    contact_points
        List of ``(host, port)`` for cluster nodes to connect to.
    system_name
        Logical name of the cluster's actor system (must match).
    client_host
        Advertised hostname for this client (for receiving responses).
    client_port
        Advertised port for this client (``0`` for auto-assignment).

    Examples
    --------
    >>> async with ClusterClient(
    ...     contact_points=[("10.0.1.10", 25520)],
    ...     system_name="my-cluster",
    ... ) as client:
    ...     counter = client.entity_ref("counter", num_shards=100)
    ...     counter.tell(ShardEnvelope("user-42", Increment(1)))
    ...     count = await client.ask(
    ...         counter,
    ...         lambda r: ShardEnvelope("user-42", GetCount(reply_to=r)),
    ...     )
    """

    def __init__(
        self,
        *,
        contact_points: list[tuple[str, int]],
        system_name: str,
        client_host: str = "127.0.0.1",
        client_port: int = 0,
        advertised_host: str | None = None,
        advertised_port: int | None = None,
        address_map: dict[tuple[str, int], tuple[str, int]] | None = None,
    ) -> None:
        self._contact_points = contact_points
        self._system_name = system_name
        self._client_host = client_host
        self._client_port = client_port
        self._advertised_host = advertised_host
        self._advertised_port = advertised_port
        self._address_map = address_map
        self._proxies: dict[str, ActorRef[ClientProxyMsg]] = {}
        self._subscriber_ref: ActorRef[TopologySubscriberMsg] | None = None
        self._last_snapshot: TopologySnapshot | None = None
        self._logger = logging.getLogger(f"casty.client.{system_name}")
        self._serializer = PickleSerializer()
        self._system: _RemoteActorSystem | None = None
        self._remote_transport: RemoteTransport | None = None

    async def __aenter__(self) -> ClusterClient:
        client_name = f"{self._system_name}-client"

        # 1. Start the actor system first
        self._system = _RemoteActorSystem(
            name=client_name,
            remote_transport=None,
            host=self._client_host,
            port=self._client_port,
        )
        await self._system.__aenter__()

        # 2. Spawn TCP transport actor (client_only — no server)
        effective_host = self._advertised_host or self._client_host
        effective_port = self._advertised_port or self._client_port
        tcp_config = TcpTransportConfig(
            host=self._client_host,
            port=self._client_port,
            self_address=(effective_host, effective_port),
            client_only=True,
            address_map=self._address_map or {},
        )
        inbound = InboundMessageHandler(
            local=self._system.local_transport,
            serializer=self._serializer,
            system_name=client_name,
        )
        tcp_ref: ActorRef[TcpTransportMsg] = self._system.spawn(
            tcp_transport(
                tcp_config, inbound,
                logger=logging.getLogger(f"casty.client.tcp.{self._system_name}"),
            ),
            "_tcp_transport",
        )

        # 3. Ask for actual port and update system address
        actual_port: int = await self._system.ask(
            tcp_ref, lambda r: GetPort(reply_to=r), timeout=5.0,
        )
        if actual_port != self._client_port:
            self._client_port = actual_port
            self._system.set_port(actual_port)
        if self._advertised_host is not None:
            self._system.set_host(self._advertised_host)

        # 4. Create RemoteTransport with the actor ref
        self._remote_transport = RemoteTransport(
            local=self._system.local_transport,
            tcp=tcp_ref,
            serializer=self._serializer,
            local_host=self._client_host,
            local_port=self._client_port,
            system_name=client_name,
            on_send_failure=self._on_send_failure,
            advertised_host=self._advertised_host,
            advertised_port=self._advertised_port,
            task_runner=self._system._ensure_task_runner(),  # pyright: ignore[reportPrivateUsage]
        )
        self._system.set_remote(self._remote_transport)

        self._subscriber_ref = self._system.spawn(
            topology_subscriber(
                contact_points=self._contact_points,
                remote_transport=self._remote_transport,
                system_name=self._system_name,
                proxies=[],
                logger=self._logger,
                on_snapshot=self._on_topology_update,
            ),
            "_topology_sub",
        )

        self._logger.info(
            "ClusterClient started on %s:%d → cluster %s",
            self._client_host,
            self._client_port,
            self._system_name,
        )
        return self

    async def __aexit__(self, *exc: object) -> None:
        if self._system is not None:
            await self._system.shutdown()

    def _on_send_failure(self, host: str, port: int) -> None:
        msg = NodeFailed(host=host, port=port)
        for proxy_ref in self._proxies.values():
            proxy_ref.tell(msg)

    def _on_topology_update(self, snapshot: TopologySnapshot) -> None:
        self._last_snapshot = snapshot

    def _require_started(self) -> tuple[_RemoteActorSystem, RemoteTransport]:
        if self._system is None or self._remote_transport is None:
            raise RuntimeError("ClusterClient is not started — use 'async with'")
        return self._system, self._remote_transport

    def lookup[M](self, key: ServiceKey[M]) -> Listing[M]:
        """Look up service instances registered in the cluster.

        Reads from the locally cached topology snapshot — no network
        round-trip required.  Returns an empty ``Listing`` if no
        topology has been received yet.

        Parameters
        ----------
        key
            Typed service key to search for.

        Returns
        -------
        Listing[M]
            Current set of instances matching the key.
        """
        _, remote_transport = self._require_started()
        entries: frozenset[ServiceEntry] = (
            self._last_snapshot.registry
            if self._last_snapshot is not None
            else frozenset()
        )
        instances = frozenset(
            ServiceInstance(
                ref=remote_transport.make_ref(
                    ActorAddress(
                        system=self._system_name,
                        path=entry.path,
                        host=entry.node.host,
                        port=entry.node.port,
                    )
                ),
                node=entry.node,
            )
            for entry in entries
            if entry.key == key.name
        )
        return Listing(key=key, instances=instances)

    def entity_ref(
        self, shard_type: str, *, num_shards: int
    ) -> ActorRef[ShardEnvelope[Any]]:
        """Get a ref that routes ``ShardEnvelope`` messages to the cluster.

        Creates a local proxy actor for the given shard type on first call.
        Subsequent calls with the same ``shard_type`` return the cached proxy.

        Parameters
        ----------
        shard_type
            Name of the sharded entity type (must match the cluster's name).
        num_shards
            Number of shards (must match the cluster's configuration).

        Returns
        -------
        ActorRef[ShardEnvelope[Any]]
            A ref that accepts ``ShardEnvelope`` messages and routes them
            to the correct cluster node.

        Examples
        --------
        >>> counter = client.entity_ref("counter", num_shards=100)
        >>> counter.tell(ShardEnvelope("entity-1", Increment(1)))
        """
        if shard_type in self._proxies:
            return cast(ActorRef[ShardEnvelope[Any]], self._proxies[shard_type])

        system, remote_transport = self._require_started()

        proxy_ref = system.spawn(
            client_proxy_behavior(
                shard_name=shard_type,
                num_shards=num_shards,
                remote_transport=remote_transport,
                system_name=self._system_name,
                logger=logging.getLogger(f"casty.client_proxy.{shard_type}"),
            ),
            f"_proxy-{shard_type}",
        )
        self._proxies[shard_type] = proxy_ref

        if self._subscriber_ref is not None:
            self._subscriber_ref.tell(RegisterProxy(proxy_ref=proxy_ref))

        return cast(ActorRef[ShardEnvelope[Any]], proxy_ref)

    async def ask[M, R](
        self,
        ref: ActorRef[M],
        msg_factory: Callable[[ActorRef[R]], M],
        *,
        timeout: float = 5.0,
    ) -> R:
        """Send a message and wait for a reply.

        Creates a temporary remotely-addressable ref so the cluster can
        respond directly to this client via TCP.

        Parameters
        ----------
        ref
            Target actor (typically from ``entity_ref``).
        msg_factory
            Factory that receives a reply-to ref and returns the message.
        timeout
            Maximum seconds to wait.

        Returns
        -------
        R
            The reply from the cluster actor.

        Raises
        ------
        asyncio.TimeoutError
            If no reply is received within ``timeout``.

        Examples
        --------
        >>> count = await client.ask(
        ...     counter_ref,
        ...     lambda r: ShardEnvelope("user-42", GetCount(reply_to=r)),
        ... )
        """
        system, remote_transport = self._require_started()

        future: asyncio.Future[R] = asyncio.get_running_loop().create_future()
        ask_id = uuid4().hex
        temp_path = f"/_ask/{ask_id}"

        def on_reply(msg: Any) -> None:
            if not future.done():
                future.set_result(msg)

        system.local_transport.register(temp_path, on_reply)
        try:
            temp_ref: ActorRef[R] = ActorRef(
                address=ActorAddress(
                    system=f"{self._system_name}-client",
                    path=temp_path,
                    host=self._advertised_host or self._client_host,
                    port=self._advertised_port or self._client_port,
                ),
                _transport=remote_transport,
            )
            message = msg_factory(temp_ref)
            ref.tell(message)
            return await asyncio.wait_for(future, timeout=timeout)
        finally:
            system.local_transport.unregister(temp_path)
