"""Shared helpers for the AICEP Phase-2 reference scripts (proposed extension, v2).
Keys are only ever loaded from file paths / env vars and are never printed."""
import base64, hashlib, json, os, tempfile, time

V = 1
ALG = "age-x25519+ed25519"
KEM = "x25519"
MIN_BUCKET = 4096


class Reject(Exception):
    """Fail-closed rejection. Message never contains key material."""


def canonical(obj) -> bytes:
    # RFC 8785-compatible for our schema: only str/int/bool/null/list/dict (no floats).
    def chk(o):
        if isinstance(o, float):
            raise ValueError("floats not allowed in canonical manifest")
        if isinstance(o, dict):
            [chk(v) for v in o.values()]
        elif isinstance(o, list):
            [chk(v) for v in o]
    chk(obj)
    return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode("utf-8")


def sha256_hex(b: bytes) -> str:
    return hashlib.sha256(b).hexdigest()


def key_id(pub_raw: bytes) -> str:
    return "ed25519:" + hashlib.sha256(pub_raw).hexdigest()[:16]


def pad(data: bytes) -> bytes:
    need = len(data) + 4
    size = MIN_BUCKET
    while size < need:
        size *= 2
    return len(data).to_bytes(4, "big") + data + b"\0" * (size - need)


def unpad(buf: bytes) -> bytes:
    if len(buf) < 4:
        raise Reject("padding: too short")
    n = int.from_bytes(buf[:4], "big")
    if n > len(buf) - 4 or any(buf[4 + n:]):
        raise Reject("padding: malformed")
    return buf[4:4 + n]


def read_secret_path(env_name: str, override: str | None = None) -> bytes:
    path = override or os.environ.get(env_name)
    if not path:
        raise SystemExit(f"missing key path: set {env_name} (a file path, never the key value)")
    with open(path, "rb") as f:
        return f.read()


TMPFS_TYPES = {"tmpfs", "ramfs"}


def _mount_fstype(path: str) -> str:
    """Filesystem type of the mount holding path (longest mountpoint prefix in /proc/self/mounts)."""
    rp = os.path.realpath(path)
    best, fs = "", ""
    try:
        with open("/proc/self/mounts") as f:
            for line in f:
                parts = line.split()
                if len(parts) < 3:
                    continue
                mp = parts[1].replace("\\040", " ")
                if (rp == mp or rp.startswith(mp.rstrip("/") + "/")) and len(mp) >= len(best):
                    best, fs = mp, parts[2]
    except OSError:
        return ""
    return fs


def require_tmpfs(path: str, what: str = "cleared plaintext input"):
    """Review fix 2: DLP-cleared plaintext may only live on tmpfs (e.g. /dev/shm). Refuse any other path,
    including symlinks that resolve off tmpfs."""
    fs = _mount_fstype(path)
    if fs not in TMPFS_TYPES:
        raise SystemExit(f"refusing {what} at {path!r}: not on tmpfs (got {fs or 'unknown'}); use /dev/shm")
    return os.path.realpath(path)


def load_private_key(file_env: str, pass_env: str, path=None, passphrase=None, keyring_service="aicep"):
    """Review fix 1: private keys must be ENCRYPTED PEM (PKCS8 + BestAvailableEncryption).
    Passphrase source, in order: explicit arg, env var <pass_env>, OS keyring (service 'aicep',
    username = <file_env>) if the optional 'keyring' package is installed. Unencrypted PEM is refused."""
    from cryptography.hazmat.primitives import serialization
    data = read_secret_path(file_env, path)
    if b"ENCRYPTED PRIVATE KEY" not in data:
        raise SystemExit(f"refusing unencrypted private key in {file_env}: use an encrypted PKCS8 PEM")
    pw = passphrase if passphrase is not None else os.environ.get(pass_env)
    if pw is None:
        try:
            import keyring  # optional
            pw = keyring.get_password(keyring_service, file_env)
        except Exception:
            pw = None
    if not pw:
        raise SystemExit(f"no passphrase for {file_env}: set {pass_env} or store it in the OS keyring")
    pw = pw.encode() if isinstance(pw, str) else pw
    try:
        return serialization.load_pem_private_key(data, password=pw)
    except (ValueError, TypeError):
        raise SystemExit(f"cannot decrypt {file_env} (wrong passphrase or corrupt key)")


