import asyncio
import time
import json
import os
import statistics
from dataclasses import dataclass, asdict
from typing import Optional, Any, Callable, List, Dict

from llama_index.core.utils import aretry_on_exceptions_with_backoff, ErrorToRetry

class TransientError(Exception):
    pass

@dataclass
class Scenario:
    name: str
    concurrency: int
    backoff: float
    work_delay: float
    fail_times: Optional[int]
    mixed_ratio: float
    repeats: int

def make_task(task_id: int, work_delay: float, fail_times: Optional[int]) -> Callable[[], Any]:
    attempts = {"n": 0}
    async def fn():
        attempts["n"] += 1
        await asyncio.sleep(work_delay)
        if fail_times is None or fail_times == 0:
            return f"ok-{task_id}"
        if fail_times == -1:
            raise TransientError(f"task {task_id} always fails")
        if attempts["n"] <= fail_times:
            raise TransientError(f"task {task_id} transient failure")
        return f"ok-{task_id}"
    return fn

async def run_once(s: Scenario) -> float:
    start = time.perf_counter()
    async def one(i: int):
        use_fail = s.fail_times
        if s.mixed_ratio < 1.0:
            cutoff = int(s.concurrency * s.mixed_ratio)
            if i >= cutoff:
                use_fail = 0
        fn = make_task(i, s.work_delay, use_fail)
        if use_fail == -1:
            max_tries = 3
        elif use_fail is None or use_fail == 0:
            max_tries = 1
        else:
            max_tries = use_fail + 1
        return await aretry_on_exceptions_with_backoff(
            fn,
            errors_to_retry=[ErrorToRetry(TransientError)],
            max_tries=max_tries,
            min_backoff_secs=s.backoff,
            max_backoff_secs=s.backoff,
        )
    try:
        await asyncio.gather(*(one(i) for i in range(s.concurrency)))
    except Exception:
        pass
    return time.perf_counter() - start

async def run_scenario(s: Scenario) -> Dict[str, Any]:
    times = []
    for _ in range(s.repeats):
        times.append(await run_once(s))
    return {
        "scenario": asdict(s),
        "n": len(times),
        "mean_s": statistics.mean(times),
        "median_s": statistics.median(times),
        "min_s": min(times),
        "max_s": max(times),
        "samples_s": times,
    }

def scenarios() -> List[Scenario]:
    base = [
        ("always_ok", None, 1.0),
        ("fail_once", 1, 1.0),
        ("fail_twice", 2, 1.0),
        ("always_fail", -1, 1.0),
        ("mixed_50pct_fail_once", 1, 0.5),
    ]
    concurrencies = [10, 50]
    backoffs = [0.01, 0.2]
    work_delay = 0.01
    reps = 3
    out = []
    for c in concurrencies:
        for b in backoffs:
            for (name, fail_times, mixed_ratio) in base:
                out.append(
                    Scenario(
                        name=f"{name}",
                        concurrency=c,
                        backoff=b,
                        work_delay=work_delay,
                        fail_times=fail_times,
                        mixed_ratio=mixed_ratio,
                        repeats=reps,
                    )
                )
    return out

def save_json(path: str, payload: Any) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2)

def load_json(path: str) -> Any:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)

def key_of(r: Dict[str, Any]) -> str:
    s = r["scenario"]
    return f'{s["name"]}|c={s["concurrency"]}|b={s["backoff"]}|d={s["work_delay"]}|f={s["fail_times"]}|m={s["mixed_ratio"]}'

def fmt(x: float) -> str:
    return f"{x:.3f}s"

def compare(before_path: str, after_path: str) -> None:
    before = load_json(before_path)["results"]
    after = load_json(after_path)["results"]
    bmap = {key_of(r): r for r in before}
    amap = {key_of(r): r for r in after}
    keys = sorted(set(bmap.keys()) & set(amap.keys()))
    header = ["scenario", "before(median)", "after(median)", "speedup", "before(mean)", "after(mean)"]
    rows = []
    for k in keys:
        br = bmap[k]
        ar = amap[k]
        bmed = br["median_s"]
        amed = ar["median_s"]
        sp = (bmed / amed) if amed > 0 else float("inf")
        rows.append([k, fmt(bmed), fmt(amed), f"{sp:.2f}x", fmt(br["mean_s"]), fmt(ar["mean_s"])])
    widths = [max(len(str(row[i])) for row in ([header] + rows)) for i in range(len(header))]
    def print_row(row):
        print(" | ".join(str(row[i]).ljust(widths[i]) for i in range(len(row))))
    print_row(header)
    print("-+-".join("-" * w for w in widths))
    for r in rows:
        print_row(r)

async def run_all(out_path: str) -> None:
    res = []
    for s in scenarios():
        res.append(await run_scenario(s))
    save_json(out_path, {"results": res})

def main():
    import sys
    if len(sys.argv) < 2 or sys.argv[1] not in {"before", "after", "compare"}:
        print("Usage: python bench_async_retry.py before|after|compare")
        raise SystemExit(2)
    mode = sys.argv[1]
    before_path = os.path.join("bench_out", "before.json")
    after_path = os.path.join("bench_out", "after.json")
    if mode == "before":
        asyncio.run(run_all(before_path))
        print(f"Wrote {before_path}")
    elif mode == "after":
        asyncio.run(run_all(after_path))
        print(f"Wrote {after_path}")
    else:
        compare(before_path, after_path)

if __name__ == "__main__":
    main()