#!/usr/bin/env python3
"""AICEP Phase-2 reference publisher (proposed extension, follows encryption-scheme.md v2).

Input: a JSON list of objects {name, type, classification, body, supersedes?} that ALREADY
passed the DLP/classification gate. Output: pub/<compartment>/<batch_uuid>/ containing
<file_uuid>.age, meta.age, manifest.json (canonical) and manifest.json.sig (detached Ed25519).

Keys: AICEP_SIGN_KEY_FILE = path to an ENCRYPTED Ed25519 private key (PKCS8 PEM); passphrase from
AICEP_SIGN_KEY_PASSPHRASE or the OS keyring. Unencrypted keys are refused. --in must be on tmpfs. Recipients are PUBLIC
age strings (AICEP_RECIPIENTS, comma-separated, or --recipient). Nothing secret is printed.
"""
import argparse, json, os, shutil, sys, tempfile, time, uuid
import pyrage
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from aicep_common import (ALG, KEM, V, FileLock, Reject, atomic_write, b64, canonical, key_id, load_json,
                          load_private_key, pad, require_tmpfs, sha256_hex)

# Review fix 5: sender-side limits mirror the receiver's (aicep_sub MAX_*), so oversize fails early.
MAX_OBJECTS = 20
MAX_OBJ_PADDED = 64 * 1024
MAX_BATCH_PADDED = 512 * 1024


def load_signer(path=None, passphrase=None) -> Ed25519PrivateKey:
    k = load_private_key("AICEP_SIGN_KEY_FILE", "AICEP_SIGN_KEY_PASSPHRASE", path, passphrase)
    if not isinstance(k, Ed25519PrivateKey):
        raise SystemExit("signing key is not Ed25519")
    return k


def publish(objects, recipients, signer, sender_id, compartment, out_root, state_path,
            ttl=3600, now=None, seq=None, batch_id=None, prev_batch_id=None, enforce_limits=True):
    """enforce_limits=False exists ONLY so tests can play a hostile sender against the receiver's caps."""
    now = int(now if now is not None else time.time())
    if enforce_limits and len(objects) > MAX_OBJECTS:
        raise Reject(f"oversized batch: {len(objects)} objects > {MAX_OBJECTS}")
    with FileLock(state_path):
        return _publish(objects, recipients, signer, sender_id, compartment, out_root, state_path,
                        ttl, now, seq, batch_id, prev_batch_id, enforce_limits)


