#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

HERE = Path(__file__).resolve().parent


def run(case_id: str, variant: str) -> dict:
    completed = subprocess.run(
        [sys.executable, str(HERE / "run_case.py"), case_id, variant],
        capture_output=True,
        text=True,
    )
    return {
        "id": case_id,
        "variant": variant,
        "returncode": completed.returncode,
        "output": completed.stdout.strip(),
        "error": completed.stderr.strip(),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--start-index", type=int, default=0)
    parser.add_argument("--end-index", type=int)
    args = parser.parse_args()
    manifest = json.loads((HERE / "manifest.json").read_text())
    cases = manifest["cases"][args.start_index : args.end_index]
    jobs = [
        (case["id"], variant)
        for case in cases
        for variant in ("baseline", "experience")
    ]
    results_name = (
        "batch-results.json"
        if args.start_index == 0 and args.end_index is None
        else f"batch-results-{args.start_index}-{args.end_index or 'end'}.json"
    )
    results_path = HERE / results_name
    results = []
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = {pool.submit(run, *job): job for job in jobs}
        for future in as_completed(futures):
            result = future.result()
            results.append(result)
            print(json.dumps(result, ensure_ascii=False), flush=True)
            results_path.write_text(
                json.dumps(results, indent=2, ensure_ascii=False) + "\n"
            )
    failed = [item for item in results if item["returncode"] != 0]
    print(json.dumps({"completed": len(results), "failed": len(failed)}))
    raise SystemExit(1 if failed else 0)


if __name__ == "__main__":
    main()