#!/usr/bin/env python3
"""Deterministic repack of a jar, replicating stripJavaArchivesHook for the
penpot backend uberjar.

Nix's stripJavaArchivesHook runs strip-nondeterminism, which drives the
perl Archive::Zip module. Two problems made it unusable here:

  * its per-member removeMember/addMember reorder loop is O(n^2) in the entry
    count (linear member search + full array rebuild per entry), so the
    54000-entry, 109 MB penpot jar spent ~32 minutes there;
  * Archive::Zip derives the DOS timestamps with localtime, so the result
    depended on the build machine's timezone.

This script reproduces the same semantics (jar.pm + zip.pm handlers) with the
stdlib zipfile module: read each entry once, sort the entry list, write it in a
single pass. The result is byte-identical across runs and timezones, and costs
seconds instead of minutes.

Semantics replicated from strip-nondeterminism:
  * entry order: META-INF/ first, then META-INF/MANIFEST.MF, then lexicographic
  * canonical time: Archive::Zip SAFE_EPOCH (1980-01-01 12:01:00 local), used
    here as a fixed UTC tuple so the output never depends on $TZ
  * .clj entries are stamped 2 DOS seconds (1 unix second) earlier, so that
    Clojure never mistakes a .clj for a stale .class
  * META-INF/MANIFEST.MF: drop Bnd-LastModified:/Built-By: lines
  * .html with a `<!-- Generated by javadoc` marker in the first 1 KB: bump
    the generation date to the canonical date
  * nested .jar entries: recurse with the same semantics
  * extra fields 0x5455 (extended timestamps), 0x000a (NTFS), 0x7875 (ux
    uid/gid) are normalized; unknown fields are kept verbatim
  * archives with a META-INF/*.SF signature are left untouched
  * javac.sh and plexus compiler-arguments files are dropped
"""

import io
import os
import re
import struct
import sys
import tempfile
import zipfile

# Archive::Zip SAFE_EPOCH: 1980-01-01 12:01:00. The DOS timestamp below is the
# same instant, but written as a fixed tuple: no localtime call on the way out.
DOS_DATE = (1980, 1, 1)
DOS_TIME = (12, 1, 0)
DOS_EPOCH_SECONDS = 315576060

MANIFEST = "META-INF/MANIFEST.MF"
MANIFEST_DIR = "META-INF/"
JAVADOC_WINDOW = 1024
JAVADOC_MARKER = b"<!-- Generated by javadoc"

# From jar.pm's extract helper: a `YYYY-MM-DD` following this class is the
# javadoc generation timestamp.
JAVADOC_DATE_RE = re.compile(rb"(<!-- Generated by javadoc[^>]*?\son\s)(\d{4}-\d{2}-\d{2})")


def order_key(name: str):
    if name == MANIFEST_DIR:
        return (0, name)
    if name == MANIFEST:
        return (1, name)
    return (2, name)


def dropped(name: str) -> bool:
    return name == "javac.sh" or re.fullmatch(
        r"org\.codehaus\.plexus\.compiler\.javac\.JavacCompiler\d+arguments", name
    ) is not None