def _publish(objects, recipients, signer, sender_id, compartment, out_root, state_path,
             ttl, now, seq, batch_id, prev_batch_id, enforce_limits=True):
    pub_raw = signer.public_key().public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw)
    skid = key_id(pub_raw)
    state = load_json(state_path, default={"last_sent_seq": {}})
    if seq is None:
        seq = int(state["last_sent_seq"].get(skid, 0)) + 1
    batch_id = batch_id or str(uuid.uuid4())
    recips = [pyrage.x25519.Recipient.from_str(r) for r in recipients]

    def enc(envelope: dict):
        padded = pad(canonical(envelope))
        if enforce_limits and len(padded) > MAX_OBJ_PADDED:
            raise Reject(f"oversized object body ({len(padded)} padded bytes > {MAX_OBJ_PADDED})")
        ct = pyrage.encrypt(padded, recips)  # in memory; plaintext never touches disk
        return ct, len(padded)

    out_files, entries, meta_items = {}, [], {}
    for o in objects:
        fid = str(uuid.uuid4())
        env = {"batch_id": batch_id, "file_id": fid, "seq": seq, "type": o["type"],
               "classification": o["classification"], "body": o["body"]}
        ct, sp = enc(env)
        out_files[f"{fid}.age"] = ct
        entries.append({"id": fid, "sha256_ct": sha256_hex(ct), "size_padded": sp})
        meta_items[fid] = {"name": o.get("name"), "type": o["type"],
                           "classification": o["classification"], "supersedes": o.get("supersedes")}
    meta_env = {"batch_id": batch_id, "file_id": "meta", "seq": seq, "type": "aicep.meta",
                "classification": "meta", "body": meta_items}
    ct, sp = enc(meta_env)
    out_files["meta.age"] = ct
    entries.append({"id": "meta", "sha256_ct": sha256_hex(ct), "size_padded": sp})
    if enforce_limits and sum(e["size_padded"] for e in entries) > MAX_BATCH_PADDED:
        raise Reject(f"oversized batch > {MAX_BATCH_PADDED} padded bytes")

    manifest = {"v": 2 if prev_batch_id else V, "alg": ALG, "kem": KEM, "sender_id": sender_id, "sender_key_id": skid,
                "recipient_key_ids": sorted(recipients), "compartment": compartment,
                "batch_id": batch_id, "seq": seq, "created": now, "expires": now + int(ttl),
                "files": sorted(entries, key=lambda e: e["id"])}
    if prev_batch_id:
        manifest["prev_batch_id"] = prev_batch_id          # v2.3: signed chain link
    mbytes = canonical(manifest)
    sig = signer.sign(mbytes)

    # Build in a staging dir next to the target, then atomic rename into pub/.
    dest_parent = os.path.join(out_root, compartment)
    os.makedirs(dest_parent, exist_ok=True)
    stage = tempfile.mkdtemp(dir=dest_parent, prefix=".staging-")
    for name, data in out_files.items():
        with open(os.path.join(stage, name), "wb") as f:
            f.write(data)
    with open(os.path.join(stage, "manifest.json"), "wb") as f:
        f.write(mbytes)
    with open(os.path.join(stage, "manifest.json.sig"), "w") as f:
        f.write(b64(sig) + "\n")
    dest = os.path.join(dest_parent, batch_id)
    if os.path.exists(dest):            # v2.3: rebuilding a staged (never-completed) batch_id
        shutil.rmtree(dest)
    os.rename(stage, dest)
    state["last_sent_seq"][skid] = max(seq, int(state["last_sent_seq"].get(skid, 0)))
    atomic_write(state_path, json.dumps(state, indent=2).encode())
    return {"batch_dir": dest, "batch_id": batch_id, "seq": seq, "sender_key_id": skid,
            "files": len(objects)}


