"""Two missing ordinary-build checks, each in its own monitored process group.

RSS polling is not an OS hard limit. prlimit enforces output/core limits only.
No driver/context address-space limit is imposed. Never kill another job.
"""
import csv
import datetime
import io
import json
import os
from pathlib import Path
import shutil
import signal
import subprocess
import tempfile
import time

OUT = Path(__file__).resolve().parent
ROOT = OUT.parents[1]
PYTHON = "/home/frank14f/anaconda3/envs/pycuda_3_10/bin/python"
TEST = ROOT / "tests/validation/multigrid/test_refinement_followup.py"

for name, selector in (("seed1", "test_nonequilibrium_seed1_eight_steps"),
                       ("rectangle", "test_offcenter_rectangle_margin_two_one_step")):
    previous = OUT / f"gpu_{name}_results.json"
    if previous.exists() and json.loads(previous.read_text())["exit_code"] == 0:
        continue  # Resume missing cases without repeating completed flow runs.
    gpus = subprocess.check_output(["nvidia-smi", "--query-gpu=index,uuid,name,memory.used,utilization.gpu",
                                    "--format=csv"], text=True)
    apps = subprocess.check_output(["nvidia-smi", "--query-compute-apps=gpu_uuid,pid,process_name,used_memory",
                                    "--format=csv"], text=True)
    gpu = next(r for r in csv.DictReader(io.StringIO(gpus), skipinitialspace=True) if r["index"] == "0")
    before = dict(utc=datetime.datetime.now(datetime.timezone.utc).isoformat(),
                  physical_device=0, uuid=gpu["uuid"], gpus=gpus, compute_apps=apps)
    before_path = OUT / f"gpu_{name}_before.json"
    if before_path.exists():
        before_path.rename(OUT / f"gpu_{name}_before_no_launch_{time.time_ns()}.json")
    before_path.write_text(json.dumps(before, indent=2) + "\n")
    if (gpu["uuid"] in apps or int(gpu["utilization.gpu [%]"].split()[0]) != 0
            or int(gpu["memory.used [MiB]"].split()[0]) > 16):
        raise SystemExit("Physical GPU 0 occupied; no job launched")
    command = ["/usr/bin/prlimit", "--fsize=2097152", "--core=0", "--", PYTHON,
               "-B", str(TEST), "NormalBuildGPU." + selector]
    scratch = tempfile.mkdtemp(prefix="celeris-followup-gpu-")
    settings = dict(PYTHONPATH=str(ROOT / "src"), CUDA_VISIBLE_DEVICES="0", CUDA_CACHE_DISABLE="1",
                    OPENBLAS_NUM_THREADS="1", PYTHONDONTWRITEBYTECODE="1", PYTHONNOUSERSITE="1",
                    TMPDIR=scratch)
    env = dict(os.environ, **settings)
    start = time.monotonic()
    peak_group = peak_gpu = 0
    peak_members = {}
    reason = None
    process = None
    try:
        with (OUT / f"gpu_{name}_stdout.txt").open("w") as log:
            process = subprocess.Popen(command, cwd=ROOT, env=env, stdout=log,
                                       stderr=subprocess.STDOUT, start_new_session=True)
            while process.poll() is None:
                group = {}
                for proc in Path("/proc").iterdir():
                    if not proc.name.isdigit():
                        continue
                    try:
                        fields = (proc / "stat").read_text().rsplit(") ", 1)[1].split()
                        if int(fields[2]) != process.pid:
                            continue
                        status = (proc / "status").read_text()
                        rss = next(int(v.split()[1])*1024 for v in status.splitlines() if v.startswith("VmRSS:"))
                        group[int(proc.name)] = dict(rss=rss, name=(proc / "comm").read_text().strip())
                    except (OSError, StopIteration, ValueError, IndexError):
                        pass
                total = sum(v["rss"] for v in group.values())
                peak_gpu = max(peak_gpu, group.get(process.pid, {}).get("rss", 0))
                if total > peak_group:
                    peak_group, peak_members = total, group
                output_bytes = sum(p.stat().st_size for p in OUT.glob("*stdout*.txt"))
                if total > 256*1024**2 or time.monotonic()-start > 120 or output_bytes > 2*1024**2:
                    reason = "process-group RSS >256MiB, 120s or combined stdout >2MiB"
                    os.killpg(process.pid, signal.SIGKILL)
                    process.wait()
                    break
                time.sleep(.02)
        result = dict(command=command, cwd=str(ROOT), environment=settings,
                      exit_code=process.returncode, elapsed_seconds=time.monotonic()-start,
                      peak_sampled_gpu_rss_bytes=peak_gpu, peak_sampled_group_rss_bytes=peak_group,
                      peak_group_members=peak_members, termination_reason=reason,
                      monitor="20ms RSS sum polling of own process group, NOT an OS hard limit")
        (OUT / f"gpu_{name}_results.json").write_text(json.dumps(result, indent=2) + "\n")
        print((OUT / f"gpu_{name}_stdout.txt").read_text(), flush=True)
        print(json.dumps(result), flush=True)
    finally:
        if process is not None and process.poll() is None:
            os.killpg(process.pid, signal.SIGKILL)
            process.wait()
        shutil.rmtree(scratch)
    if process.returncode:
        raise SystemExit(process.returncode)