"""Spawn or collect the deployed beam-dory-code app. No candidate execution locally.

MODAL_PROFILE=aside modal deploy modal_app/code_eval.py::APP --env main
MODAL_PROFILE=aside python3 scripts/spawn_code_eval.py spawn --run smoke-20261030 --configs base --limit 2 --runtime-s 360
MODAL_PROFILE=aside python3 scripts/spawn_code_eval.py spawn --run full-20261030
MODAL_PROFILE=aside python3 scripts/spawn_code_eval.py collect --run full-20261030
"""

import argparse
import json
from pathlib import Path

import modal

CONFIGS = ["base", "rE-t050-default", "rE-t050-low", "rE-t050-high", "rA-t030-default",
           "rC-t050-default", "rC-t050-low", "rC-t050-high"]
LOCAL = Path(__file__).resolve().parents[1] / "data/eval/code"


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("action", choices=["spawn", "collect", "status"])
    parser.add_argument("--run", required=True)
    parser.add_argument("--configs", default=",".join(CONFIGS))
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--runtime-s", type=int, default=2700)
    args = parser.parse_args()
    configs = args.configs.split(",")
    assert set(configs) <= set(CONFIGS)
    path = LOCAL / args.run
    path.mkdir(parents=True, exist_ok=True)
    calls_path = path / "calls.json"
    if args.action == "spawn":
        if calls_path.exists():
            raise FileExistsError(calls_path)
        fn = modal.Function.from_name("beam-dory-code", "evaluate", environment_name="main")
        calls = {}
        for config in configs:
            call = fn.spawn(config, args.run, args.limit, args.runtime_s)
            calls[config] = call.object_id
            calls_path.write_text(json.dumps(calls, indent=2) + "\n")
            print(config, call.object_id, flush=True)
        return
    calls = json.loads(calls_path.read_text())
    volume = modal.Volume.from_name("beam-dory-data", environment_name="main")
    for config, call_id in calls.items():
        if config not in configs:
            continue
        try:
            result = modal.FunctionCall.from_id(call_id).get(timeout=0)
            print(config, "generation:", result["state"], flush=True)
        except (TimeoutError, modal.exception.TimeoutError):
            print(config, "generation: running", flush=True)
        remote = f"eval/code/{args.run}/{config}"
        dest = path / config
        dest.mkdir(exist_ok=True)
        for name in ("manifest.json", "progress.json", "summary.json", "results.jsonl", "prompts.jsonl"):
            try:
                contents = b"".join(volume.read_file(f"{remote}/{name}"))
                (dest / name).write_bytes(contents)
                if name in ("progress.json", "summary.json"):
                    print(config, name, contents.decode().strip(), flush=True)
            except FileNotFoundError:
                pass
        manifest_file = dest / "manifest.json"
        if manifest_file.exists():
            manifest = json.loads(manifest_file.read_text())
            grade_id = manifest.get("grade_call_id")
            if grade_id:
                try:
                    graded = modal.FunctionCall.from_id(grade_id).get(timeout=0)
                    print(config, "grading: done", graded["benches"], flush=True)
                except (TimeoutError, modal.exception.TimeoutError):
                    print(config, "grading: running", flush=True)
        # Raw generations contain model-written code as inert JSON strings.
        # Collection deliberately does not import, compile, or execute any of it.
        if args.action == "collect":
            for name in ("generations.jsonl", "sanitized.jsonl"):
                try:
                    with (dest / name).open("wb") as handle:
                        for chunk in volume.read_file(f"{remote}/{name}"):
                            handle.write(chunk)
                except FileNotFoundError:
                    (dest / name).unlink(missing_ok=True)


if __name__ == "__main__":
    main()