"""Force-kill a worker immediately after a committed checkpoint; verify exact recovery."""
from pathlib import Path
import argparse
import json
import os
import subprocess
import sys
import tempfile
import time
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--worker")
    args = parser.parse_args()
    import torch
    from safetensors.torch import load_file
    from nexora import training
    from nexora.data import prepare
    if args.worker:
        folder = Path(args.worker)
        save = training.save_checkpoint
        def checkpoint(*a, **kw):
            receipt = save(*a, **kw)
            if receipt["step"] == 3:
                (folder / "ready").write_text("committed")
                time.sleep(120)
            return receipt
        training.save_checkpoint = checkpoint
        training.train(folder / "config.json", folder / "data", folder / "interrupted")
        return
    with tempfile.TemporaryDirectory(prefix="nexora-recovery-") as d:
        folder = Path(d)
        cfg = {"model": {"hidden_size": 32, "layers": 1, "heads": 4, "kv_heads": 2, "intermediate_size": 64, "max_context": 32},
               "training": {"steps": 9, "batch_size": 2, "sequence_length": 16, "learning_rate": .001, "seed": 73, "eval_every": 3, "checkpoint_every": 3, "device": "cpu", "threads": 2}}
        (folder / "config.json").write_text(json.dumps(cfg))
        prepare([{"id": "a", "text": "Training for recovery uses a small deterministic sequence of original engineering text.", "source": "original", "license": "owner-authored", "domain": "text"},
                 {"id": "b", "text": "Separate validation helps detect changes in the sequence of restored parameter updates.", "source": "original", "license": "owner-authored", "domain": "text", "split": "validation"}], folder / "data")
        training.train(folder / "config.json", folder / "data", folder / "full")
        with (folder / "worker.log").open("w") as log:
            worker = subprocess.Popen([sys.executable, str(Path(__file__).resolve()), "--worker", d], stdout=log, stderr=subprocess.STDOUT)
            deadline = time.monotonic()+90
            try:
                while not (folder / "ready").exists():
                    if worker.poll() is not None or time.monotonic() > deadline:
                        raise RuntimeError("Worker did not produce committed checkpoint")
                    time.sleep(.1)
                worker.kill()
                code = worker.wait(timeout=10)
            finally:
                if worker.poll() is None:
                    worker.kill()
                    worker.wait(timeout=10)
        training.train(folder / "config.json", folder / "data", folder / "interrupted", resume=True)
        a = load_file(str(folder / "full/model.safetensors"))
        b = load_file(str(folder / "interrupted/model.safetensors"))
        exact = all(torch.equal(a[k], b[k]) for k in a)
        if not exact:
            raise AssertionError("Forced-kill recovery diverged")
        report = {"status": "VALIDATED", "worker_forcibly_killed": True, "worker_exit_code": code,
                  "checkpoint_step": 3, "final_step": 9, "all_parameters_bitwise_equal": exact,
                  "limitations": "Single-process CPU recovery after checkpoint commit; not distributed kill recovery or mid-write power-loss durability"}
        Path("reports/recovery.json").write_text(json.dumps(report, indent=2))
        print(json.dumps(report))


if __name__ == "__main__":
    main()