"""Slice shared engine_*.jsonl by a [t_start_unix, t_end_unix] window.

Used between B3 policy runs and the analyzer so each policy gets
its own per-window engine_state directory.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--input-dir", type=Path, required=True)
    p.add_argument("--output-dir", type=Path, required=True)
    p.add_argument("--window", type=Path, required=True,
                   help="run_window.json with t_start_unix / t_end_unix")
    args = p.parse_args()

    window = json.loads(args.window.read_text())
    ts = float(window["t_start_unix"])
    te = float(window["t_end_unix"])

    args.output_dir.mkdir(parents=True, exist_ok=True)
    print(f"window {ts:.3f} .. {te:.3f}")
    for src in sorted(args.input_dir.glob("engine_*.jsonl")):
        n_in = n_out = 0
        dst = args.output_dir / src.name
        with src.open() as fi, dst.open("w") as fo:
            for line in fi:
                n_in += 1
                try:
                    r = json.loads(line)
                except json.JSONDecodeError:
                    continue
                t = r.get("t_unix", 0)
                if ts <= t <= te:
                    fo.write(line)
                    n_out += 1
        print(f"  {src.name}: {n_out}/{n_in}")


if __name__ == "__main__":
    main()