def canonical_tuple(entry_name: str):
    if entry_name.endswith(".clj"):
        # 2 DOS seconds earlier, i.e. one unix second (DOS has 2 s granularity).
        seconds = DOS_TIME[0] * 3600 + DOS_TIME[1] * 60 + DOS_TIME[2] - 2
        return DOS_DATE + (seconds // 3600, (seconds % 3600) // 60, seconds % 60)
    return DOS_DATE + DOS_TIME


def normalize_manifest(data: bytes) -> bytes:
    out = []
    for line in data.splitlines(keepends=True):
        if line.startswith(b"Bnd-LastModified:") or line.startswith(b"Built-By:"):
            continue
        out.append(line)
    return b"".join(out)


def normalize_javadoc(data: bytes) -> bytes:
    head = data[:JAVADOC_WINDOW]
    if JAVADOC_MARKER not in head:
        return data

    def repl(match):
        return match.group(1) + b"1980-01-01"

    return JAVADOC_DATE_RE.sub(repl, data, count=1)


def normalize_extra(extra: bytes, filename: str) -> bytes:
    """Normalize known timestamp/identity extra fields, keep the rest as-is."""
    want = canonical_tuple(filename)
    out = bytearray()
    pos = 0
    while pos + 4 <= len(extra):
        field_id, field_len = struct.unpack_from("<HH", extra, pos)
        data = extra[pos + 4 : pos + 4 + field_len]
        if field_id == 0x5455:  # extended timestamp
            if len(data) >= 5:
                data = data[:1] + b"\x00\x00\x00\x00" + data[5:]
        elif field_id == 0x000a:  # NTFS timestamp
            if len(data) >= 4:
                reserved = data[:4]
                tags = bytearray()
                tag_pos = 4
                while tag_pos + 4 <= len(data):
                    tag, tag_size = struct.unpack_from("<HH", data, tag_pos)
                    tags += data[tag_pos : tag_pos + 4]
                    if tag == 0x0001 and tag_size == 24:
                        # 32 bytes of 100ns ticks since 1601, i.e. 12 bytes of
                        # zeros followed by the unix seconds in the low word.
                        ticks = b"\x00" * 12 + struct.pack("<I", int(DOS_EPOCH_SECONDS))
                        tags += ticks
                    else:
                        tags += data[tag_pos + 4 : tag_pos + 4 + tag_size]
                    tag_pos += 4 + tag_size
                data = reserved + bytes(tags)
        elif field_id == 0x7875:  # Info-ZIP new unix extra field
            if len(data) >= 3 and data[0] == 1:
                uid_len, gid_len = data[1], data[2 + data[1]]
                data = (
                    b"\x01"
                    + bytes([uid_len])
                    + b"\x00" * uid_len
                    + bytes([gid_len])
                    + b"\x00" * gid_len
                )
        out += struct.pack("<HH", field_id, len(data)) + data
        pos += 4 + field_len
    return bytes(out)


def normalize_member(name: str, data: bytes) -> bytes:
    if name == MANIFEST:
        return normalize_manifest(data)
    if name.endswith(".html"):
        return normalize_javadoc(data)
    if name.endswith(".jar"):
        return normalize_jar_bytes(data)
    return data


def _rewrite(source, target):
    zin = zipfile.ZipFile(source)
    if any(n.startswith("META-INF/") and n.endswith(".SF") for n in zin.namelist()):
        # Signed jar: normalization would break the signature.
        return False, 0

    with zipfile.ZipFile(target, "w") as zout:
        written = 0
        for item in sorted(zin.infolist(), key=lambda i: order_key(i.filename)):
            if dropped(item.filename):
                continue
            data = normalize_member(item.filename, zin.read(item.filename))
            date_time = canonical_tuple(item.filename)
            info = zipfile.ZipInfo(item.filename, date_time=date_time)
            info.compress_type = item.compress_type
            info.external_attr = item.external_attr
            info.internal_attr = item.internal_attr
            info.create_system = 3
            info.extra = normalize_extra(item.extra, item.filename)
            zout.writestr(info, data)
            written += 1
    return True, written


def normalize_jar_bytes(data: bytes) -> bytes:
    buffer = io.BytesIO()
    with zipfile.ZipFile(io.BytesIO(data)) as zin:
        if any(n.startswith("META-INF/") and n.endswith(".SF") for n in zin.namelist()):
            return data
        with zipfile.ZipFile(buffer, "w") as zout:
            for item in sorted(zin.infolist(), key=lambda i: order_key(i.filename)):
                if dropped(item.filename):
                    continue
                info = zipfile.ZipInfo(item.filename, date_time=canonical_tuple(item.filename))
                info.compress_type = item.compress_type
                info.external_attr = item.external_attr
                info.internal_attr = item.internal_attr
                info.create_system = 3
                info.extra = normalize_extra(item.extra, item.filename)
                zout.writestr(info, normalize_member(item.filename, zin.read(item.filename)))
    return buffer.getvalue()


def main() -> int:
    if len(sys.argv) != 2:
        print(f"usage: {sys.argv[0]} <jar>", file=sys.stderr)
        return 2

    path = sys.argv[1]
    tmp_fd, tmp_path = tempfile.mkstemp(dir=os.path.dirname(path) or ".", suffix=".jar")
    os.close(tmp_fd)
    try:
        handled, entries = _rewrite(path, tmp_path)
        if not handled:
            print(f"jar-strip: {path}: signed jar, leaving untouched", file=sys.stderr)
            return 0
        if entries == 0:
            print(f"jar-strip: {path}: no entries written", file=sys.stderr)
            return 1
        # Self-validation: reopen and verify every CRC before trusting the output.
        with zipfile.ZipFile(tmp_path) as check:
            bad = check.testzip()
        if bad is not None:
            print(f"jar-strip: {path}: CRC mismatch at {bad}", file=sys.stderr)
            return 1
        os.replace(tmp_path, path)
        print(f"jar-strip: {path}: normalized {entries} entries")
        return 0
    finally:
        if os.path.exists(tmp_path):
            os.unlink(tmp_path)


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