"""Discovery example: a bounded asyncio worker with graceful shutdown.

Declared runtime: Python 3.14.

The sample is deliberately self-contained. Persistence and transport are
protocols so lifecycle, cancellation, retry, and ownership behavior remain
visible without requiring a particular database or message broker.
"""

from __future__ import annotations

import asyncio
import contextlib
import dataclasses
import datetime as dt
import enum
import json
import logging
import random
import signal
import sys
import uuid
from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine
from typing import Any, Protocol, TypeVar


Json = None | bool | int | float | str | list["Json"] | dict[str, "Json"]
T = TypeVar("T")
UTC = dt.UTC


class JobState(enum.StrEnum):
    PENDING = "pending"
    RUNNING = "running"
    RETRY_WAIT = "retry_wait"
    SUCCEEDED = "succeeded"
    DEAD = "dead"


@dataclasses.dataclass(frozen=True, slots=True)
class Failure:
    name: str
    message: str
    code: str
    retryable: bool


@dataclasses.dataclass(frozen=True, slots=True)
class Job:
    id: uuid.UUID
    tenant_id: uuid.UUID
    operation_key: str
    payload: Json
    state: JobState
    attempts: int
    max_attempts: int
    available_at: dt.datetime
    lease_owner: str | None = None
    lease_token: uuid.UUID | None = None
    lease_expires_at: dt.datetime | None = None


@dataclasses.dataclass(frozen=True, slots=True)
class Claim:
    job: Job
    owner: str
    token: uuid.UUID


@dataclasses.dataclass(frozen=True, slots=True)
class WorkerPolicy:
    concurrency: int = 4
    input_capacity: int = 32
    lease_seconds: float = 30.0
    heartbeat_seconds: float = 10.0
    attempt_timeout_seconds: float = 20.0
    shutdown_timeout_seconds: float = 25.0
    poll_seconds: float = 0.25
    base_retry_seconds: float = 1.0
    max_retry_seconds: float = 60.0

    def validate(self) -> None:
        if self.concurrency < 1:
            raise ValueError("concurrency must be positive")
        if self.input_capacity < self.concurrency:
            raise ValueError("input capacity must be at least concurrency")
        if self.heartbeat_seconds >= self.lease_seconds / 2:
            raise ValueError("heartbeat must be less than half of lease duration")
        if self.attempt_timeout_seconds <= 0 or self.shutdown_timeout_seconds <= 0:
            raise ValueError("timeouts must be positive")


class WorkerError(Exception):
    code = "WORKER_ERROR"
    retryable = True


class PermanentJobError(WorkerError):
    code = "PERMANENT_JOB_ERROR"
    retryable = False


class LeaseLostError(WorkerError):
    code = "LEASE_LOST"


class WorkerStoppingError(WorkerError):
    code = "WORKER_STOPPING"


def serialize_failure(error: BaseException) -> Failure:
    if isinstance(error, asyncio.CancelledError):
        return Failure("CancelledError", "Job attempt was cancelled.", "CANCELLED", True)
    if isinstance(error, WorkerError):
        return Failure(type(error).__name__, str(error), error.code, error.retryable)
    return Failure(type(error).__name__, str(error), "UNEXPECTED_ERROR", True)


def utc_now() -> dt.datetime:
    return dt.datetime.now(UTC)


def retry_delay(policy: WorkerPolicy, attempt: int) -> float:
    exponential = min(
        policy.base_retry_seconds * (2 ** max(attempt - 1, 0)),
        policy.max_retry_seconds,
    )
    return exponential * random.uniform(0.8, 1.2)


class JobStore(Protocol):
    async def enqueue(
        self,
        *,
        tenant_id: uuid.UUID,
        operation_key: str,
        payload: Json,
        max_attempts: int,
    ) -> tuple[Job, bool]: ...

    async def claim_next(
        self,
        *,
        tenant_id: uuid.UUID,
        owner: str,
        lease_seconds: float,
        now: dt.datetime,
    ) -> Claim | None: ...

    async def renew(
        self,
        *,
        claim: Claim,
        lease_seconds: float,
        now: dt.datetime,
    ) -> bool: ...

    async def complete(self, *, claim: Claim, result: Json, now: dt.datetime) -> bool: ...

    async def fail(
        self,
        *,
        claim: Claim,
        failure: Failure,
        available_at: dt.datetime | None,
        now: dt.datetime,
    ) -> bool: ...

    async def release_expired(
        self, *, tenant_id: uuid.UUID, now: dt.datetime, limit: int
    ) -> int: ...


class Transport(Protocol):
    async def deliver(self, job: Job) -> Json: ...


@dataclasses.dataclass(slots=True)
class LogContext:
    logger: logging.Logger
    tenant_id: uuid.UUID
    worker_id: str

    def write(self, level: int, event: str, **fields: Json) -> None:
        self.logger.log(
            level,
            json.dumps(
                {
                    "event": event,
                    "tenant_id": str(self.tenant_id),
                    "worker_id": self.worker_id,
                    **fields,
                },
                sort_keys=True,
            ),
        )


