"""Sample fresh native builds, including parsing, expansion, optimization and JIT.

Run with the same Python environment as benchmark-native-output.py. The sampler
runs on a separate core; no generation timeout or memory cutoff is imposed.
Generation-only peaks are also recorded to match the historical F13 measurements,
which begin with the expression already loaded. Nested-expression comparisons
use the full build window, including parsing and expansion.
"""
import argparse
import importlib.util
import json
import os
from pathlib import Path
import subprocess
import sys
import time

HERE = Path(__file__).resolve().parent
spec = importlib.util.spec_from_file_location("memory_helper", HERE / "f13-output.py")
helper = importlib.util.module_from_spec(spec)
spec.loader.exec_module(helper)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--provenance", type=Path, required=True)
    parser.add_argument("--cpu", type=int, default=10)
    parser.add_argument("--sampler-cpu", type=int, default=11)
    args = parser.parse_args()
    args.output = args.output.resolve()
    args.output.mkdir(parents=True, exist_ok=False)
    os.sched_setaffinity(0, {args.cpu, args.sampler_cpu})
    results = []
    for run in range(1, 4):
        for dataset, input_form in [("f13", "expanded"), ("nested", "original"), ("nested", "expanded")]:
            key = f"{dataset}-{input_form}"
            work = args.output / f"{key}-{run}"
            command = [sys.executable, "-I", str(HERE / "benchmark-native-output.py"),
                       "--dataset", dataset, "--input-form", input_form,
                       "--output", str(work), "--provenance", str(args.provenance.resolve()),
                       "--cpu", str(args.cpu), "--build-only", "--profile-build", "--label", "Symbolica"]
            os.sched_setaffinity(0, {args.cpu, args.sampler_cpu})
            with (args.output / f"{key}-{run}.log").open("w") as log:
                process = subprocess.Popen(command, stdout=log, stderr=log)
                os.sched_setaffinity(0, {args.sampler_cpu})
                while not (work / "build.ready").exists():
                    assert process.poll() is None, f"Worker failed: {key}-{run}"
                    time.sleep(.01)
                baseline = sum(rss for _, rss in helper.process_tree(process.pid))
                full_peak = baseline
                generation_peak = None
                generation_baseline = None
                samples = generation_samples = 0
                (work / "build.start").touch()
                while True:
                    rss = sum(value for _, value in helper.process_tree(process.pid))
                    full_peak = max(full_peak, rss)
                    samples += 1
                    if generation_peak is None and (work / "generation.ready").exists():
                        generation_peak = generation_baseline = rss
                        (work / "generation.start").touch()
                    if generation_peak is not None:
                        generation_peak = max(generation_peak, rss)
                        generation_samples += 1
                    if (work / "build.done").exists():
                        (work / "build.ack").touch()
                        break
                    assert process.poll() is None, f"Worker failed: {key}-{run}"
                    time.sleep(.01)
                assert process.wait() == 0, f"Worker validation failed: {key}-{run}"
            record = json.loads((work / "build.json").read_text())
            assert record["candidate"]["verification"]["verified"]
            assert record["candidate"]["serialization_verification"]["verified"]
            record["memory"] = {
                "baseline_rss_bytes": baseline, "full_build_peak_tree_rss_bytes": full_peak,
                "generation_baseline_rss_bytes": generation_baseline,
                "generation_peak_tree_rss_bytes": generation_peak,
                "sample_count": samples, "generation_sample_count": generation_samples,
                "sample_interval_seconds": .01, "sampler_cpu": args.sampler_cpu,
                "includes": "worker and all live descendants; profiler excluded",
                "full_build_window": "before parsing through optimization, JIT and first evaluation",
                "generation_window": "after parsing and explicit expansion through optimization, JIT and first evaluation",
            }
            record.update(case=key, run=run, command=command)
            results.append(record)
            (args.output / "results.json").write_text(json.dumps(
                {"completed": len(results) == 9, "results": results}, indent=2) + "\n")
            print(f"{key}-{run}: generation {generation_peak / 2**20:.1f} MiB; "
                  f"full build {full_peak / 2**20:.1f} MiB; validated", flush=True)


if __name__ == "__main__":
    main()