#!/usr/bin/env python3
"""Prepare or roll back the Qortal empty-cache runtime overlay.

The official JAR is never overwritten. A package-owned contract binds its exact
bytes, the independently reviewed replacement class, and the still-separate
upstream provenance disposition. Any drift refuses before Core can launch.
"""
import hashlib
import io
import json
import os
from pathlib import Path
import stat
import sys
import zipfile

SCHEMA = "archipelago.qortal.empty-cache-overlay.v1"
VERSION = "6.1.9"
TAG = "v6.1.9"
UPSTREAM_COMMIT = "108bf191d42d710ec617f535af30cfd82fc03c87"
TARGET = "org/qortal/controller/arbitrary/ArbitraryDataCacheManager.class"
PATCH_SHA256 = "7fccb0afc2113abe19af0e1fa2e70cb7a61fbd8c410ad372e4e9d28ef31ae0c8"
MAX_JAR_BYTES = 256 * 1024 * 1024
MAX_CLASS_BYTES = 1024 * 1024
MAX_CONTRACT_BYTES = 16 * 1024


class Refusal(ValueError):
    pass


def sha(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def unique_object(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise Refusal("duplicate JSON property")
        result[key] = value
    return result


def read_regular(path: Path, limit: int) -> bytes:
    fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK)
    with os.fdopen(fd, "rb") as stream:
        if not stat.S_ISREG(os.fstat(stream.fileno()).st_mode):
            raise Refusal("input must be a regular file")
        data = stream.read(limit + 1)
    if len(data) > limit:
        raise Refusal("input exceeds package bound")
    return data


def parse_json(raw: bytes):
    try:
        return json.loads(raw.decode("utf-8"), object_pairs_hook=unique_object)
    except (UnicodeError, json.JSONDecodeError, RecursionError) as exc:
        raise Refusal("invalid strict JSON") from exc


def load_json(path: Path, limit: int):
    return parse_json(read_regular(path, limit))


def validate_contract(doc):
    expected_keys = {
        "schema", "qortal_version", "upstream_tag", "upstream_commit",
        "original_jar_sha256", "target_entry", "replacement_class_sha256",
        "source_patch_sha256", "behavior_review_status",
        "upstream_provenance_status", "provenance_receipt_sha256",
    }
    if type(doc) is not dict or set(doc) != expected_keys:
        raise Refusal("overlay contract schema drift")
    fixed = {
        "schema": SCHEMA,
        "qortal_version": VERSION,
        "upstream_tag": TAG,
        "upstream_commit": UPSTREAM_COMMIT,
        "target_entry": TARGET,
        "source_patch_sha256": PATCH_SHA256,
        "behavior_review_status": "accepted",
        "upstream_provenance_status": "accepted",
    }
    for key, value in fixed.items():
        if doc.get(key) != value:
            raise Refusal("overlay contract is not accepted for this package")
    for key in ("original_jar_sha256", "replacement_class_sha256", "provenance_receipt_sha256"):
        value = doc.get(key)
        if type(value) is not str or len(value) != 64 or any(c not in "0123456789abcdef" for c in value):
            raise Refusal("overlay contract digest is invalid")
    return doc


def write_exclusive(path: Path, data: bytes, mode: int = 0o600):
    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, mode)
    try:
        with os.fdopen(fd, "wb") as stream:
            stream.write(data)
            stream.flush()
            os.fsync(stream.fileno())
    except BaseException:
        try:
            path.unlink()
        except FileNotFoundError:
            pass
        raise


def build_runtime(original: bytes, replacement: bytes) -> bytes:
    if replacement[:4] != b"\xca\xfe\xba\xbe":
        raise Refusal("replacement is not a Java class")
    source = io.BytesIO(original)
    output = io.BytesIO()
    with zipfile.ZipFile(source) as src:
        names = src.namelist()
        if len(set(names)) != len(names) or names.count(TARGET) != 1:
            raise Refusal("official JAR has ambiguous entries")
        if any(name.upper().startswith("META-INF/") and
               name.upper().endswith((".SF", ".RSA", ".DSA", ".EC")) for name in names):
            raise Refusal("refuse to invalidate a signed JAR")
        with zipfile.ZipFile(output, "w") as dst:
            for item in src.infolist():
                dst.writestr(item, replacement if item.filename == TARGET else src.read(item.filename))
    runtime = output.getvalue()
    with zipfile.ZipFile(io.BytesIO(original)) as src, zipfile.ZipFile(io.BytesIO(runtime)) as dst:
        changed = [name for name in names if src.read(name) != dst.read(name)]
    if changed != [TARGET]:
        raise Refusal("runtime overlay changed unexpected entries")
    return runtime


def prepare(original_path: Path, class_path: Path, contract_path: Path,
            output_path: Path, receipt_path: Path):
    contract_raw = read_regular(contract_path, MAX_CONTRACT_BYTES)
    contract = validate_contract(parse_json(contract_raw))
    original = read_regular(original_path, MAX_JAR_BYTES)
    replacement = read_regular(class_path, MAX_CLASS_BYTES)
    if sha(original) != contract["original_jar_sha256"]:
        raise Refusal("official JAR bytes differ from package contract")
    if sha(replacement) != contract["replacement_class_sha256"]:
        raise Refusal("replacement class bytes differ from package contract")
    runtime = build_runtime(original, replacement)
    receipt = {
        "schema": SCHEMA,
        "original_path": str(original_path),
        "original_jar_sha256": sha(original),
        "runtime_path": str(output_path),
        "runtime_jar_sha256": sha(runtime),
        "contract_sha256": sha(contract_raw),
    }
    write_exclusive(output_path, runtime)
    try:
        write_exclusive(receipt_path, (json.dumps(receipt, sort_keys=True) + "\n").encode("ascii"))
    except BaseException:
        output_path.unlink()
        raise


def rollback(original_path: Path, output_path: Path, receipt_path: Path):
    receipt = load_json(receipt_path, MAX_CONTRACT_BYTES)
    expected_keys = {"schema", "original_path", "original_jar_sha256",
                     "runtime_path", "runtime_jar_sha256", "contract_sha256"}
    if type(receipt) is not dict or set(receipt) != expected_keys or receipt.get("schema") != SCHEMA:
        raise Refusal("rollback receipt schema drift")
    if receipt.get("original_path") != str(original_path) or receipt.get("runtime_path") != str(output_path):
        raise Refusal("rollback path drift")
    original = read_regular(original_path, MAX_JAR_BYTES)
    runtime = read_regular(output_path, MAX_JAR_BYTES)
    if sha(original) != receipt.get("original_jar_sha256") or sha(runtime) != receipt.get("runtime_jar_sha256"):
        raise Refusal("rollback byte drift")
    output_path.unlink()
    receipt_path.unlink()


def main():
    try:
        if len(sys.argv) == 7 and sys.argv[1] == "prepare":
            prepare(*(Path(value) for value in sys.argv[2:]))
        elif len(sys.argv) == 5 and sys.argv[1] == "rollback":
            rollback(*(Path(value) for value in sys.argv[2:]))
        else:
            print("usage: prepare-runtime-jar.py prepare ORIGINAL CLASS CONTRACT OUTPUT RECEIPT\n"
                  "       prepare-runtime-jar.py rollback ORIGINAL OUTPUT RECEIPT", file=sys.stderr)
            return 64
    except (OSError, ValueError, zipfile.BadZipFile, RecursionError):
        print("Qortal runtime overlay refused: package provenance, official JAR, replacement class, or rollback bytes drifted.", file=sys.stderr)
        return 78
    return 0


if __name__ == "__main__":
    sys.exit(main())