#!/usr/bin/env python3
"""Quarantine contaminated memory entries and emit a minimal clean-state bundle.

Nothing is ever deleted: every original line ends up byte-exact in either
clean_store.jsonl or quarantine.jsonl. The clean store is exactly the verified
checkpoint prefix; later entries are quarantined (contaminated ones as
"contaminated", the rest as "after_checkpoint_unverified" for manual review).

Exit codes: 0 done, 1 verification failed / refused, 2 usage or output exists.
Stdlib only, no network, no subprocess. Writes only into --out-dir.
"""
import argparse
import hashlib
import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import verify_checkpoint as vc  # noqa: E402

OUT_NAMES = ("clean_store.jsonl", "quarantine.jsonl", "clean_bundle.json")


class Refused(Exception):
    pass


def detect_flagged(entries):
    return [e["id"] for e in entries if e.get("flags")]


def bundle_hash(bundle):
    body = {k: v for k, v in bundle.items() if k != "bundle_sha256"}
    return hashlib.sha256(json.dumps(body, sort_keys=True, separators=(",", ":")).encode()).hexdigest()


def plan(store_bytes, checkpoint, checkpoint_file_sha, contaminated_ids=None):
    """Pure function: returns (clean_bytes, quarantine_bytes, bundle dict)."""
    problems = vc.verify(store_bytes, checkpoint)
    if problems:
        raise Refused("store failed verification: " + "; ".join(problems))
    lines = vc.split_lines(store_bytes)
    entries = [json.loads(l) for l in lines]
    n = checkpoint["entry_count"]
    ids = list(contaminated_ids) if contaminated_ids is not None else detect_flagged(entries)
    known = {e["id"] for e in entries}
    unknown = [i for i in ids if i not in known]
    if unknown:
        raise Refused("unknown ids: " + ", ".join(unknown))
    prefix_ids = {e["id"] for e in entries[:n]}
    inside = [i for i in ids if i in prefix_ids]
    if inside:
        raise Refused("ids inside the immutable checkpoint prefix, refusing to touch it: " + ", ".join(inside))
    bad = set(ids)
    clean = b"".join(lines[:n])
    q_lines = []
    retained = [e["id"] for e in entries[:n]]
    quarantined, contaminated = [], []
    for pos in range(n, len(lines)):
        e = entries[pos]
        reason = "contaminated" if e["id"] in bad else "after_checkpoint_unverified"
        rec = {"id": e["id"], "original_position": pos + 1, "reason": reason,
               "original_line": lines[pos].decode("utf-8")}
        q_lines.append(json.dumps(rec, sort_keys=True, ensure_ascii=False) + "\n")
        quarantined.append(e["id"])
        if reason == "contaminated":
            contaminated.append(e["id"])
    bundle = {
        "schema": "bilateral-reset-bundle/1",
        "checkpoint_ref": {"file_sha256": checkpoint_file_sha, "head_hash": checkpoint["head_hash"],
                           "entry_count": n, "prefix_sha256": checkpoint["prefix_sha256"]},
        "retained_ids": retained,
        "quarantined_ids": quarantined,
        "contaminated_ids": contaminated,
        "reverify": "python3 verifier/verify_checkpoint.py clean_store.jsonl <checkpoint.json> must exit 0; "
                    "then recompute bundle_sha256 over this JSON without that key (sorted keys, compact separators).",
    }
    bundle["bundle_sha256"] = bundle_hash(bundle)
    return clean, "".join(q_lines).encode("utf-8"), bundle


def main(argv):
    ap = argparse.ArgumentParser(prog="quarantine.py", description=__doc__.splitlines()[0])
    ap.add_argument("store")
    ap.add_argument("checkpoint")
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--ids", help="comma-separated contaminated ids (default: entries with non-empty flags)")
    ap.add_argument("--force", action="store_true", help="overwrite existing outputs")
    try:
        a = ap.parse_args(argv)
    except SystemExit as e:
        return 2 if e.code else 0
    try:
        with open(a.store, "rb") as f:
            data = f.read()
        with open(a.checkpoint, "rb") as f:
            cp_bytes = f.read()
        cp = vc.load_checkpoint(a.checkpoint)
    except (OSError, ValueError) as e:
        sys.stderr.write("error: %s\n" % e)
        return 2
    outs = [os.path.join(a.out_dir, n) for n in OUT_NAMES]
    cp_real = os.path.realpath(a.checkpoint)
    store_real = os.path.realpath(a.store)
    for o in outs:
        if os.path.realpath(o) in (cp_real, store_real):
            sys.stderr.write("refused: output would overwrite the checkpoint or the input store\n")
            return 1
    existing = [o for o in outs if os.path.exists(o)]
    if existing and not a.force:
        sys.stderr.write("refused: outputs exist (use --force): " + ", ".join(existing) + "\n")
        return 2
    ids = [i for i in a.ids.split(",") if i] if a.ids is not None else None
    try:
        clean, quar, bundle = plan(data, cp, hashlib.sha256(cp_bytes).hexdigest(), ids)
    except Refused as e:
        sys.stderr.write("refused: %s\n" % e)
        return 1
    os.makedirs(a.out_dir, exist_ok=True)
    for path, content in zip(outs, (clean, quar, (json.dumps(bundle, indent=2, sort_keys=True) + "\n").encode())):
        with open(path, "wb") as f:
            f.write(content)
    print("clean=%d quarantined=%d bundle_sha256=%s" % (len(bundle["retained_ids"]), len(bundle["quarantined_ids"]), bundle["bundle_sha256"]))
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))