"""Bounded parallel execution of independent real KNP runs, with separate outputs."""
from concurrent.futures import ThreadPoolExecutor, as_completed
import argparse
import json
import os
from pathlib import Path
import re
import subprocess
import sys
import time

ap = argparse.ArgumentParser()
ap.add_argument('plan', type=Path)
ap.add_argument('--jobs', type=int, default=min(4, os.cpu_count() or 1))
ap.add_argument('--summary', type=Path, required=True)
a = ap.parse_args()
if not 1 <= a.jobs <= (os.cpu_count() or 1):
    ap.error('--jobs must be between 1 and available logical CPUs')
plan = json.loads(a.plan.read_text())
jobs = plan['jobs']
names = [job['name'] for job in jobs]
if len(set(names)) != len(names) or not all(re.fullmatch(r'[A-Za-z0-9_-]+', name) for name in names):
    ap.error('Run names must be unique and contain only letters, digits, underscores or hyphens')
root = Path('/home/int/run/cats-dogs')
if any((root / name).exists() for name in names):
    ap.error('A planned run directory already exists; choose new names')


def execute(job):
    if not isinstance(job['args'], list) or not all(isinstance(x, str) for x in job['args']):
        raise ValueError('args must be a string list')
    result = subprocess.run([sys.executable, str(Path(__file__).with_name('run.py')),
                             job['name'], *job['args']], capture_output=True, text=True)
    row = {'name': job['name'], 'returncode': result.returncode}
    if result.returncode:
        row['error'] = result.stderr[-3000:]
    else:
        row['metrics'] = json.loads((root / job['name'] / 'metrics.json').read_text())
    return row


started = time.monotonic()
rows = []
with ThreadPoolExecutor(max_workers=a.jobs) as pool:
    for future in as_completed([pool.submit(execute, job) for job in jobs]):
        row = future.result()
        rows.append(row)
        print(json.dumps(row), flush=True)
summary = {'jobs': a.jobs, 'logical_cpus': os.cpu_count(), 'runs': rows,
           'elapsed_seconds': round(time.monotonic() - started, 3),
           'plan': str(a.plan), 'all_succeeded': all(row['returncode'] == 0 for row in rows)}
a.summary.write_text(json.dumps(summary, indent=2) + '\n')
print('BATCH', json.dumps(summary), flush=True)
sys.exit(0 if summary['all_succeeded'] else 1)