"""Reference solution — Async Pipeline with Backpressure.

DO NOT READ BEFORE YOU HAVE RUN THE PROBLEM UNDER THE CLOCK.
Teaching text: ../../../WARMUP.md § Chapter 10.

An unbounded queue does not absorb overload. It converts a throughput problem
into a latency problem and then into an out-of-memory crash: by Little's law,
10,000 items queued and served at 100/s means every new arrival waits 100
seconds, so by the time you notice, everything in the queue is already useless
-- and the crash loses in-flight work that a rejection would have let the
client retry. The bound is what makes the producer FEEL the consumer's
slowness.

Two async-specific rules this design depends on:

  * CancelledError inherits BaseException, not Exception. So `except
    Exception` correctly does NOT swallow it -- and if you catch it explicitly
    for cleanup you MUST re-raise, or the worker becomes uncancellable and the
    shutdown deadline turns into a hang.

  * Waiting forever on a full queue turns a bounded queue back into an
    unbounded one, because the backlog just moves into the callers. Hence
    shed_after: wait a bounded time, then shed.
"""

from __future__ import annotations

import asyncio


class QueueFull(Exception):
    """The item was shed rather than queued."""


class Pipeline:
    def __init__(self, handler, workers=4, max_queue=100,
                 item_timeout=None, shed_after=None):
        if workers <= 0:
            raise ValueError("workers must be positive")
        if max_queue <= 0:
            raise ValueError("max_queue must be positive")
        self._handler = handler
        self._n_workers = workers
        self._max_queue = max_queue
        self._item_timeout = item_timeout
        self._shed_after = shed_after
        self._queue = None
        self._workers = []
        self._running = False
        self.errors = []
        self.stats = {"accepted": 0, "shed": 0, "done": 0,
                      "failed": 0, "timeout": 0}

    # ---- lifecycle -------------------------------------------------------
    async def start(self):
        if self._running:
            return
        self._queue = asyncio.Queue(maxsize=self._max_queue)   # the bound IS
        self._running = True                                   # the backpressure
        self._workers = [asyncio.create_task(self._worker())
                         for _ in range(self._n_workers)]

    @property
    def running(self):
        return self._running

    @property
    def qsize(self):
        return self._queue.qsize() if self._queue else 0

    # ---- submission ------------------------------------------------------
    async def submit(self, item, block=True):
        if not self._running:
            raise RuntimeError("pipeline is not running")
        if not block:
            try:
                self._queue.put_nowait(item)
            except asyncio.QueueFull:
                self.stats["shed"] += 1
                raise QueueFull("queue full — shed") from None
        elif self._shed_after is None:
            await self._queue.put(item)              # pure backpressure
        else:
            try:
                await asyncio.wait_for(self._queue.put(item), self._shed_after)
            except (asyncio.TimeoutError, TimeoutError):
                self.stats["shed"] += 1
                raise QueueFull("queue full — shed after waiting") from None
        self.stats["accepted"] += 1

    # ---- workers ---------------------------------------------------------
    async def _worker(self):
        while True:
            item = await self._queue.get()
            try:
                if item is None:                     # sentinel: drain complete
                    return
                try:
                    if self._item_timeout is None:
                        await self._handler(item)
                    else:
                        await asyncio.wait_for(self._handler(item),
                                               self._item_timeout)
                    self.stats["done"] += 1
                except (asyncio.TimeoutError, TimeoutError):
                    self.stats["timeout"] += 1
                except asyncio.CancelledError:
                    # BaseException, not Exception. Observe it and RE-RAISE;
                    # swallowing it makes the worker uncancellable.
                    raise
                except Exception as exc:
                    self.stats["failed"] += 1
                    self.errors.append(exc)          # collected, not swallowed
            finally:
                self._queue.task_done()

    # ---- shutdown --------------------------------------------------------
    async def shutdown(self, drain_timeout=10.0):
        """Idempotent: stop accepting, drain, then cancel at a hard deadline."""
        if not self._running:
            return self.stats
        self._running = False                        # 1. stop accepting

        for _ in self._workers:                      # 2. one sentinel each, so
            await self._queue.put(None)              #    every worker exits once

        done, pending = await asyncio.wait(self._workers, timeout=drain_timeout)
        for task in pending:                         # 3. the hard deadline
            task.cancel()
        if pending:
            await asyncio.gather(*pending, return_exceptions=True)
        self._workers = []
        return self.stats