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.listen import Listener
from utils.logging import logger

JOBS_CHANGED_CHANNEL = "jobs_changed"
POLL_INTERVAL_SECONDS = 60.0
RESTART_SECONDS = 10.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
        try:
            await handler(ctx)
        except (FloodWait, FloodPremiumWait) as exc:
            wait = exc.value if isinstance(exc.value, int) else 1
            await repository.defer_flood_wait(self.pool, ctx.job_id, wait)
            logger.warning(
                f"[yellow]FloodWait {wait}s on job {ctx.job_id} ({job.kind}), "
                f"deferred[/]"
            )
        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))
        else:
            await repository.finish(self.pool, ctx.job_id, JobStatus.DONE)

    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:
        while True:
            try:
                await self._serve()
            except Exception:
                logger.exception(
                    f"[red]Job consumer for account {self.account_id} failed,[/] "
                    f"restarting in {RESTART_SECONDS:.0f}s"
                )
            await asyncio.sleep(RESTART_SECONDS)

    async def _serve(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).[/]")
        listener = Listener(
            f"jobs {self.account_id}",
            {JOBS_CHANGED_CHANNEL: lambda _payload: self._wake.set()},
            self._wake.set,
        )
        await listener.start()
        logger.info(f"[green]Job consumer running for account {self.account_id}.[/]")
        try:
            while True:
                self._wake.clear()
                await self.drain()
                with contextlib.suppress(TimeoutError):
                    await asyncio.wait_for(
                        self._wake.wait(), timeout=await self._idle_seconds()
                    )
        finally:
            await listener.stop()

    async def _idle_seconds(self) -> float:
        due = await repository.seconds_to_next_due(self.pool, self.account_id)
        return POLL_INTERVAL_SECONDS if due is None else min(due, POLL_INTERVAL_SECONDS)