"""
Distributed Chat Cluster Router with Sharded Ring Buffers for 1,000,000 CCU.
Supports Hybrid Redis Cluster Pub/Sub bridging and In-Memory high-throughput sharding.
"""

from __future__ import annotations

import asyncio
import inspect
import time
from typing import Any, Callable, Coroutine, Dict, List, Optional, Set, Tuple

from server.chat.chat_types import ChatMessageDTO

# Type alias for subscriber callback: def or async def callback(message: ChatMessageDTO) -> Any
SubscriberCallback = Callable[[ChatMessageDTO], Any]


class ShardedChannelRegistry:
    """
    Sharded subscriber registry that minimizes lock contention when
    managing up to 1,000,000 concurrent client subscriptions.
    """

    def __init__(self, num_shards: int = 64) -> None:
        self.num_shards = num_shards
        # List of shards: each shard is a dict of channel_key -> Dict of (client_id, callback)
        self.shards: List[Dict[str, Dict[int, SubscriberCallback]]] = [{} for _ in range(num_shards)]
        # Pre-partitioned cached callback lists for zero-reflection hot-path execution
        self.cached_sync: List[Dict[str, List[SubscriberCallback]]] = [{} for _ in range(num_shards)]
        self.cached_async: List[Dict[str, List[SubscriberCallback]]] = [{} for _ in range(num_shards)]

    def _get_shard_index(self, channel_key: str) -> int:
        return hash(channel_key) % self.num_shards

    def _rebuild_cache(self, shard_idx: int, channel_key: str) -> None:
        sync_list: List[SubscriberCallback] = []
        async_list: List[SubscriberCallback] = []
        for cb in self.shards[shard_idx][channel_key].values():
            if inspect.iscoroutinefunction(cb):
                async_list.append(cb)
            else:
                sync_list.append(cb)
        self.cached_sync[shard_idx][channel_key] = sync_list
        self.cached_async[shard_idx][channel_key] = async_list

    def subscribe(self, channel_key: str, client_id: int, callback: SubscriberCallback) -> None:
        shard_idx = self._get_shard_index(channel_key)
        shard = self.shards[shard_idx]
        if channel_key not in shard:
            shard[channel_key] = {}
        shard[channel_key][client_id] = callback
        self._rebuild_cache(shard_idx, channel_key)

    def unsubscribe(self, channel_key: str, client_id: int) -> None:
        shard_idx = self._get_shard_index(channel_key)
        shard = self.shards[shard_idx]
        if channel_key not in shard:
            return
        shard[channel_key].pop(client_id, None)
        if not shard[channel_key]:
            del shard[channel_key]
            self.cached_sync[shard_idx].pop(channel_key, None)
            self.cached_async[shard_idx].pop(channel_key, None)
        else:
            self._rebuild_cache(shard_idx, channel_key)

    def get_partitioned_callbacks(self, channel_key: str) -> Tuple[List[SubscriberCallback], List[SubscriberCallback]]:
        shard_idx = self._get_shard_index(channel_key)
        return (
            self.cached_sync[shard_idx].get(channel_key, []),
            self.cached_async[shard_idx].get(channel_key, [])
        )

    def get_callbacks(self, channel_key: str) -> List[SubscriberCallback]:
        shard_idx = self._get_shard_index(channel_key)
        sync_cbs = self.cached_sync[shard_idx].get(channel_key, [])
        async_cbs = self.cached_async[shard_idx].get(channel_key, [])
        return sync_cbs + async_cbs

    def get_subscribers(self, channel_key: str) -> List[Tuple[int, SubscriberCallback]]:
        shard_idx = self._get_shard_index(channel_key)
        shard = self.shards[shard_idx]
        subscribers = shard.get(channel_key)
        return list(subscribers.items()) if subscribers else []

    def total_subscribers_count(self) -> int:
        count = 0
        for shard in self.shards:
            for subs in shard.values():
                count += len(subs)
        return count


class RedisShardedPubSubBridge:
    """
    Redis 7 Sharded Pub/Sub bridge interface (SPUBLISH / SSUBSCRIBE).
    Routes messages using Redis cluster hash tags {...} to confine pub/sub
    traffic to the slot-owning node, preventing cluster-wide broadcast storms.
    """

    def __init__(
        self,
        publish_handler: Optional[Callable[[str, str], Coroutine[Any, Any, int]]] = None,
        subscribe_handler: Optional[Callable[[str, SubscriberCallback], Coroutine[Any, Any, None]]] = None,
    ) -> None:
        self._publish_handler = publish_handler
        self._subscribe_handler = subscribe_handler
        self.published_messages: List[Tuple[str, str]] = []

    @staticmethod
    def to_sharded_channel(channel_key: str) -> str:
        """
        Wraps channel key into Redis Cluster hash tag format.
        e.g. 'world' -> '{world}', 'zone:barrow' -> '{zone:barrow}', 'guild:101' -> '{guild:101}'.
        """
        if channel_key.startswith("{") and channel_key.endswith("}"):
            return channel_key
        return f"{{{channel_key}}}"

    async def spublish(self, channel_key: str, message: ChatMessageDTO | str) -> int:
        """Dispatches SPUBLISH command to Redis 7 shard slot."""
        sharded_ch = self.to_sharded_channel(channel_key)
        payload = message if isinstance(message, str) else message.raw_content
        self.published_messages.append((sharded_ch, payload))
        if self._publish_handler is not None:
            return await self._publish_handler(sharded_ch, payload)
        return 1

    async def ssubscribe(self, channel_key: str, callback: SubscriberCallback) -> None:
        """Registers SSUBSCRIBE on the Redis 7 shard channel."""
        sharded_ch = self.to_sharded_channel(channel_key)
        if self._subscribe_handler is not None:
            await self._subscribe_handler(sharded_ch, callback)

    async def sunsubscribe(self, channel_key: str) -> None:
        """Removes SSUBSCRIBE registration."""
        pass