class FileLock:
    """Review fix 5: exclusive fcntl lock on <path>.lock for load->check->write critical sections."""

    def __init__(self, path):
        self.path = os.path.abspath(path) + ".lock"

    def __enter__(self):
        os.makedirs(os.path.dirname(self.path), exist_ok=True)
        import fcntl
        self.f = open(self.path, "a+")
        fcntl.flock(self.f, fcntl.LOCK_EX)
        return self

    def __exit__(self, *exc):
        import fcntl
        fcntl.flock(self.f, fcntl.LOCK_UN)
        self.f.close()


def load_json(path, default=None):
    if default is not None and not os.path.exists(path):
        return default
    with open(path, "rb") as f:
        return json.loads(f.read())


def atomic_write(path: str, data: bytes):
    d = os.path.dirname(os.path.abspath(path))
    os.makedirs(d, exist_ok=True)
    fd, tmp = tempfile.mkstemp(dir=d)
    with os.fdopen(fd, "wb") as f:
        f.write(data)
    os.replace(tmp, path)


def b64(b: bytes) -> str:
    return base64.b64encode(b).decode()


def unb64(s: str) -> bytes:
    return base64.b64decode(s, validate=True)


# ---------------- v2.3 delivery guarantees (proposed Phase 2 extension) ----------------
SUPPORTED_V = (1, 2)          # v2 = manifest may carry signed prev_batch_id
# QoP floor: the ONLY accepted quality of protection is signed + encrypted. Unsigned, unencrypted,
# or sign-only / encrypt-only batches are never ingested; they go to dlq with reason "qop".
QOP_FLOOR = "sign+encrypt"
REASONS_CRYPTO = {"folder_mismatch", "sig_fail", "hash_mismatch", "revoked", "equivocation", "qop", "replay"}
REASONS_TRANSIENT = {"ingest_error", "io_error"}
REASONS_OTHER = {"expired", "gap_timeout", "size_limit", "rejected"}
MAX_ATTEMPTS_TRANSIENT, MAX_ATTEMPTS_CRYPTO = 3, 1
RETRY_SCHEDULE = (3600, 4 * 3600, 12 * 3600)   # sender resend offsets after publish (no ack yet)
HOLD_WINDOW = 1800                               # receiver reorder window before human hold


def classify_reason(msg: str) -> str:
    m = msg.lower()
    for key, code in (("signature", "sig_fail"), ("hash", "hash_mismatch"), ("revok", "revoked"),
                      ("replay", "replay"), ("expired", "expired"), ("oversized", "size_limit"),
                      ("equivocation", "equivocation"), ("unsigned", "qop"), ("not canonical", "sig_fail"),
                      ("unlisted", "hash_mismatch"), ("binding", "hash_mismatch")):
        if key in m:
            return code
    return "rejected"


class LedgerError(Exception):
    """Ledger integrity failure: fail closed."""


MAX_TTL = 7 * 86400              # W1: maximum signed validity window (expires - created)