class Outbox:
    """v2.3 sender outbox. Seq is RESERVED and logged 'staged' before anything is built or uploaded;
    'published' only after the transport durably accepted it. A reserved seq is reused only for the
    same batch_id (crash recovery). Resends follow RETRY_SCHEDULE; then 'unconfirmed' alert.
    Receipts are cumulative (acked_through_seq)."""

    def __init__(self, path, signer, sender_id, ledger_path, now=time.time, folder=None):
        from aicep_common import Ledger, check_slug
        self.folder = check_slug(folder)        # v2.5: this sender's registered pub/<folder>
        self.path, self.signer, self.sender_id, self.now = path, signer, sender_id, now
        self.ledger = Ledger(ledger_path)
        pub_raw = signer.public_key().public_bytes(serialization.Encoding.Raw, serialization.PublicFormat.Raw)
        self.skid = key_id(pub_raw)

    def _load(self):
        return load_json(self.path, default={"reserved_seq": 0, "batches": {}, "last_published": None})

    def _save(self, st):
        atomic_write(self.path, json.dumps(st, indent=2, sort_keys=True).encode())

    def reserve(self):
        st = self._load()
        st["reserved_seq"] += 1
        bid = str(uuid.uuid4())
        st["batches"][bid] = {"seq": st["reserved_seq"], "status": "staged", "staged_at": int(self.now())}
        self._save(st)
        self.ledger.append("staged", batch_id=bid, seq=st["reserved_seq"], t=int(self.now()))
        return bid, st["reserved_seq"]

    def send(self, objects, recipients, compartment, adapter, work_root, ttl=3600, batch_id=None):
        st = self._load()
        if batch_id is None:
            batch_id, seq = self.reserve()
        else:                                          # resume a staged batch after a crash
            ent = st["batches"].get(batch_id)
            if not ent or ent["status"] != "staged":
                raise ValueError("only a staged batch_id can be resumed")
            seq = ent["seq"]
        st = self._load()
        prev = st.get("last_published") or "genesis"   # v2.4 (V3): chain starts at the first batch
        now = int(self.now())
        res = publish(objects, recipients, self.signer, self.sender_id, compartment, work_root,
                      os.path.join(work_root, ".pubstate.json"), ttl=ttl, now=now, seq=seq,
                      batch_id=batch_id, prev_batch_id=prev)
        adapter.publish(res["batch_dir"], {"compartment": compartment, "batch_id": batch_id, "seq": seq,
                                           "folder": self.folder})
        st = self._load()
        st["batches"][batch_id].update(status="published", published_at=now, expires=now + ttl,
                                       resends=0, batch_dir=res["batch_dir"], compartment=compartment)
        st["last_published"] = batch_id
        self._save(st)
        self.ledger.append("published", batch_id=batch_id, seq=seq, t=now)
        return res

    def staged(self):
        return [b for b, e in self._load()["batches"].items() if e["status"] == "staged"]

    def due(self, adapter=None):
        """Run the retry schedule. Returns actions taken: ('resend', bid) / ('unconfirmed', bid)."""
        from aicep_common import RETRY_SCHEDULE
        st, now, acts = self._load(), int(self.now()), []
        for bid, e in sorted(st["batches"].items(), key=lambda kv: kv[1]["seq"]):
            if e["status"] != "published":
                continue
            k = e.get("resends", 0)
            final = e["published_at"] + RETRY_SCHEDULE[-1] + 3600     # 1 h grace after last resend
            if now >= e["expires"] or (k >= len(RETRY_SCHEDULE) and now >= final):
                e["status"] = "unconfirmed"; acts.append(("unconfirmed", bid))
                self.ledger.append("unconfirmed", batch_id=bid, seq=e["seq"], t=now)
                continue
            if k < len(RETRY_SCHEDULE) and now >= e["published_at"] + RETRY_SCHEDULE[k]:
                e["resends"] = k + 1
                if adapter is not None:
                    adapter.publish(e["batch_dir"], {"compartment": e["compartment"], "batch_id": bid, "seq": e["seq"],
                                                                  "folder": self.folder})
                acts.append(("resend", bid))
                self.ledger.append("resend", batch_id=bid, seq=e["seq"], attempt=k + 1, t=now)
        self._save(st)
        return acts

    def on_receipt(self, receipt: dict, sig_b64: str, receiver_pub_raw: bytes):
        """Verify the receiver's signed cumulative receipt; mark every seq <= acked_through_seq acked."""
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
        from aicep_common import unb64
        Ed25519PublicKey.from_public_bytes(receiver_pub_raw).verify(unb64(sig_b64), canonical(receipt))
        st, n = self._load(), int(receipt.get("acked_through_seq", 0))
        for bid, e in st["batches"].items():
            if e["status"] in ("published", "unconfirmed") and e["seq"] <= n:
                e["status"] = "acked"
        self._save(st)
        self.ledger.append("receipt", acked_through_seq=n, peer_ledger_head=receipt.get("ledger_head"),
                           t=int(self.now()))
        return n


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--in", dest="inp", required=True, help="JSON list of DLP-cleared objects")
    ap.add_argument("--out", required=True, help="pub/ root (batch lands in <out>/<folder>/<batch_id>)")
    ap.add_argument("--folder", required=True, help="v2.5: this sender's registered folder slug, e.g. agent-alpha")
    ap.add_argument("--compartment", required=True)
    ap.add_argument("--sender-id", required=True)
    ap.add_argument("--recipient", action="append", default=[], help="age1... public recipient")
    ap.add_argument("--state", default=os.environ.get("AICEP_PUB_STATE", "pub_state.json"))
    ap.add_argument("--ttl", type=int, default=3600)
    a = ap.parse_args()
    from aicep_common import check_slug
    try:
        check_slug(a.folder)
    except ValueError as e:
        sys.exit(str(e))
    recips = a.recipient or [r for r in os.environ.get("AICEP_RECIPIENTS", "").split(",") if r]
    if not recips:
        sys.exit("no recipients: pass --recipient or set AICEP_RECIPIENTS")
    inp = require_tmpfs(a.inp)          # review fix 2: cleared plaintext only on tmpfs
    res = publish(load_json(inp), recips, load_signer(), a.sender_id, a.compartment,
                  os.path.join(a.out, a.folder), a.state, a.ttl)
    print(json.dumps(res))  # publish log row: ids and counts only, no secrets


if __name__ == "__main__":
    main()