class ChatClusterRouter:
    """
    Core distributed fan-out engine capable of horizontal cluster scaling.
    Combines local sharded non-blocking delivery with Redis Pub/Sub cluster bridge.
    """

    def __init__(
        self,
        num_shards: int = 64,
        batch_fanout_size: int = 1000,
        redis_bridge: Optional[RedisShardedPubSubBridge] = None,
    ) -> None:
        self.registry = ShardedChannelRegistry(num_shards=num_shards)
        self.batch_fanout_size = batch_fanout_size
        self.redis_bridge = redis_bridge
        self._dispatch_latencies_ms: List[float] = []
        self._total_messages_routed: int = 0

    def attach_redis_bridge(self, bridge: RedisShardedPubSubBridge) -> None:
        """Attaches an external Redis 7 Sharded Pub/Sub bridge."""
        self.redis_bridge = bridge

    def _safe_invoke(
        self,
        cb: SubscriberCallback,
        message: ChatMessageDTO
    ) -> Tuple[bool, Optional[Coroutine[Any, Any, Any]]]:
        """
        Invokes subscriber callback safely:
        If cb is a coroutine function (async def), returns coroutine for batch gathering.
        If cb is a regular synchronous callable, invokes it directly without coroutine overhead.
        Eliminates unneeded coroutines and TypeError exceptions.
        """
        try:
            if inspect.iscoroutinefunction(cb):
                coro = cb(message)
                return True, coro
            res = cb(message)
            if asyncio.iscoroutine(res):
                return True, res
            return True, None
        except Exception:
            return False, None

    async def broadcast_to_channel(
        self,
        channel_key: str,
        message: ChatMessageDTO
    ) -> int:
        """
        Dispatches a chat message to all subscribers of a channel key.
        Executes fan-out in parallel batches to achieve ultra-low p99 latency.
        """
        t_start = time.perf_counter()
        sync_cbs, async_cbs = self.registry.get_partitioned_callbacks(channel_key)
        delivered_count = 0

        # Fast sync dispatch path (zero allocations, zero inspect reflection)
        for cb in sync_cbs:
            try:
                res = cb(message)
                if res is not None and asyncio.iscoroutine(res):
                    async_cbs.append(cb)
                else:
                    delivered_count += 1
            except Exception:
                pass

        # Async coroutines path
        if async_cbs:
            pending_coroutines: List[Coroutine[Any, Any, Any]] = []
            for cb in async_cbs:
                try:
                    coro = cb(message)
                    try:
                        # Drive immediate coroutines inline (completes in <1µs if no await)
                        coro.send(None)
                    except StopIteration:
                        delivered_count += 1
                    except Exception:
                        pass
                    else:
                        # Suspended on actual async I/O
                        pending_coroutines.append(coro)
                except Exception:
                    pass

            if pending_coroutines:
                results = await asyncio.gather(*pending_coroutines, return_exceptions=True)
                delivered_count += sum(1 for r in results if not isinstance(r, Exception))

        # Redis 7 Sharded Pub/Sub bridge fanout hook
        if self.redis_bridge is not None:
            try:
                await self.redis_bridge.spublish(channel_key, message)
            except Exception:
                pass

        elapsed_ms = (time.perf_counter() - t_start) * 1000.0
        self._record_latency(elapsed_ms)
        self._total_messages_routed += 1
        return delivered_count

    def _record_latency(self, latency_ms: float) -> None:
        if len(self._dispatch_latencies_ms) >= 1000:
            self._dispatch_latencies_ms.pop(0)
        self._dispatch_latencies_ms.append(latency_ms)

    def get_latency_stats(self) -> Dict[str, float]:
        """Calculates p50, p95, p99 latency metrics for observability."""
        if not self._dispatch_latencies_ms:
            return {"p50": 0.0, "p95": 0.0, "p99": 0.0, "avg": 0.0}

        sorted_latencies = sorted(self._dispatch_latencies_ms)
        n = len(sorted_latencies)
        p50 = sorted_latencies[int(n * 0.50)]
        p95 = sorted_latencies[min(int(n * 0.95), n - 1)]
        p99 = sorted_latencies[min(int(n * 0.99), n - 1)]
        avg = sum(sorted_latencies) / n

        return {
            "p50": round(p50, 3),
            "p95": round(p95, 3),
            "p99": round(p99, 3),
            "avg": round(avg, 3),
            "total_routed": self._total_messages_routed
        }