class TaskSupervisor:
    """Own tasks and collect failures without losing references."""

    def __init__(self, log: LogContext) -> None:
        self._log = log
        self._tasks: set[asyncio.Task[Any]] = set()
        self._closed = False

    def create(self, awaitable: Coroutine[Any, Any, T], *, name: str) -> asyncio.Task[T]:
        if self._closed:
            awaitable.close()
            raise RuntimeError("task supervisor is closed")
        task = asyncio.create_task(awaitable, name=name)
        self._tasks.add(task)
        task.add_done_callback(self._done)
        return task

    def _done(self, task: asyncio.Task[Any]) -> None:
        self._tasks.discard(task)
        if task.cancelled():
            return
        try:
            error = task.exception()
        except asyncio.CancelledError:
            return
        if error is not None:
            self._log.write(
                logging.ERROR,
                "background_task_failed",
                task=task.get_name(),
                error=repr(error),
            )

    async def cancel_and_wait(self, reason: str) -> None:
        self._closed = True
        tasks = tuple(self._tasks)
        for task in tasks:
            task.cancel(reason)
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)

    async def wait(self) -> None:
        tasks = tuple(self._tasks)
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)

    @property
    def count(self) -> int:
        return len(self._tasks)


class LeaseHeartbeat:
    def __init__(
        self,
        *,
        store: JobStore,
        claim: Claim,
        policy: WorkerPolicy,
        log: LogContext,
    ) -> None:
        self._store = store
        self._claim = claim
        self._policy = policy
        self._log = log
        self._task: asyncio.Task[None] | None = None
        self._lost = asyncio.Event()
        self._stopping = asyncio.Event()

    async def __aenter__(self) -> LeaseHeartbeat:
        if self._task is not None:
            raise RuntimeError("heartbeat already started")
        self._task = asyncio.create_task(
            self._run(),
            name=f"heartbeat:{self._claim.job.id}",
        )
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        traceback: Any,
    ) -> None:
        self._stopping.set()
        task = self._task
        if task is None:
            return
        task.cancel("heartbeat context closed")
        with contextlib.suppress(asyncio.CancelledError):
            await task
        self._task = None

    async def _run(self) -> None:
        try:
            while not self._stopping.is_set():
                await asyncio.sleep(self._policy.heartbeat_seconds)
                renewed = await self._store.renew(
                    claim=self._claim,
                    lease_seconds=self._policy.lease_seconds,
                    now=utc_now(),
                )
                if not renewed:
                    self._lost.set()
                    self._log.write(
                        logging.WARNING,
                        "lease_lost",
                        job_id=str(self._claim.job.id),
                    )
                    return
        except asyncio.CancelledError:
            raise
        except Exception as error:
            self._lost.set()
            self._log.write(
                logging.ERROR,
                "heartbeat_failed",
                job_id=str(self._claim.job.id),
                error=repr(error),
            )

    async def wait_lost(self) -> None:
        await self._lost.wait()


async def race_lease_and_delivery(
    *,
    heartbeat: LeaseHeartbeat,
    deliver: Awaitable[Json],
) -> Json:
    delivery_task = asyncio.create_task(deliver, name="job-delivery")
    lease_task = asyncio.create_task(heartbeat.wait_lost(), name="lease-loss")
    try:
        done, _ = await asyncio.wait(
            {delivery_task, lease_task},
            return_when=asyncio.FIRST_COMPLETED,
        )
        if lease_task in done:
            delivery_task.cancel("lease lost")
            with contextlib.suppress(asyncio.CancelledError):
                await delivery_task
            raise LeaseLostError("lease was lost during delivery")
        lease_task.cancel("delivery completed")
        with contextlib.suppress(asyncio.CancelledError):
            await lease_task
        return await delivery_task
    finally:
        for task in (delivery_task, lease_task):
            if not task.done():
                task.cancel("race cleanup")


