import asyncio
import logging

from app.core.config import settings
from app.models.health import HealthSnapshot
from app.models.proxy import ProxyDocument
from app.repositories.health import record_health_sample
from app.repositories.proxy import list_proxies
from app.services.mitmproxy.health import check_proxy

logger = logging.getLogger(__name__)


class HealthMonitor:
    def __init__(self):
        self.task: asyncio.Task[None] | None = None
        self.subscribers: dict[int, set[asyncio.Queue[HealthSnapshot]]] = {}

    async def start(self) -> None:
        if self.task is None or self.task.done():
            self.task = asyncio.create_task(self.run())

    async def stop(self) -> None:
        if self.task is None:
            return
        self.task.cancel()
        await asyncio.gather(self.task, return_exceptions=True)
        self.task = None

    async def run(self) -> None:
        loop = asyncio.get_running_loop()
        next_cycle = loop.time()
        while True:
            delay = next_cycle - loop.time()
            if delay > 0:
                await asyncio.sleep(delay)

            cycle_started = loop.time()
            try:
                await self.run_cycle()
            except asyncio.CancelledError:
                raise
            except Exception:
                logger.exception("Health monitor cycle failed")
                next_cycle = loop.time() + min(
                    1.0,
                    float(settings.HEALTHCHECK_INTERVAL_SECONDS),
                )
                continue

            next_cycle = cycle_started + settings.HEALTHCHECK_INTERVAL_SECONDS

    async def run_cycle(self) -> None:
        proxies = [
            proxy
            for proxy in await list_proxies()
            if proxy.status == "running" and proxy.healthcheck.enabled
        ]
        async with asyncio.TaskGroup() as group:
            for proxy in proxies:
                group.create_task(self.run_check(proxy))

    async def run_check(self, proxy: ProxyDocument) -> None:
        try:
            snapshot = await check_proxy(proxy)
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception(
                "Health check failed for proxy %s",
                proxy.proxy_port,
            )
            return

        try:
            await record_health_sample(snapshot)
        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception(
                "Health snapshot persistence failed for proxy %s",
                proxy.proxy_port,
            )

        await self.publish(proxy.proxy_port, snapshot)

    async def publish(self, proxy_port: int, snapshot: HealthSnapshot) -> None:
        for queue in list(self.subscribers.get(proxy_port, set())):
            try:
                queue.put_nowait(snapshot)
            except asyncio.QueueFull:
                logger.debug(
                    "Dropped health update for slow subscriber on proxy %s",
                    proxy_port,
                )
                continue

    def subscribe(self, proxy_port: int) -> asyncio.Queue[HealthSnapshot]:
        queue: asyncio.Queue[HealthSnapshot] = asyncio.Queue(maxsize=8)
        listeners = self.subscribers.setdefault(proxy_port, set())
        listeners.add(queue)
        return queue

    def unsubscribe(
        self,
        proxy_port: int,
        queue: asyncio.Queue[HealthSnapshot],
    ) -> None:
        listeners = self.subscribers.get(proxy_port)
        if listeners is None:
            return
        listeners.discard(queue)
        if not listeners:
            self.subscribers.pop(proxy_port, None)


health_monitor = HealthMonitor()