class Ledger:
    """Append-only, hash-chained audit log (JSON lines). Each entry carries prev = hash of previous entry.
    head() goes into every receipt so both sides can cross-check history. Never contains content or keys."""
    GENESIS = "0" * 64

    def __init__(self, path):
        self.path = path

    @property
    def anchor_path(self):
        return self.path + ".anchor"

    def entries(self):
        if not os.path.exists(self.path):
            return []
        with open(self.path, "rb") as f:
            return [json.loads(l) for l in f if l.strip()]

    def anchor(self):
        """v2.4 (V6): separately written {head, count}; detects tail truncation."""
        return load_json(self.anchor_path, default={"head": self.GENESIS, "count": 0})

    def head(self):
        e = self.entries()
        return e[-1]["hash"] if e else self.GENESIS

    def _anchor_ok(self):
        """W2/W3: the anchor must match what is on disk; no anchor is only OK for an empty ledger."""
        es = self.entries()
        if not os.path.exists(self.anchor_path):
            return not es
        a = self.anchor()
        return a.get("count") == len(es) and a.get("head") == (es[-1]["hash"] if es else self.GENESIS)

    def append(self, event: str, **fields):
        try:
            ok = self._anchor_ok()
        except Exception:
            ok = False
        if not ok:
            raise LedgerError(f"ledger anchor mismatch on {os.path.basename(self.path)}: refusing to append "
                              "(possible truncation); a human must run --reanchor-ledger --operator NAME")
        return self._append(event, **fields)

    def _append(self, event: str, **fields):
        prev = self.head()
        body = {"event": event, "prev": prev, "t": int(fields.pop("t", time.time())), **fields}
        body["hash"] = sha256_hex(canonical({k: v for k, v in body.items() if k != "hash"}))
        os.makedirs(os.path.dirname(os.path.abspath(self.path)), exist_ok=True)
        with open(self.path, "ab") as f:
            f.write(canonical(body) + b"\n")
            f.flush(); os.fsync(f.fileno())
        n = len(self.entries())
        atomic_write(self.anchor_path, canonical({"head": body["hash"], "count": n}))
        return body["hash"]

    def links_ok(self) -> bool:
        try:
            prev = self.GENESIS
            for e in self.entries():
                if e.get("prev") != prev or e.get("hash") != sha256_hex(canonical({k: v for k, v in e.items() if k != "hash"})):
                    return False
                prev = e["hash"]
            return True
        except Exception:
            return False

    def reanchor(self, operator: str, t=None):
        """W3: HUMAN-ONLY re-anchor (legacy ledger, or crash between line append and anchor write).
        Internal links must still verify; the event is logged as reanchored{operator, old_count, new_count}."""
        operator = operator.strip() if isinstance(operator, str) else ""
        if not operator:
            raise PermissionError("reanchor requires a named human operator")
        if not self.links_ok():
            raise LedgerError("ledger links do not verify; refusing to re-anchor")
        old = self.anchor() if os.path.exists(self.anchor_path) else {"head": None, "count": None}
        n = len(self.entries())
        return self._append("reanchored", operator=operator, old_count=old.get("count"),
                            old_head=old.get("head"), new_count=n + 1,
                            **({"t": t} if t is not None else {}))

    def verify(self, expected_head=None) -> bool:
        """Internal links + head anchor (V6). expected_head: e.g. ledger_head from the last signed receipt."""
        try:
            es = self.entries()
            prev = self.GENESIS
            for e in es:
                h = e.get("hash")
                if e.get("prev") != prev or h != sha256_hex(canonical({k: v for k, v in e.items() if k != "hash"})):
                    return False
                prev = h
            a = self.anchor()
        except Exception:
            return False                      # malformed line / anchor = not verified
        if a.get("head") != prev or a.get("count") != len(es):
            return False
        if expected_head is not None and expected_head not in {e["hash"] for e in es} | {self.GENESIS}:
            return False
        return True


# ---------------- v2.4 shared manifest gate (V1, V5, V7, V11) ----------------
CLOCK_SKEW = 300


def _is_int(x):
    return isinstance(x, int) and not isinstance(x, bool)


def check_schema(m):
    """V5: type-check signed fields up front; malformed = Reject (never ValueError/KeyError)."""
    for k in ("created", "expires", "seq"):
        if not _is_int(m.get(k)):
            raise Reject(f"manifest schema: {k} must be an integer")
    files = m.get("files")
    if not isinstance(files, list) or not files:
        raise Reject("manifest schema: files must be a non-empty list")
    for f in files:
        if not (isinstance(f, dict) and isinstance(f.get("id"), str) and f["id"]
                and "/" not in f["id"] and not f["id"].startswith(".")
                and _is_int(f.get("size_padded")) and f["size_padded"] > 0
                and isinstance(f.get("sha256_ct"), str)):
            raise Reject("manifest schema: malformed file entry")
    if "prev_batch_id" in m and not isinstance(m["prev_batch_id"], str):
        raise Reject("manifest schema: prev_batch_id must be a string")


