# clients/conduit_manager.py (Updated)
import asyncio
from typing import Any, Dict, List, Optional, Tuple

from common.logging import get_logger, span
from common.platform.twitch.model import (
    EventSubConduit,
    EventSubConduitShard,
    EventSubTransport,
    EventSubSubscription,
    EventSubSubscriptionRequest,
)
from common.platform.twitch.webhook import WebhookManager
from common.platform.twitch.websocket import WebSocketManager
from .api import TwitchApi

logger = get_logger(__name__)


async def _wait_for_disconnect(shard_id, event):
    await event.wait()
    return shard_id


class ConduitManager:
    """
    Manages EventSub Conduits and their
    associated transports (WebSockets and Webhooks).

    Args:
        twitch_api (TwitchApi): Instance to manage EventSub subscriptions and conduits.
        websocket_manager (WebSocketManager): Manager for WebSocket shards.
        webhook_manager (WebhookManager): Manager for Webhook shards.
        max_shards (int): Maximum number of WebSocket shards.
        max_conduits (int): Maximum number of conduits.

    Methods
    -------
    create_conduit(shard_count: int) -> Optional[EventSubConduit]:
        Create a new EventSub Conduit.
    monitor_websocket_shards(conduit_id: str):
        Monitor WebSocket shards and restart them if disconnected.
    assign_transports_to_conduit_shards(conduit_id: str, transport_configs: List[EventSubTransport]) -> List[EventSubConduitShard]:
        Assign transports (Webhooks or WebSockets) to shards within a Conduit.
    list_conduits() -> List[EventSubConduit]:
        List all EventSub Conduits.
    delete_conduit(conduit_id: str) -> bool:
        Delete an EventSub Conduit.
    shutdown():
        Shutdown all managed transports (WebSockets and Webhooks).
    """

    def __init__(
        self,
        twitch_api: TwitchApi,
        websocket_manager: Optional[WebSocketManager] = None,
        webhook_manager: Optional[WebhookManager] = None,
        max_shards: int = 10,
        max_conduits: int = 1,
    ):
        """
        Initialize the ConduitManager.

        Args:
            twitch_api (TwitchApi): Instance to manage EventSub subscriptions and conduits.
            websocket_manager (WebSocketManager): Manager for WebSocket shards.
            webhook_manager (WebhookManager): Manager for Webhook shards.
        """
        self.twitch_api = twitch_api
        self.conduits: Dict[str, EventSubConduit] = {}
        self.active_conduit_id: Optional[str] = None
        if max_shards < 1:
            raise ValueError("max_shards must be greater than 0")
        if max_conduits < 1:
            raise ValueError("max_conduits must be greater than 0")
        if max_shards > 20000:
            raise ValueError("max_shards must be less than or equal to 20000")
        if max_conduits > 5:
            raise ValueError("max_conduits must be less than 5")
        self.max_shards = max_shards
        self.max_conduits = max_conduits
        self.websocket_manager = websocket_manager
        self.webhook_manager = webhook_manager

    async def initialize(self):
        conduits = await self.list_conduits()
        for conduit in conduits:
            self.conduits[conduit.id] = conduit
        if conduits and not self.active_conduit_id:
            self.active_conduit_id = conduits[0].id

    async def __aenter__(self):
        await self.initialize()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        return self.shutdown()

    async def create_conduit(self, shard_count: int) -> Optional[EventSubConduit]:
        """
        Create a new EventSub Conduit.

        Args:
            shard_count (int): Number of shards to create.

        Returns:
            Optional[EventSubConduit]: The created Conduit.
        """
        attributes = {
            "shard_count": shard_count,
            "max_shards": self.max_shards,
            "current_conduits": len(self.conduits),
            "max_conduits": self.max_conduits,
        }

        with span(logger, "create_conduit", attributes):
            logger.info(f"Creating new EventSub Conduit with {shard_count} shards")

            # Validate parameters
            if shard_count < 1:
                logger.error("Shard count must be greater than 0")
                return None
            if shard_count > self.max_shards:
                logger.error(f"Shard count exceeds maximum of {self.max_shards}")
                return None
            if len(self.conduits) >= self.max_conduits:
                logger.error(
                    f"Maximum number of conduits ({self.max_conduits}) reached"
                )
                return None

            # Create the conduit
            logger.debug(
                f"Calling Twitch API to create conduit with {shard_count} shards"
            )
            conduit = await self.twitch_api.create_eventsub_conduit(shard_count)

            if conduit:
                self.conduits[conduit.id] = conduit
                self.active_conduit_id = conduit.id
                logger.info(
                    f"Successfully created Conduit with ID: {conduit.id} and Shard Count: {conduit.shard_count}"
                )
                return conduit
            else:
                logger.error("Failed to create conduit through Twitch API")
                return None

    async def ensure_conduit(
        self, preferred_conduit_id: Optional[str], shard_count: int
    ) -> Optional[EventSubConduit]:
        """Ensure that a usable conduit exists, creating one if necessary.

        Args:
            preferred_conduit_id: Conduit ID from configuration (may be missing).
            shard_count: Desired shard count when creating a new conduit.

        Returns:
            EventSubConduit instance if available, otherwise None.
        """

        preferred = (preferred_conduit_id or "").strip()

        conduits = await self.list_conduits()

        if preferred:
            match = next((c for c in conduits if c.id == preferred), None)
            if match:
                self.active_conduit_id = match.id
                return match

            if conduits:
                fallback = conduits[0]
                self.active_conduit_id = fallback.id
                logger.warning(
                    "Preferred conduit missing; using available conduit",
                    extra={
                        "preferred": preferred,
                        "fallback": fallback.id,
                        "available": [c.id for c in conduits],
                    },
                )
                return fallback

        if conduits:
            chosen = conduits[0]
            self.active_conduit_id = chosen.id
            logger.info(
                "No preferred conduit configured; using first available",
                extra={"selected": chosen.id},
            )
            return chosen

        logger.info("No conduits exist; creating new conduit")
        created = await self.create_conduit(shard_count)
        if not created:
            logger.error("Failed to create a conduit during ensure_conduit")
        return created

    async def ensure_shard_monitoring_subscription(
        self, conduit_id: str
    ) -> Optional[EventSubSubscription]:
        """Ensure a conduit.shard.disabled subscription exists for the conduit."""

        attributes = {"conduit_id": conduit_id}
        with span(logger, "ensure_shard_monitoring_subscription", attributes):
            try:
                subscriptions = await self.twitch_api.list_eventsub_subscriptions()
            except Exception as exc:  # pragma: no cover - defensive logging
                logger.error(
                    "Failed to list EventSub subscriptions while ensuring shard monitoring",
                    extra={"error": str(exc)},
                )
                return None

            for subscription in subscriptions:
                transport = subscription.transport or {}
                if (
                    subscription.type == "conduit.shard.disabled"
                    and transport.get("method") == "conduit"
                    and transport.get("conduit_id") == conduit_id
                ):
                    logger.debug(
                        "Found existing conduit.shard.disabled subscription",
                        extra={"subscription_id": subscription.id},
                    )
                    return subscription

            logger.info(
                "Creating conduit.shard.disabled subscription for conduit",
                extra={"conduit_id": conduit_id},
            )

            try:
                client_id = (self.twitch_api.config.twitch_client_id or "").strip()
                if not client_id:
                    logger.error(
                        "Cannot create conduit.shard.disabled subscription without Twitch client ID",
                        extra={"conduit_id": conduit_id},
                    )
                    return None

                request = EventSubSubscriptionRequest(
                    type="conduit.shard.disabled",
                    version="1",
                    condition={"conduit_id": conduit_id, "client_id": client_id},
                    transport=EventSubTransport(
                        method="conduit", conduit_id=conduit_id
                    ),
                )
                return await self.twitch_api.create_eventsub_subscription(request)
            except Exception as exc:
                logger.error(
                    "Failed to create conduit.shard.disabled subscription",
                    extra={"conduit_id": conduit_id, "error": str(exc)},
                )
                return None

    async def handle_disabled_shard(self, conduit_id: str, shard_id: str) -> bool:
        """Recover a disabled shard by provisioning a new transport session."""

        attributes = {"conduit_id": conduit_id, "shard_id": str(shard_id)}
        with span(logger, "handle_disabled_shard", attributes):
            shard_id_str = str(shard_id)

            if self.websocket_manager is None:
                logger.error("Cannot handle disabled shard without WebSocket manager")
                return False

            if self.active_conduit_id and conduit_id != self.active_conduit_id:
                logger.warning(
                    "Handling shard for non-active conduit",
                    extra={
                        "requested_conduit": conduit_id,
                        "active_conduit": self.active_conduit_id,
                    },
                )

            await self.websocket_manager.stop_shard(shard_id_str)

            transport = EventSubTransport(method="websocket")
            updated_shards = await self.assign_transports_to_conduit_shards(
                conduit_id=conduit_id,
                transport_configs={shard_id_str: transport},
            )

            if updated_shards:
                logger.info(
                    "Successfully reactivated disabled shard",
                    extra={"conduit_id": conduit_id, "shard_id": shard_id_str},
                )
                return True

            logger.error(
                "Failed to reactivate disabled shard",
                extra={"conduit_id": conduit_id, "shard_id": shard_id_str},
            )
            return False

    async def reactivate_disabled_shards(self, conduit_id: str) -> Dict[str, Any]:
        """Inspect conduit shards and restart any that are disabled."""

        attributes = {"conduit_id": conduit_id}
        with span(logger, "reactivate_disabled_conduit_shards", attributes):
            try:
                shards = await self.twitch_api.get_conduit_shards(conduit_id)
            except Exception as exc:  # pragma: no cover - defensive logging
                logger.error(
                    "Failed to fetch conduit shards while attempting reactivation",
                    extra={"conduit_id": conduit_id, "error": str(exc)},
                )
                return {"attempted": 0, "success": 0, "results": {}}

            disabled: List[Tuple[str, str]] = []
            for shard in shards:
                shard_id = str(shard.id)
                status = (shard.status or "").lower()
                transport_method = (
                    shard.transport.method
                    if getattr(shard, "transport", None)
                    else None
                )
                session_id = (
                    shard.transport.session_id
                    if getattr(shard, "transport", None)
                    else None
                )

                if status and status != "enabled":
                    disabled.append((shard_id, status))
                elif transport_method == "websocket" and not session_id:
                    disabled.append((shard_id, "missing_session"))

            if not disabled:
                logger.debug(
                    "No disabled conduit shards detected",
                    extra={"conduit_id": conduit_id},
                )
                return {"attempted": 0, "success": 0, "results": {}}

            results: Dict[str, Dict[str, Any]] = {}
            success_count = 0

            for shard_id, reason in disabled:
                try:
                    recovered = await self.handle_disabled_shard(conduit_id, shard_id)
                except Exception as exc:  # pragma: no cover - defensive logging
                    logger.error(
                        "Exception while reactivating disabled shard",
                        extra={
                            "conduit_id": conduit_id,
                            "shard_id": shard_id,
                            "reason": reason,
                            "error": str(exc),
                        },
                    )
                    recovered = False

                results[shard_id] = {"recovered": recovered, "reason": reason}
                if recovered:
                    success_count += 1

            logger.info(
                "Attempted conduit shard reactivation",
                extra={
                    "conduit_id": conduit_id,
                    "attempted": len(disabled),
                    "successful": success_count,
                    "shards": results,
                },
            )

            return {
                "attempted": len(disabled),
                "success": success_count,
                "results": results,
            }

    async def monitor_websocket_shards(self, conduit_id: str):
        """
        Monitor WebSocket shards and restart them if disconnected.

        Args:
            conduit_id (str): The ID of the conduit whose shards to monitor
        """
        with span(logger, "monitor_websocket_shards", {"conduit_id": conduit_id}):
            logger.info(
                f"Starting WebSocket shard monitoring for conduit: {conduit_id}"
            )

            if not self.websocket_manager:
                logger.error("WebSocketManager is not initialized")
                return

            connection_closed_events = self.websocket_manager.connection_closed_events
            if not connection_closed_events:
                logger.error("No connection closed events found")
                return

            # Log the number of shards being monitored
            logger.info(f"Monitoring {len(connection_closed_events)} WebSocket shards")

            # Create tasks for all existing shards
            connection_closed_events_coros = [
                asyncio.create_task(_wait_for_disconnect(shard_id, event))
                for shard_id, event in connection_closed_events.items()
            ]

            while True:
                if not connection_closed_events_coros:
                    logger.warning(
                        "No shard monitoring tasks available, stopping monitor"
                    )
                    break

                logger.debug(
                    f"Waiting for shard disconnection events ({len(connection_closed_events_coros)} shards monitored)"
                )
                done, pending = await asyncio.wait(
                    connection_closed_events_coros, return_when=asyncio.FIRST_COMPLETED
                )

                # Update the list of pending tasks
                connection_closed_events_coros = list(pending)

                for task in done:
                    try:
                        shard_id = task.result()
                        logger.info(
                            f"WebSocket shard {shard_id} disconnected. Initiating restart..."
                        )

                        if self.websocket_manager:
                            logger.debug(f"Starting WebSocket shard {shard_id}")
                            await self.websocket_manager.start_shard(shard_id)

                            logger.debug(f"Configuring transport for shard {shard_id}")
                            transport_config = EventSubTransport(method="websocket")
                            await self.assign_transports_to_conduit_shards(
                                conduit_id, {shard_id: transport_config}
                            )
                            logger.info(
                                f"Successfully restarted WebSocket shard {shard_id}"
                            )

                            # Add new monitoring task for this shard
                            if (
                                shard_id
                                in self.websocket_manager.connection_closed_events
                            ):
                                new_event = (
                                    self.websocket_manager.connection_closed_events[
                                        shard_id
                                    ]
                                )
                                new_task = asyncio.create_task(
                                    _wait_for_disconnect(shard_id, new_event)
                                )
                                connection_closed_events_coros.append(new_task)
                                logger.debug(
                                    f"Added new monitoring task for shard {shard_id}"
                                )
                    except Exception as e:
                        logger.error(
                            f"Error handling shard {shard_id} reconnection: {str(e)}"
                        )

    async def assign_transports_to_conduit_shards(
        self,
        conduit_id: str,
        transport_configs: List[EventSubTransport]
        | Dict[str, EventSubTransport]
        | Dict[int, EventSubTransport],
    ) -> List[EventSubConduitShard]:
        """
        Assign transports (Webhooks or WebSockets) to shards within a Conduit.

        Args:
            conduit_id (str): The ID of the Conduit.
            transport_configs (List[EventSubTransport]): List of transport configurations per shard.

        Returns:
            List[EventSubConduitShard]: Updated shard details.
        """
        if isinstance(transport_configs, dict):
            transport_items: List[Tuple[str, EventSubTransport]] = [
                (str(shard_id), transport)
                for shard_id, transport in transport_configs.items()
            ]
        else:
            transport_items = [
                (str(shard_id), transport)
                for shard_id, transport in enumerate(transport_configs)
            ]

        attributes = {
            "conduit_id": conduit_id,
            "transport_count": len(transport_items),
            "transport_methods": [transport.method for _, transport in transport_items],
        }

        with span(logger, "assign_transports_to_conduit_shards", attributes):
            logger.info(
                f"Assigning {len(transport_items)} transports to conduit {conduit_id}"
            )

            shards: List[EventSubConduitShard] = []
            for shard_id_str, transport in transport_items:
                logger.debug(
                    f"Processing transport {transport.method} for shard {shard_id_str}"
                )
                method = transport.method

                if method == "webhook":
                    callback = transport.callback
                    secret = transport.secret
                    if not callback or not secret:
                        logger.error(
                            f"Webhook configuration missing 'callback' or 'secret' for shard {shard_id_str}"
                        )
                        continue
                    if not self.webhook_manager:
                        logger.error(
                            "WebhookManager is not initialized; cannot assign webhook shard"
                        )
                        continue
                    logger.info(
                        f"Adding webhook shard {shard_id_str} with callback URL: {callback}"
                    )
                    await self.webhook_manager.add_webhook_shard(
                        shard_id=shard_id_str, callback_url=callback, secret=secret
                    )

                elif method == "websocket":
                    # Start a new WebSocket shard
                    logger.info(f"Starting WebSocket shard {shard_id_str}")
                    await self.websocket_manager.start_shard(shard_id=shard_id_str)
                    try:
                        welcome_event = self.websocket_manager.welcome_events[
                            shard_id_str
                        ]
                        logger.debug(
                            f"Waiting for welcome event on shard {shard_id_str}"
                        )
                        await asyncio.wait_for(welcome_event.wait(), timeout=10)
                        logger.info(
                            f"WebSocket shard {shard_id_str} connected successfully"
                        )
                    except asyncio.TimeoutError:
                        logger.error(
                            f"Timeout waiting for WebSocket shard {shard_id_str} to connect"
                        )
                        continue
                    except KeyError:
                        logger.error(
                            f"No welcome event found for WebSocket shard {shard_id_str}"
                        )
                        continue

                    transport.session_id = self.websocket_manager.session_ids[
                        shard_id_str
                    ]
                    logger.debug(
                        f"Assigned session ID {transport.session_id} to shard {shard_id_str}"
                    )
                    # The WebSocketManager will handle assigning subscriptions once connected
                else:
                    logger.error(
                        f"Unsupported transport method '{method}' for shard {shard_id_str}"
                    )
                    continue

                shard = EventSubConduitShard(id=shard_id_str, transport=transport)
                shards.append(shard)
                logger.debug(
                    f"Added shard {shard_id_str} to list with transport method {method}"
                )

            # Update the conduit with the assigned shards
            if not shards:
                logger.warning(f"No valid shards to assign to conduit {conduit_id}")
                return []

            logger.info(f"Updating conduit {conduit_id} with {len(shards)} shards")
            try:
                updated_shards = await self.twitch_api.update_eventsub_conduit_shards(
                    conduit_id=conduit_id, shards=shards
                )
                logger.info(
                    f"Successfully assigned {len(updated_shards)} transports to conduit {conduit_id}"
                )
                return updated_shards
            except Exception as e:
                logger.error(
                    f"Failed to assign transports to shards in conduit {conduit_id}: {e}"
                )
                return []

    async def list_conduits(self) -> List[EventSubConduit]:
        """
        List all EventSub Conduits.

        Returns:
            List[EventSubConduit]: List of Conduits.
        """
        conduits = await self.twitch_api.list_eventsub_conduits()
        for conduit in conduits:
            self.conduits[conduit.id] = conduit
        logger.info(f"Retrieved {len(conduits)} conduits.")
        return conduits

    async def delete_conduit(self, conduit_id: str) -> bool:
        """
        Delete an EventSub Conduit.

        Args:
            conduit_id (str): The ID of the Conduit to delete.

        Returns:
            bool: True if deletion was successful, False otherwise.
        """
        success = await self.twitch_api.delete_eventsub_conduit(conduit_id)
        if success:
            del self.conduits[conduit_id]
            logger.info(f"Deleted Conduit with ID: {conduit_id}")
        return success

    async def shutdown(self):
        """
        Shutdown all managed transports (WebSockets and Webhooks).
        """
        if self.webhook_manager:
            await self.webhook_manager.shutdown()
        if self.websocket_manager:
            await self.websocket_manager.stop_all()
        logger.info("ConduitManager shutdown complete.")
