import asyncio
import contextlib

import asyncpg
from pyrogram import Client
from pyrogram.errors import FloodPremiumWait, FloodWait

from userbot.modules.jobs import repository
from userbot.modules.jobs.context import JobContext
from userbot.modules.jobs.registry import JOB_HANDLERS, JobHandler
from utils.db.models import Job, JobStatus
from utils.env import env
from utils.logging import logger

JOBS_CHANGED_CHANNEL = "jobs_changed"
POLL_INTERVAL_SECONDS = 60.0


class JobConsumer:
    def __init__(
        self,
        client: Client | None,
        pool: asyncpg.Pool,
        account_id: int,
        registry: dict[str, JobHandler] | None = None,
    ) -> None:
        self.client = client
        self.pool = pool
        self.account_id = account_id
        self.registry = registry if registry is not None else JOB_HANDLERS
        self._wake = asyncio.Event()

    async def execute(self, job: Job) -> None:
        ctx = JobContext(self.client, self.pool, self.account_id, job)
        handler = self.registry.get(job.kind)
        if handler is None:
            await repository.finish(
                self.pool, ctx.job_id, JobStatus.FAILED, f"no handler for {job.kind!r}"
            )
            return
        while True:
            try:
                await handler(ctx)
            except (FloodWait, FloodPremiumWait) as exc:
                await repository.bump_flood_wait(self.pool, ctx.job_id)
                wait = exc.value if isinstance(exc.value, int) else 1
                logger.warning(
                    f"[yellow]FloodWait {wait}s on job {ctx.job_id} ({job.kind})[/]"
                )
                await asyncio.sleep(wait)
            except Exception as exc:
                logger.exception(f"Job {ctx.job_id} ({job.kind}) failed")
                await repository.finish(
                    self.pool, ctx.job_id, JobStatus.FAILED, str(exc)
                )
                return
            else:
                await repository.finish(self.pool, ctx.job_id, JobStatus.DONE)
                return

    async def drain(self) -> None:
        while True:
            job = await repository.claim_next(self.pool, self.account_id)
            if job is None:
                return
            await self.execute(job)

    async def run(self) -> None:
        requeued = await repository.requeue_running(self.pool, self.account_id)
        if requeued:
            logger.info(f"[yellow]Requeued {requeued} stale running job(s).[/]")
        pruned = await repository.prune_finished(self.pool, self.account_id)
        if pruned:
            logger.info(f"[yellow]Pruned {pruned} finished job(s).[/]")
        conn = await asyncpg.connect(dsn=env.db.connection_url)
        await conn.add_listener(JOBS_CHANGED_CHANNEL, lambda *_: self._wake.set())
        logger.info(f"[green]Job consumer running for account {self.account_id}.[/]")
        try:
            while True:
                await self.drain()
                self._wake.clear()
                with contextlib.suppress(TimeoutError):
                    await asyncio.wait_for(
                        self._wake.wait(), timeout=POLL_INTERVAL_SECONDS
                    )
        finally:
            with contextlib.suppress(Exception):
                await conn.close()