#!/usr/bin/env python3
"""Asyncio cancellation and timeout verification tool.

This script demonstrates and tests a few operationally useful patterns:
- cooperative cancellation with CancelledError
- timeout boundaries using asyncio.wait_for()
- cleanup behavior after cancellation
- distinguishing timeout from cancellation and other failures

It is intentionally safe: it does not call external services, write files,
or require secrets. Use it as a starting point for validating your own
asyncio task behavior before production use.
"""

from __future__ import annotations

import argparse
import asyncio
from dataclasses import dataclass
from typing import Awaitable, Callable, Optional


@dataclass
class RunResult:
    name: str
    outcome: str
    detail: str


async def cooperative_worker(delay: float, cleanup_delay: float = 0.0) -> str:
    """A sample coroutine that responds to cancellation and performs cleanup.

    Args:
        delay: Time to simulate useful work.
        cleanup_delay: Optional time spent in cleanup after cancellation.

    Returns:
        A short success string.

    Raises:
        asyncio.CancelledError: Re-raised after cleanup so cancellation remains visible.
    """
    try:
        await asyncio.sleep(delay)
        return f"completed after {delay:.2f}s"
    except asyncio.CancelledError:
        # Cleanup should be idempotent and quick where possible.
        if cleanup_delay > 0:
            await asyncio.sleep(cleanup_delay)
        raise


async def blocking_shape_worker(duration: float) -> str:
    """A placeholder for CPU-bound or blocking work.

    This intentionally yields once at the start and then simulates work using
    sleep so the example remains safe. In real code, move blocking work out of
    the event loop using asyncio.to_thread() or a process pool.
    """
    await asyncio.sleep(0)
    await asyncio.sleep(duration)
    return f"blocking-shape task finished in {duration:.2f}s"


async def run_with_timeout(
    name: str,
    coro_factory: Callable[[], Awaitable[str]],
    timeout: float,
) -> RunResult:
    """Run a coroutine with a timeout and classify the outcome."""
    try:
        result = await asyncio.wait_for(coro_factory(), timeout=timeout)
        return RunResult(name=name, outcome="success", detail=result)
    except TimeoutError:
        return RunResult(
            name=name,
            outcome="timeout",
            detail=f"operation exceeded timeout of {timeout:.2f}s",
        )
    except asyncio.CancelledError:
        return RunResult(name=name, outcome="cancelled", detail="task was cancelled")
    except Exception as exc:  # noqa: BLE001 - report unexpected failures clearly
        return RunResult(name=name, outcome="error", detail=f"{type(exc).__name__}: {exc}")


async def manual_cancel_demo(work_delay: float, cancel_after: float, cleanup_delay: float) -> RunResult:
    """Create a task, cancel it, and confirm that cancellation propagates."""
    task = asyncio.create_task(cooperative_worker(work_delay, cleanup_delay=cleanup_delay))
    await asyncio.sleep(cancel_after)
    task.cancel()

    try:
        result = await task
        return RunResult(name="manual-cancel", outcome="unexpected-success", detail=result)
    except asyncio.CancelledError:
        return RunResult(
            name="manual-cancel",
            outcome="cancelled",
            detail=f"task cancelled after {cancel_after:.2f}s and re-raised CancelledError",
        )
    except Exception as exc:  # noqa: BLE001
        return RunResult(name="manual-cancel", outcome="error", detail=f"{type(exc).__name__}: {exc}")


async def timeout_demo(work_delay: float, timeout: float, cleanup_delay: float) -> RunResult:
    """Show how wait_for() behaves when the timeout expires."""
    return await run_with_timeout(
        name="timeout-demo",
        coro_factory=lambda: cooperative_worker(work_delay, cleanup_delay=cleanup_delay),
        timeout=timeout,
    )


async def sibling_cancellation_demo(work_delays: list[float], timeout: float) -> RunResult:
    """Run multiple tasks and cancel remaining work on timeout.

    This models a fan-out request where one slow dependency should not keep the
    whole request alive indefinitely.
    """
    tasks = [asyncio.create_task(cooperative_worker(delay)) for delay in work_delays]
    try:
        done, pending = await asyncio.wait(tasks, timeout=timeout)
        if pending:
            for task in pending:
                task.cancel()
            await asyncio.gather(*pending, return_exceptions=True)
            return RunResult(
                name="sibling-cancel",
                outcome="timeout",
                detail=f"{len(done)} task(s) finished, {len(pending)} task(s) cancelled after timeout",
            )
        results = await asyncio.gather(*done)
        return RunResult(
            name="sibling-cancel",
            outcome="success",
            detail=f"all tasks completed: {', '.join(results)}",
        )
    except asyncio.CancelledError:
        for task in tasks:
            task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)
        raise


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Verify asyncio cancellation and timeout handling patterns.",
    )
    parser.add_argument(
        "--work-delay",
        type=float,
        default=2.0,
        help="Simulated work duration in seconds (default: 2.0)",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=1.0,
        help="Timeout in seconds for timeout demonstrations (default: 1.0)",
    )
    parser.add_argument(
        "--cancel-after",
        type=float,
        default=0.5,
        help="Delay before manual cancellation is requested (default: 0.5)",
    )
    parser.add_argument(
        "--cleanup-delay",
        type=float,
        default=0.1,
        help="Optional cleanup delay after cancellation (default: 0.1)",
    )
    parser.add_argument(
        "--sibling-delays",
        type=str,
        default="0.4,0.8,2.5",
        help="Comma-separated task delays for sibling cancellation demo (default: 0.4,0.8,2.5)",
    )
    return parser.parse_args()


def validate_positive(name: str, value: float) -> float:
    if value < 0:
        raise ValueError(f"{name} must be non-negative")
    return value


def parse_delays(raw: str) -> list[float]:
    values: list[float] = []
    for item in raw.split(","):
        stripped = item.strip()
        if not stripped:
            continue
        try:
            delay = float(stripped)
        except ValueError as exc:
            raise ValueError(f"invalid delay value: {stripped!r}") from exc
        validate_positive("sibling delay", delay)
        values.append(delay)
    if not values:
        raise ValueError("at least one sibling delay must be provided")
    return values


async def main_async() -> int:
    args = parse_args()

    work_delay = validate_positive("--work-delay", args.work_delay)
    timeout = validate_positive("--timeout", args.timeout)
    cancel_after = validate_positive("--cancel-after", args.cancel_after)
    cleanup_delay = validate_positive("--cleanup-delay", args.cleanup_delay)

    try:
        sibling_delays = parse_delays(args.sibling_delays)
    except ValueError as exc:
        print(f"Input error: {exc}")
        return 2

    results: list[RunResult] = []
    results.append(await manual_cancel_demo(work_delay, cancel_after, cleanup_delay))
    results.append(await timeout_demo(work_delay, timeout, cleanup_delay))
    results.append(await sibling_cancellation_demo(sibling_delays, timeout))

    print("Asyncio cancellation and timeout verification results:\n")
    for result in results:
        print(f"- {result.name}: {result.outcome} — {result.detail}")

    return 0


def main() -> None:
    try:
        raise SystemExit(asyncio.run(main_async()))
    except KeyboardInterrupt:
        print("Interrupted by user")
        raise SystemExit(130)


if __name__ == "__main__":
    main()