#!/usr/bin/env python3
"""
Async batch processor with checkpointing, idempotency, retries, and graceful shutdown.

Usage:
    python batch_processor.py --items job-a job-b ... [--checkpoint checkpoint.json]

Or import and use the BatchProcessor class directly.
"""

from __future__ import annotations

import asyncio
import json
import os
import signal
import sys
import tempfile
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

# ---------------------------------------------------------------------------
# Public exception
# ---------------------------------------------------------------------------

class ProcessingError(Exception):
    """Raised by ``process()`` when an external call fails."""


# ---------------------------------------------------------------------------
# Simulated flaky worker (the "external call")
# ---------------------------------------------------------------------------

async def process(item: str) -> str:
    """
    Simulate a flaky external call.

    Succeeds ~80 % of the time, raises ``ProcessingError`` otherwise.
    Latency is uniform random in [10 ms, 80 ms].
    """
    import random
    await asyncio.sleep(random.uniform(0.010, 0.080))
    if random.random() < 0.20:          # 20 % failure rate
        raise ProcessingError(f"transient failure for {item!r}")
    return f"done:{item}"


# ---------------------------------------------------------------------------
# Checkpoint data model
# ---------------------------------------------------------------------------

@dataclass
class ItemResult:
    status: str          # "succeeded" | "failed"
    attempts: int = 0
    error: str = ""


@dataclass
class Checkpoint:
    """Persisted progress.  All fields are JSON-serialisable."""
    completed: dict[str, ItemResult] = field(default_factory=dict)

    # Convenience accessors
    @property
    def succeeded(self) -> list[str]:
        return [k for k, v in self.completed.items() if v.status == "succeeded"]

    @property
    def failed(self) -> list[str]:
        return [k for k, v in self.completed.items() if v.status == "failed"]

    @property
    def summary(self) -> dict[str, int]:
        return {
            "succeeded": len(self.succeeded),
            "failed": len(self.failed),
        }

    def is_done(self, item: str) -> bool:
        return item in self.completed

    def mark_succeeded(self, item: str) -> None:
        self.completed[item] = ItemResult(status="succeeded")

    def mark_failed(self, item: str, error: str = "") -> None:
        self.completed[item] = ItemResult(status="failed", error=error)


# ---------------------------------------------------------------------------
# Checkpoint I/O (atomic writes)
# ---------------------------------------------------------------------------

def load_checkpoint(path: Path) -> Checkpoint:
    """Load checkpoint from disk, or return an empty one if absent."""
    if not path.exists():
        return Checkpoint()
    raw = json.loads(path.read_text(encoding="utf-8"))
    completed: dict[str, ItemResult] = {}
    for k, v in raw.get("completed", {}).items():
        completed[k] = ItemResult(**v)
    return Checkpoint(completed=completed)


async def save_checkpoint(cp: Checkpoint, path: Path) -> None:
    """
    Atomically persist *cp* to *path*.

    Writes to a temp file in the same directory then ``os.replace()`` —
    this is atomic on every POSIX filesystem and on Windows when the dest
    is on the same volume.  A crash mid-write can never corrupt the target.
    """
    parent = path.parent
    parent.mkdir(parents=True, exist_ok=True)
    loop = asyncio.get_event_loop()

    def _write() -> None:
        fd, tmp = tempfile.mkstemp(suffix=".tmp", dir=parent)
        try:
            payload = json.dumps(asdict(cp), indent=2, ensure_ascii=False) + "\n"
            os.write(fd, payload.encode("utf-8"))
            os.close(fd)
            # Atomic rename — the only point where the target is visible.
            os.replace(tmp, str(path))
        except BaseException:
            os.close(fd) if not None else None
            try:
                os.unlink(tmp)
            except OSError:
                pass
            raise

    await loop.run_in_executor(None, _write)


# ---------------------------------------------------------------------------
# Batch processor
# ---------------------------------------------------------------------------