class Worker:
    def __init__(
        self,
        *,
        tenant_id: uuid.UUID,
        store: JobStore,
        transport: Transport,
        policy: WorkerPolicy,
        logger: logging.Logger,
    ) -> None:
        policy.validate()
        self.tenant_id = tenant_id
        self.store = store
        self.transport = transport
        self.policy = policy
        self.worker_id = f"worker-{uuid.uuid4()}"
        self.log = LogContext(logger, tenant_id, self.worker_id)
        self.supervisor = TaskSupervisor(self.log)
        self.queue: asyncio.Queue[Claim] = asyncio.Queue(maxsize=policy.input_capacity)
        self.stopping = asyncio.Event()
        self._stop_reason = "shutdown requested"

    def request_stop(self, reason: str) -> None:
        if self.stopping.is_set():
            return
        self._stop_reason = reason
        self.stopping.set()
        self.log.write(logging.INFO, "shutdown_requested", reason=reason)

    async def run(self) -> None:
        await self.store.release_expired(
            tenant_id=self.tenant_id,
            now=utc_now(),
            limit=100,
        )
        try:
            async with asyncio.TaskGroup() as group:
                group.create_task(self._poll(), name="poller")
                for index in range(self.policy.concurrency):
                    group.create_task(self._consume(index), name=f"consumer:{index}")
                await self.stopping.wait()
                await self._drain_for_shutdown()
        except* WorkerStoppingError:
            pass
        finally:
            await self.supervisor.cancel_and_wait(self._stop_reason)

    async def _poll(self) -> None:
        while not self.stopping.is_set():
            if self.queue.full():
                await asyncio.sleep(self.policy.poll_seconds)
                continue
            claim = await self.store.claim_next(
                tenant_id=self.tenant_id,
                owner=self.worker_id,
                lease_seconds=self.policy.lease_seconds,
                now=utc_now(),
            )
            if claim is None:
                await asyncio.sleep(self.policy.poll_seconds)
                continue
            await self.queue.put(claim)
            self.log.write(
                logging.INFO,
                "job_claimed",
                job_id=str(claim.job.id),
                attempt=claim.job.attempts,
            )
        raise WorkerStoppingError(self._stop_reason)

    async def _consume(self, index: int) -> None:
        while True:
            if self.stopping.is_set() and self.queue.empty():
                raise WorkerStoppingError(self._stop_reason)
            try:
                claim = await asyncio.wait_for(self.queue.get(), timeout=0.25)
            except TimeoutError:
                continue
            try:
                await self._execute(claim, index)
            finally:
                self.queue.task_done()

    async def _execute(self, claim: Claim, consumer: int) -> None:
        self.log.write(
            logging.INFO,
            "job_started",
            job_id=str(claim.job.id),
            consumer=consumer,
        )
        try:
            async with LeaseHeartbeat(
                store=self.store,
                claim=claim,
                policy=self.policy,
                log=self.log,
            ) as heartbeat:
                async with asyncio.timeout(self.policy.attempt_timeout_seconds):
                    result = await race_lease_and_delivery(
                        heartbeat=heartbeat,
                        deliver=self.transport.deliver(claim.job),
                    )
            completed = await self.store.complete(
                claim=claim,
                result=result,
                now=utc_now(),
            )
            if not completed:
                raise LeaseLostError(f"completion rejected for {claim.job.id}")
            self.log.write(logging.INFO, "job_succeeded", job_id=str(claim.job.id))
        except asyncio.CancelledError:
            await self._record_failure(
                claim,
                Failure("CancelledError", "worker stopped", "CANCELLED", True),
            )
            raise
        except TimeoutError:
            await self._record_failure(
                claim,
                Failure("TimeoutError", "attempt timed out", "ATTEMPT_TIMEOUT", True),
            )
        except Exception as error:
            await self._record_failure(claim, serialize_failure(error))

    async def _record_failure(self, claim: Claim, failure: Failure) -> None:
        exhausted = claim.job.attempts >= claim.job.max_attempts
        retryable = failure.retryable and not exhausted
        available_at = (
            utc_now() + dt.timedelta(seconds=retry_delay(self.policy, claim.job.attempts))
            if retryable
            else None
        )
        recorded = await self.store.fail(
            claim=claim,
            failure=dataclasses.replace(failure, retryable=retryable),
            available_at=available_at,
            now=utc_now(),
        )
        self.log.write(
            logging.WARNING,
            "job_retry_scheduled" if retryable else "job_dead",
            job_id=str(claim.job.id),
            code=failure.code,
            recorded=recorded,
            available_at=available_at.isoformat() if available_at else None,
        )

    async def _drain_for_shutdown(self) -> None:
        try:
            async with asyncio.timeout(self.policy.shutdown_timeout_seconds):
                await self.queue.join()
                await self.supervisor.wait()
        except TimeoutError:
            self.log.write(
                logging.WARNING,
                "shutdown_deadline_exceeded",
                queued=self.queue.qsize(),
                background=self.supervisor.count,
            )
            await self.supervisor.cancel_and_wait("shutdown deadline exceeded")


@contextlib.asynccontextmanager
async def installed_signal_handlers(worker: Worker) -> AsyncIterator[None]:
    loop = asyncio.get_running_loop()
    installed: list[signal.Signals] = []

    def stop_for(sig: signal.Signals) -> None:
        worker.request_stop(sig.name)

    if sys.platform != "win32":
        for sig in (signal.SIGTERM, signal.SIGINT):
            loop.add_signal_handler(sig, stop_for, sig)
            installed.append(sig)
    try:
        yield
    finally:
        for sig in installed:
            loop.remove_signal_handler(sig)


async def serve(worker: Worker) -> None:
    async with installed_signal_handlers(worker):
        await worker.run()


def configure_logging() -> logging.Logger:
    logger = logging.getLogger("discovery.worker")
    logger.setLevel(logging.INFO)
    if not logger.handlers:
        handler = logging.StreamHandler()
        handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
        logger.addHandler(handler)
    return logger


async def main(store: JobStore, transport: Transport, tenant_id: uuid.UUID) -> None:
    worker = Worker(
        tenant_id=tenant_id,
        store=store,
        transport=transport,
        policy=WorkerPolicy(),
        logger=configure_logging(),
    )
    await serve(worker)