# clients/webhook_manager.py

import asyncio
from typing import Callable, Dict, Any, Optional

import uvicorn
from fastapi import FastAPI

from common.logging import get_logger

logger = get_logger(__name__)


class WebhookServer:
    """
    A simple HTTP server for handling incoming webhook events.
    """

    def __init__(self, secret: str, event_handler: Callable[[Dict[str, Any]], None]):
        """
        Initialize the WebhookServer.

        Args:
            secret (str): The secret for signature verification.
            event_handler (Callable[[Dict[str, Any]], None]): Function to handle incoming events.
        """
        self.api = None
        self.secret = secret
        self.event_handler = event_handler
        self.server: Optional[uvicorn.Server] = None

    async def start(self, host: str = "", port: int = 8080):
        """
        Start the HTTP server.

        Args:
            host (str, optional): Host to bind the server to. Defaults to "".
            port (int, optional): Port to bind the server to. Defaults to 8080.
        """
        self.api = FastAPI()

        @self.api.post("/")
        async def handle_event(payload: Dict[str, Any]):
            self.event_handler(payload)

        logger.info(f"Starting Webhook server on {host}:{port}")
        config = uvicorn.Config(self.api, host=host, port=port)
        self.server = uvicorn.Server(config=config)
        await asyncio.create_task(self.server.serve())

    async def stop(self):
        """
        Stop the HTTP server.
        """
        if self.server:
            await self.server.shutdown()
            await self.server.shutdown()
            self.server = None
            logger.info("Webhook server has been stopped.")


class WebhookManager:
    def __init__(
        self,
        event_handler: Callable[[Dict[str, Any]], None],
        host: str = "0.0.0.0",
        port_base: int = 8080,  # Base port for assigning different webhook servers
    ):
        """
        Initialize the WebhookManager.

        Args:
            event_handler (Callable[[Dict[str, Any]], None]): Function to handle incoming events.
            host (str, optional): Host to bind the webhook servers to. Defaults to "0.0.0.0".
            port_base (int, optional): Starting port number for webhook servers. Each server will use a unique port.
        """
        self.event_handler = event_handler
        self.host = host
        self.port_base = port_base
        self.webhook_servers: Dict[str, WebhookServer] = {}
        self.next_port = port_base

    async def add_webhook_shard(
        self, shard_id: str, callback_url: str, secret: str
    ) -> None:
        """
        Add a new Webhook shard by starting a WebhookServer.

        Args:
            shard_id (str): The shard identifier.
            callback_url (str): The callback URL for the webhook.
            secret (str): The secret for signature verification.
        """
        if shard_id in self.webhook_servers:
            logger.warning(f"Webhook shard {shard_id} is already managed.")
            return

        # Assign a unique port for this webhook server
        port = self.next_port
        self.next_port += 1

        server = WebhookServer(secret=secret, event_handler=self.event_handler)
        await server.start(host=self.host, port=port)
        self.webhook_servers[shard_id] = server
        logger.info(
            f"Webhook shard {shard_id} assigned to port {port} with callback {callback_url}"
        )

        # Note: Ensure that the `callback_url` points to the server's public address.
        # For local development, tools like Ngrok can expose the local webhook server to the internet.

    async def remove_webhook_shard(self, shard_id: str) -> None:
        """
        Remove an existing Webhook shard by stopping its WebhookServer.

        Args:
            shard_id (str): The shard identifier.
        """
        server = self.webhook_servers.get(shard_id)
        if not server:
            logger.warning(f"Webhook shard {shard_id} is not managed.")
            return
        await server.stop()
        del self.webhook_servers[shard_id]
        logger.info(f"Webhook shard {shard_id} has been removed and stopped.")

    async def shutdown(self):
        """
        Shutdown all Webhook shards gracefully.
        """
        tasks = [server.stop() for server in self.webhook_servers.values()]
        await asyncio.gather(*tasks, return_exceptions=True)
        self.webhook_servers.clear()
        logger.info("All Webhook shards have been shut down.")


__all__ = ["WebhookServer", "WebhookManager"]
