"""Worker runner: long-running process that claims and processes jobs."""

from __future__ import annotations

import asyncio
import logging
import signal

from pycharter.worker.config import WorkerConfig
from pycharter.worker.job_sources.base import JobSource
from pycharter.worker.job_sources.database import DatabaseJobSource
from pycharter.worker.processor import JobProcessor

logger = logging.getLogger(__name__)


class Worker:
    """PyCharter worker: long-running job processor."""

    def __init__(
        self,
        config: WorkerConfig | None = None,
        *,
        job_source: JobSource | None = None,
    ) -> None:
        self._config = config or WorkerConfig()
        self._custom_job_source = job_source
        self._job_source: JobSource | None = None
        self._processor: JobProcessor | None = None
        self._shutdown_event = asyncio.Event()
        self._in_flight = 0
        self._started = False

    @property
    def config(self) -> WorkerConfig:
        return self._config

    async def start(self) -> None:
        """Initialize job source and processor. Call before run()."""
        config = self._config
        db_url = config.resolve_db_url()
        if self._custom_job_source is not None:
            self._job_source = self._custom_job_source
        else:
            self._job_source = DatabaseJobSource(db_url, config=config)
        self._processor = JobProcessor(self._job_source, db_url)
        self._started = True
        logger.info(
            "Worker %s started (concurrency=%d, poll=%dms)",
            config.worker_id,
            config.concurrency,
            config.poll_interval_ms,
        )

    async def run(self) -> None:
        """Run pollers until shutdown. Call start() first."""
        if not self._started or self._job_source is None or self._processor is None:
            raise RuntimeError("Worker.start() must be called before run()")

        loop = asyncio.get_running_loop()
        for sig in (signal.SIGTERM, signal.SIGINT):
            try:
                loop.add_signal_handler(sig, lambda s=sig: self._request_shutdown(s))
            except NotImplementedError:
                pass

        config = self._config
        tasks = []
        for idx in range(config.concurrency):
            task = asyncio.create_task(
                self._poll_loop(idx),
                name=f"worker-poller-{idx}",
            )
            tasks.append(task)

        logger.info(
            "Worker %s running with %d poller(s)",
            config.worker_id,
            config.concurrency,
        )

        await self._shutdown_event.wait()
        logger.info("Shutdown requested, draining in-flight jobs...")
        for t in tasks:
            t.cancel()
        await self._drain(config.drain_timeout_s)
        await self._job_source.close()
        logger.info("Worker %s stopped", config.worker_id)

    async def shutdown(self) -> None:
        """Programmatically trigger a graceful shutdown."""
        self._shutdown_event.set()

    def _request_shutdown(self, sig: signal.Signals | None = None) -> None:
        logger.info("Received %s, initiating shutdown", sig)
        self._shutdown_event.set()

    async def _drain(self, timeout_s: int) -> None:
        if self._in_flight == 0:
            return
        logger.info("Waiting for %d in-flight job(s)...", self._in_flight)
        loop = asyncio.get_running_loop()
        deadline = loop.time() + timeout_s
        while self._in_flight > 0:
            remaining = deadline - loop.time()
            if remaining <= 0:
                logger.warning(
                    "Drain timeout (%ds): %d job(s) still in flight",
                    timeout_s,
                    self._in_flight,
                )
                break
            await asyncio.sleep(min(0.5, remaining))

    async def _poll_loop(self, idx: int) -> None:
        assert self._job_source is not None
        assert self._processor is not None
        config = self._config
        poll_interval = config.poll_interval_ms / 1000.0
        try:
            while not self._shutdown_event.is_set():
                claimed = await self._job_source.claim_next(config.worker_id)
                if claimed is not None:
                    self._in_flight += 1
                    try:
                        await self._processor.process_one(claimed)
                    except Exception as exc:
                        logger.exception(
                            "Poller-%d: unhandled error processing %s: %s",
                            idx,
                            claimed.event_id,
                            exc,
                        )
                        try:
                            await self._job_source.fail(
                                claimed.event_id, str(exc), retry=True
                            )
                        except Exception:
                            logger.exception(
                                "Failed to mark job %s as failed", claimed.event_id
                            )
                    finally:
                        self._in_flight -= 1
                else:
                    await asyncio.sleep(poll_interval)
        except asyncio.CancelledError:
            pass