class BatchProcessor:
    """
    Process a list of items with bounded concurrency, retries, checkpointing,
    and graceful SIGINT handling.
    """

    def __init__(
        self,
        items: list[str],
        *,
        checkpoint_path: str | Path = "checkpoint.json",
        max_concurrency: int = 8,
        max_retries: int = 3,
        base_backoff: float = 0.1,
    ) -> None:
        self.items = list(items)
        self.checkpoint_path = Path(checkpoint_path)
        self.max_concurrency = max_concurrency
        self.max_retries = max_retries
        self.base_backoff = base_backoff

        # Load any existing checkpoint for idempotency
        self.cp = load_checkpoint(self.checkpoint_path)

        # Semaphore for bounded concurrency
        self._semaphore: asyncio.Semaphore | None = None

        # Graceful-shutdown flag — set by the signal handler
        self._shutdown_requested = asyncio.Event()

        # Track in-flight tasks so we can wait for them on shutdown
        self._in_flight: set[asyncio.Task[Any]] = set()

        # Wall-clock start time (set in ``run``)
        self._start_time: float = 0.0

    # ------------------------------------------------------------------
    # Signal handling
    # ------------------------------------------------------------------

    def _install_signal_handlers(self) -> None:
        loop = asyncio.get_event_loop()
        for sig in (signal.SIGINT, signal.SIGTERM):
            loop.add_signal_handler(sig, self._on_signal, sig)

    def _remove_signal_handlers(self) -> None:
        loop = asyncio.get_event_loop()
        for sig in (signal.SIGINT, signal.SIGTERM):
            try:
                loop.remove_signal_handler(sig)
            except ValueError:
                pass

    def _on_signal(self, sig: signal.Signals) -> None:
        # First signal: request graceful shutdown.
        # Subsequent signals during the drain window force immediate exit.
        if not self._shutdown_requested.is_set():
            print(f"\n[BatchProcessor] {sig.name} received — draining in-flight tasks …", file=sys.stderr)
            self._shutdown_requested.set()
        else:
            print(f"\n[BatchProcessor] {sig.name} received again — forcing exit.", file=sys.stderr)
            sys.exit(1)

    # ------------------------------------------------------------------
    # Core processing
    # ------------------------------------------------------------------

    async def _process_one(self, item: str) -> None:
        """Retry *item* up to max_retries with exponential backoff."""
        attempts = 0
        last_error = ""
        while attempts < self.max_retries:
            attempts += 1
            try:
                result = await process(item)
                self.cp.mark_succeeded(item)
                await save_checkpoint(self.cp, self.checkpoint_path)
                print(f"  ✓ {item} (attempt {attempts})")
                return
            except ProcessingError as exc:
                last_error = str(exc)
                if attempts < self.max_retries:
                    backoff = self.base_backoff * (2 ** (attempts - 1))
                    print(f"  ✗ {item} attempt {attempts}/{self.max_retries} failed — retry in {backoff:.2f}s", file=sys.stderr)
                    try:
                        await asyncio.sleep(backoff)
                    except asyncio.CancelledError:
                        # We were cancelled mid-backoff — do NOT mark failed.
                        raise

        # Exhausted retries → terminal failure
        self.cp.mark_failed(item, error=last_error)
        await save_checkpoint(self.cp, self.checkpoint_path)
        print(f"  ☠ {item} failed after {attempts} attempts", file=sys.stderr)

    async def _run_item(self, item: str) -> None:
        """Wrap ``_process_one`` with concurrency tracking and shutdown guard."""
        task = asyncio.current_task()
        assert task is not None
        self._in_flight.add(task)
        try:
            # Check shutdown before starting work on this item.
            if self._shutdown_requested.is_set():
                return
            await self._process_one(item)
        finally:
            self._in_flight.discard(task)

    # ------------------------------------------------------------------
    # Public entry point
    # ------------------------------------------------------------------

    async def run(self) -> dict[str, Any]:
        """
        Run the batch.  Returns a summary dict suitable for JSON output.
        """
        self._start_time = time.monotonic()
        self._install_signal_handlers()

        # Items already in the checkpoint are skipped.
        pending = [it for it in self.items if not self.cp.is_done(it)]
        skipped = len(self.items) - len(pending)

        self._semaphore = asyncio.Semaphore(self.max_concurrency)

        try:
            async def _bounded(item: str) -> None:
                async with self._semaphore:
                    await self._run_item(item)

            workers = [asyncio.create_task(_bounded(it)) for it in pending]

            # Wait until either all workers finish or shutdown is requested.
            if workers:
                done, pending_tasks = await asyncio.wait(
                    workers,
                    return_when=asyncio.FIRST_COMPLETED,
                )
                # Check if we should drain or abort.
                if self._shutdown_requested.is_set():
                    # Cancel any still-pending workers (they won't start new items).
                    for t in pending_tasks:
                        t.cancel()
                    # Wait for all to finish/cancel.
                    if pending_tasks:
                        await asyncio.gather(*pending_tasks, return_exceptions=True)
                else:
                    # Some finished; continue with the rest.
                    if pending_tasks:
                        done2, _ = await asyncio.wait(pending_tasks)
                        await asyncio.gather(*done2, return_exceptions=True)

            # Final flush (idempotent — same content).
            await save_checkpoint(self.cp, self.checkpoint_path)

        finally:
            self._remove_signal_handlers()

        elapsed_ms = int((time.monotonic() - self._start_time) * 1000)
        summary = {
            "succeeded": len(self.cp.succeeded),
            "failed": len(self.cp.failed),
            "skipped": skipped,
            "total": len(self.items),
            "elapsed_ms": elapsed_ms,
        }
        return summary


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------

def main(argv: list[str] | None = None) -> int:
    import argparse

    parser = argparse.ArgumentParser(description="Async batch processor with checkpointing")
    parser.add_argument("items", nargs="+", help="Item identifiers to process")
    parser.add_argument("--checkpoint", default="checkpoint.json", help="Checkpoint file path")
    parser.add_argument("--concurrency", type=int, default=8, help="Max concurrent items")
    parser.add_argument("--retries", type=int, default=3, help="Max retries per item")
    args = parser.parse_args(argv)

    processor = BatchProcessor(
        args.items,
        checkpoint_path=args.checkpoint,
        max_concurrency=args.concurrency,
        max_retries=args.retries,
    )

    summary = asyncio.run(processor.run())
    print(json.dumps(summary))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())