def verify_manifest(mbytes, sig, registry, now):
    """Registry gate + Ed25519 over exact canonical bytes + schema. Returns (manifest, registry entry).
    Raises Reject. Used by ingest() AND by the Receiver before any hold/dedup/chain decision (V1, V11)."""
    from cryptography.exceptions import InvalidSignature
    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
    try:
        m = json.loads(mbytes)
    except Exception:
        raise Reject("manifest/signature unparseable")
    try:
        if not isinstance(m, dict) or canonical(m) != mbytes:
            raise Reject("manifest is not canonical")
    except ValueError:
        raise Reject("manifest is not canonical")
    skid = m.get("sender_key_id")
    ent = registry.get(skid) if isinstance(skid, str) else None
    if not ent:
        raise Reject("unknown signer key")
    if ent.get("state") != "verified" or not ent.get("verified_by") or not ent.get("verified_at"):
        raise Reject("signer key not verified (pending)")
    if ent.get("revoked_at") is not None and now >= int(ent["revoked_at"]):
        raise Reject("signer key revoked")
    if ent.get("expires") is not None and now >= int(ent["expires"]):
        raise Reject("signer key expired in registry")
    pub_raw = unb64(ent["ed25519_pub"])
    if key_id(pub_raw) != skid or ent.get("sender_id") != m.get("sender_id"):
        raise Reject("registry entry does not match manifest sender")
    try:
        Ed25519PublicKey.from_public_bytes(pub_raw).verify(sig, mbytes)
    except InvalidSignature:
        raise Reject("bad signature")
    if m.get("v") not in SUPPORTED_V or (m.get("alg"), m.get("kem")) != (ALG, KEM):
        raise Reject("unsupported version/alg")
    check_schema(m)
    # V7: receiver-clock bounds on the sender-chosen created
    if m["created"] > now + CLOCK_SKEW:
        raise Reject("manifest created in the future (clock skew bound)")
    if m["expires"] <= m["created"]:
        raise Reject("manifest expires <= created")
    if m["expires"] - m["created"] > MAX_TTL:
        raise Reject(f"manifest ttl exceeds MAX_TTL ({MAX_TTL} s)")
    if ent.get("revoked_at") is not None and m["created"] >= int(ent["revoked_at"]):
        raise Reject("signed at/after key revocation")
    return m, ent


# ---------------- v2.5 per-agent folders ----------------
import re as _re
SLUG_RE = _re.compile(r"^[a-z0-9-]{3,64}$")


def valid_slug(s) -> bool:
    """Folder slug: lowercase [a-z0-9-], 3-64 chars, no dots or slashes (e.g. agent-alpha)."""
    return isinstance(s, str) and bool(SLUG_RE.match(s))


def check_slug(s):
    if not valid_slug(s):
        raise ValueError(f"invalid folder slug {s!r}: need lowercase [a-z0-9-], 3-64 chars")
    return s


def validate_registry_folders(registry):
    """Every entry has a valid `folder` slug, unique across the registry. Raises ValueError."""
    seen = {}
    for skid, ent in registry.items():
        f = check_slug(ent.get("folder"))
        if f in seen:
            raise ValueError(f"folder {f!r} registered to two keys")
        seen[f] = skid
    return seen


def active_folders(registry, now):
    """Folders the receiver watches: verified, not revoked, not expired senders only. {slug: skid}."""
    out = {}
    for skid, ent in registry.items():
        if ent.get("state") != "verified" or not ent.get("verified_by") or not ent.get("verified_at"):
            continue
        if ent.get("revoked_at") is not None and now >= int(ent["revoked_at"]):
            continue
        if ent.get("expires") is not None and now >= int(ent["expires"]):
            continue
        if valid_slug(ent.get("folder")):
            out[ent["folder"]] = skid
    return out