#!/usr/bin/env python3
"""
Async batch processor with checkpointing, idempotency and graceful SIGINT.
Stdlib only, Python 3.11+.
"""
import asyncio
import json
import os
import signal
import time
import random

CHECKPOINT_PATH = "checkpoint.json"

class ProcessingError(Exception):
    pass

async def process(item: str) -> str:
    """Flaky external call ~80% success, 10-80ms latency."""
    await asyncio.sleep(random.uniform(0.01, 0.08))
    if random.random() < 0.2:
        raise ProcessingError(f"flaky failure for {item}")
    return f"ok:{item}"

def load_checkpoint():
    if not os.path.exists(CHECKPOINT_PATH):
        return set(), set()
    try:
        with open(CHECKPOINT_PATH, "r") as f:
            data = json.load(f)
        return set(data.get("completed", [])), set(data.get("failed", []))
    except Exception:
        return set(), set()

def save_checkpoint(completed, failed):
    data = {"completed": list(completed), "failed": list(failed)}
    tmp = CHECKPOINT_PATH + ".tmp"
    with open(tmp, "w") as f:
        json.dump(data, f)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, CHECKPOINT_PATH)   # atomic

async def handle_item(item, completed, failed, lock, sem, stop_event):
    async with sem:
        if stop_event.is_set():
            return
        backoff = 0.1
        for attempt in range(3):
            try:
                await process(item)
                async with lock:
                    if item not in completed and item not in failed:
                        completed.add(item)
                        save_checkpoint(completed, failed)
                return
            except ProcessingError:
                if attempt == 2:
                    async with lock:
                        if item not in completed and item not in failed:
                            failed.add(item)
                            save_checkpoint(completed, failed)
                    return
                await asyncio.sleep(backoff)
                backoff *= 2

async def run_batch(items):
    completed, failed = load_checkpoint()
    initially_done = len(completed) + len(failed)
    pending = [i for i in items if i not in completed and i not in failed]

    stop_event = asyncio.Event()
    lock = asyncio.Lock()
    sem = asyncio.Semaphore(8)

    loop = asyncio.get_running_loop()
    def _sigint():
        if not stop_event.is_set():
            stop_event.set()
    try:
        loop.add_signal_handler(signal.SIGINT, _sigint)
    except (NotImplementedError, ValueError):
        pass

    start = time.monotonic()
    q = asyncio.Queue()
    for it in pending:
        q.put_nowait(it)

    async def worker():
        while True:
            if stop_event.is_set() and q.empty():
                break
            try:
                item = await asyncio.wait_for(q.get(), timeout=0.1)
            except asyncio.TimeoutError:
                if q.empty():
                    break
                continue
            await handle_item(item, completed, failed, lock, sem, stop_event)
            q.task_done()

    workers = [asyncio.create_task(worker()) for _ in range(8)]
    await asyncio.gather(*workers, return_exceptions=True)

    elapsed_ms = int((time.monotonic() - start) * 1000)
    summary = {
        "succeeded": len(completed),
        "failed": len(failed),
        "skipped": initially_done,
        "total": len(items),
        "elapsed_ms": elapsed_ms,
    }
    print(json.dumps(summary))
    return summary

def main():
    items = [f"job-{i}" for i in range(100)]
    asyncio.run(run_batch(items))

# ----------------------------------------------------------------------
# Tests
async def run_tests():
    if os.path.exists(CHECKPOINT_PATH):
        os.remove(CHECKPOINT_PATH)
    random.seed(0)
    items = [f"job-{i}" for i in range(50)]

    calls = {}
    active = 0
    max_active = 0
    orig = process

    async def counting(item):
        nonlocal active, max_active
        active += 1
        if active > max_active:
            max_active = active
        calls[item] = calls.get(item, 0) + 1
        try:
            return await orig(item)
        finally:
            active -= 1

    globals()["process"] = counting

    # watcher for checkpoint validity
    errors = []
    async def watcher():
        while True:
            try:
                with open(CHECKPOINT_PATH, "r") as f:
                    json.load(f)
            except FileNotFoundError:
                pass
            except Exception as e:
                errors.append(e)
                break
            await asyncio.sleep(0.001)
    w = asyncio.create_task(watcher())

    # first partial run -> simulate interruption
    try:
        await asyncio.wait_for(run_batch(items), timeout=0.05)
    except asyncio.TimeoutError:
        pass

    # second run to finish
    await run_batch(items)
    w.cancel()
    try:
        await w
    except asyncio.CancelledError:
        pass

    completed, failed = load_checkpoint()
    assert len(completed) + len(failed) == len(items), "not all done"
    assert not errors, f"checkpoint corrupt: {errors}"
    assert max_active <= 8, f"concurrency exceeded: {max_active}"
    # idempotency: no item processed again after it was completed/failed
    # (calls count includes retries, but completed items must not be retried)
    print("tests passed")

if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "test":
        asyncio.run(run_tests())
    else:
        main()