import asyncio
import contextlib
import signal

import asyncpg
import uvloop

from dependencies.container import container
from userbot.modules.uptime import repository as uptime
from utils.listen import Listener
from utils.logging import logger, setup_logging
from utils.read.bots import BOTS_CHANGED_CHANNEL
from utils.read.watches import WATCHES_CHANGED_CHANNEL
from utils.storage import ContentAddressedStorage
from watcher.registry import BotRegistry

setup_logging()

RETRY_INTERVAL = 300
UPTIME_INTERVAL = 60


async def _retry_loop(registry: BotRegistry) -> None:
    while True:
        await asyncio.sleep(RETRY_INTERVAL)
        try:
            await registry.sync()
        except Exception:
            logger.exception("[red]Bot sync failed[/]")


async def _keep_uptime(registry: BotRegistry, pool: asyncpg.Pool) -> None:
    while True:
        for account_id in registry.connected_accounts:
            with contextlib.suppress(Exception):
                await uptime.heartbeat(pool, account_id)
        await asyncio.sleep(UPTIME_INTERVAL)


def _change_listener(registry: BotRegistry, tasks: set[asyncio.Task]) -> Listener:
    def on_change(_payload: str) -> None:
        task = asyncio.create_task(registry.sync())
        tasks.add(task)
        task.add_done_callback(tasks.discard)

    def on_watch(_payload: str) -> None:
        task = asyncio.create_task(registry.refresh_watches())
        tasks.add(task)
        task.add_done_callback(tasks.discard)

    def on_reconnect() -> None:
        on_change("")
        on_watch("")

    return Listener(
        "watcher",
        {BOTS_CHANGED_CHANNEL: on_change, WATCHES_CHANGED_CHANNEL: on_watch},
        on_reconnect,
    )


async def runner() -> None:
    pool = await container.get(asyncpg.Pool)
    storage = await container.get(ContentAddressedStorage)
    registry = BotRegistry(pool, storage)
    tasks: set[asyncio.Task] = set()
    listener = _change_listener(registry, tasks)
    stop = asyncio.Event()
    asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, stop.set)
    retry = asyncio.create_task(_retry_loop(registry))
    beat = asyncio.create_task(_keep_uptime(registry, pool))
    try:
        await registry.sync()
        if not registry.running:
            logger.warning("[yellow]No bots yet. Add a bot token in the web UI.[/]")
        await listener.start()
        logger.info("[green]Bot watcher running.[/]")
        await stop.wait()
    finally:
        for task in (retry, beat):
            task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await task
        await listener.stop()
        await registry.close()
        await container.close()


def main() -> None:
    uvloop.install()
    logger.info("Starting bot watcher...")
    with contextlib.suppress(KeyboardInterrupt):
        asyncio.run(runner())
    logger.info("[red]Bot watcher stopped.[/]")