codezip.py

#!/usr/bin/env python3
"""
codezip.py — Lossless compression system for Python code files (v2)

Design principles
------------------
1. Fully lossless: every candidate output is decompressed and compared
   byte-for-byte inside the compressor; only verified outputs compete
   and RAW storage is the guaranteed fallback.
2. Code-aware preprocessing: common Python tokens are substituted with
   1-byte dictionary codes (ESCAPE + code) before a general-purpose
   compressor runs. The mapping is bijective, fully reversible, and —
   since v2 — implemented with C-level primitives (bytes regex for
   encoding, bytes.find loop for decoding) instead of per-byte Python
   loops.
3. Candidate race: several codecs (zlib / lzma / zstandard-if-installed,
   each raw and preprocessed) compress in parallel via threads
   (the underlying C codecs release the GIL) and the smallest verified
   payload wins.
4. Speed levels: ``fast`` (default-ish quick), ``balanced`` (default,
   close to max ratio at a fraction of the time), ``max`` (previous
   v1 extreme settings for the best ratio).
5. Optional C extensions (``zstandard``) are used when importable and
   otherwise skipped — the stdlib-only path always works.

Format (.czc)
-------------
MAGIC(4B) + VERSION(1B) + METHOD(1B) + FILENAME_LEN(2B) + FILENAME(UTF-8)
+ ORIGINAL_SIZE(8B) + SHA256(32B) + PAYLOAD

Method codes 0-5 are identical to v1, so v2 reads every v1 file.
Codes 6+ (zstandard) require v2 to decompress.
"""

import argparse
import bz2
import concurrent.futures
import hashlib
import lzma
import os
import re
import shutil
import struct
import subprocess
import time
import zlib

MAGIC = b"CZC1"
VERSION = 1

# ---------------------------------------------------------------------------
# 1. Code-specific dictionary (common Python patterns, greedy longest match)
# ---------------------------------------------------------------------------
PY_DICTIONARY = [
    "        return ",
    "    return ",
    "return ",
    "def __init__(self",
    "def __init__(self, ",
    "self.",
    "@staticmethod\n",
    "@property\n",
    "@classmethod\n",
    "@dataclass\n",
    "if __name__ == '__main__':",
    'if __name__ == "__main__":',
    "import ",
    "from ",
    "class ",
    "def ",
    "elif ",
    "else:",
    "if ",
    "for ",
    "while ",
    "try:",
    "except ",
    "finally:",
    "with ",
    "as ",
    "lambda ",
    "yield ",
    "raise ",
    "assert ",
    "pass\n",
    "break\n",
    "continue\n",
    "and ",
    "or ",
    "not ",
    "in ",
    "is ",
    "print(",
    "len(",
    "range(",
    "str(",
    "int(",
    "float(",
    "bool(",
    "list(",
    "dict(",
    "set(",
    "tuple(",
    "            ",
    "        ",
    "    ",
    "self, ",
    "__init__",
    "__main__",
    "__name__",
    "__repr__",
    "__str__",
    '"""',
    "# ",
    "## ",
    "### ",
    ": str",
    ": int",
    ": bool",
    ": float",
    ": bytes",
    ": List[",
    ": Dict[",
    ": Optional[",
    ": Union[",
    "-> None:",
    "-> bool:",
    "-> str:",
    "-> int:",
    "-> Optional[",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from typing import ",
    "from dataclasses import ",
    "async def ",
    "await ",
    "    if ",
    "    for ",
    "    def ",
    "        self.",
    "raise ValueError(",
    "raise TypeError(",
    "raise Exception(",
    "logging.",
    "logger.",
    ".append(",
    ".extend(",
    ".format(",
]

# NOTE: dedup while preserving order (sorting a set() would make the token
# order depend on PYTHONHASHSEED and break cross-process determinism).
_seen = set()
_deduped = []
for _t in PY_DICTIONARY:
    if _t not in _seen:
        _seen.add(_t)
        _deduped.append(_t)
_TOKENS = sorted(_deduped, key=len, reverse=True)[:255]
_TOKEN_BYTES = [t.encode("utf-8") for t in _TOKENS]
ESCAPE = 0x00
_NUL = b"\x00"

# Single bytes-regex: literal NUL first, then tokens longest-first.
# re alternation is ordered, so longest match wins; scanning is C-level.
_RX = re.compile(_NUL + b"|" + b"|".join(re.escape(t) for t in _TOKEN_BYTES))
_SUB_OF = {b"\x00": b"\x00\x00"}
for _i, _tb in enumerate(_TOKEN_BYTES):
    _SUB_OF[_tb] = bytes((ESCAPE, _i + 1))


def preprocess_encode(data: bytes) -> bytes:
    """Replace dictionary tokens with ESCAPE+code (2 bytes). Literal
    0x00 bytes in the original data are escaped as 0x00 0x00."""
    if not data:
        return b""
    return _RX.sub(lambda m: _SUB_OF[m.group(0)], data)


def preprocess_decode(data: bytes) -> bytes:
    """Inverse of preprocess_encode. Pure-C scan via bytes.find; the
    Python loop only visits escape positions (rare), not every byte."""
    if not data or data.find(_NUL) < 0:
        return data if isinstance(data, bytes) else bytes(data)
    out = bytearray()
    pos, n = 0, len(data)
    find = data.find
    while True:
        j = find(_NUL, pos)
        if j < 0:
            out += data[pos:]
            break
        out += data[pos:j]
        if j + 1 >= n:
            raise ValueError("Truncated escape sequence in preprocessed data")
        code = data[j + 1]
        if code == ESCAPE:
            out.append(ESCAPE)
        else:
            out += _TOKEN_BYTES[code - 1]
        pos = j + 2
    return bytes(out)


# ---------------------------------------------------------------------------
# 2. Codecs (optional C extensions degrade gracefully to stdlib-only)
# ---------------------------------------------------------------------------
try:
    import zstandard as _zstd

    _HAS_ZSTD = True
except ImportError:  # stdlib-only environment: zstd candidates are skipped
    _zstd = None
    _HAS_ZSTD = False

METHOD_CODE = {
    "RAW": 0,
    "Z": 1,
    "L": 2,
    "PZ": 3,
    "PL": 4,
    "LX": 5,
    "ZS": 6,
    "PZS": 7,
}
CODE_METHOD = {v: k for k, v in METHOD_CODE.items()}

_LZMA_PRESET_MAX = 9 | lzma.PRESET_EXTREME
_LZMA2_FILTERS = [
    {
        "id": lzma.FILTER_LZMA2,
        "preset": 9 | lzma.PRESET_EXTREME,
        "lc": 3,
        "lp": 0,
        "pb": 0,
        "mf": lzma.MF_BT4,
        "depth": 512,
    },
]


def _raw_compress(data: bytes) -> bytes:
    return data


def _zlib_compress(data: bytes, level: int) -> bytes:
    return zlib.compress(data, level)


def _lzma_compress(data: bytes, preset: int) -> bytes:
    return lzma.compress(data, preset=preset)


def _lzma2_extreme_compress(data: bytes) -> bytes:
    return lzma.compress(data, format=lzma.FORMAT_XZ, filters=_LZMA2_FILTERS)


def _zstd_compress(data: bytes, level: int) -> bytes:
    assert _zstd is not None, "zstandard package not installed"
    return _zstd.ZstdCompressor(level=level).compress(data)


def _zstd_decompress(payload: bytes) -> bytes:
    assert _zstd is not None, "zstandard package not installed"
    return _zstd.ZstdDecompressor().decompress(payload)


def decompress_bytes(method: str, payload: bytes) -> bytes:
    if method == "RAW":
        return payload
    if method == "Z":
        return zlib.decompress(payload)
    if method == "L":
        return lzma.decompress(payload)
    if method == "PZ":
        return preprocess_decode(zlib.decompress(payload))
    if method == "PL":
        return preprocess_decode(lzma.decompress(payload))
    if method == "LX":
        # LX payloads are XZ-format LZMA2 of preprocessed data; the plain
        # lzma decoder handles both XZ and legacy LZMA containers.
        return preprocess_decode(lzma.decompress(payload))
    if method == "ZS":
        if not _HAS_ZSTD:
            raise ValueError("ZS payload needs the 'zstandard' package")
        return _zstd_decompress(payload)
    if method == "PZS":
        if not _HAS_ZSTD:
            raise ValueError("PZS payload needs the 'zstandard' package")
        return preprocess_decode(_zstd_decompress(payload))
    raise ValueError(f"Unknown compression method: {method}")


# A candidate spec is (method, use_preprocessed_input, codec, codec_arg).
# Same method code may appear twice at different levels (e.g. PL@6, PL@9e);
# the winner is picked purely by verified payload size.
def _candidate_specs(level: str):
    specs = [("RAW", False, "raw", 0)]
    if level == "fast":
        specs += [("Z", False, "zlib", 1), ("PZ", True, "zlib", 1)]
        if _HAS_ZSTD:
            specs += [("ZS", False, "zstd", 3), ("PZS", True, "zstd", 3)]
    elif level == "max":
        specs += [
            ("Z", False, "zlib", 9),
            ("PZ", True, "zlib", 9),
            ("L", False, "lzma", 6),
            ("PL", True, "lzma", 6),
            ("PL", True, "lzma", _LZMA_PRESET_MAX),
            ("LX", True, "lzma2x", 0),
        ]
        if _HAS_ZSTD:
            specs += [
                ("ZS", False, "zstd", 10),
                ("PZS", True, "zstd", 10),
                ("PZS", True, "zstd", 19),
            ]
    else:  # "balanced" (default)
        specs += [
            ("Z", False, "zlib", 9),
            ("PZ", True, "zlib", 9),
            ("L", False, "lzma", 6),
            ("PL", True, "lzma", 6),
        ]
        if _HAS_ZSTD:
            specs += [("ZS", False, "zstd", 10), ("PZS", True, "zstd", 10)]
    return specs


def _run_spec(spec, data: bytes, pre: bytes | None):
    """Compress + immediately verify one candidate. Returns
    (payload_size, spec_order, method, payload) or None."""
    order, (method, use_pre, codec, arg) = spec
    src = pre if use_pre else data
    assert src is not None
    try:
        if codec == "raw":
            payload = src
        elif codec == "zlib":
            payload = _zlib_compress(src, arg)
        elif codec == "lzma":
            payload = _lzma_compress(src, arg)
        elif codec == "lzma2x":
            payload = _lzma2_extreme_compress(src)
        elif codec == "zstd":
            payload = _zstd_compress(src, arg)
        else:
            return None
        if decompress_bytes(method, payload) != data:
            return None
    except Exception:
        return None
    return (len(payload), order, method, payload)


def compress_data(data: bytes, level: str = "balanced"):
    """Race all candidates for *level* in parallel threads and return
    the (method, payload) of the smallest verified output."""
    if level not in ("fast", "balanced", "max"):
        raise ValueError(f"Unknown level: {level}")
    specs = _candidate_specs(level)
    needs_pre = any(use_pre for _, use_pre, _, _ in specs)
    pre = preprocess_encode(data) if needs_pre and data else data
    ordered = list(enumerate(specs))
    workers = min(len(specs), os.cpu_count() or 4)
    best = None
    if len(specs) == 1 or len(data) < 1024:
        # Thread startup is not worth it for a single candidate or tiny input.
        for spec in ordered:
            res = _run_spec(spec, data, pre)
            if res is not None and (best is None or (res[0], res[1]) < (best[0], best[1])):
                best = res
    else:
        with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as ex:
            futs = [ex.submit(_run_spec, spec, data, pre) for spec in ordered]
            for f in concurrent.futures.as_completed(futs):
                res = f.result()
                if res is not None and (best is None or (res[0], res[1]) < (best[0], best[1])):
                    best = res
    if best is None:  # unreachable in practice (RAW always verifies)
        return "RAW", data
    return best[2], best[3]


def decompress_payload(method: str, payload: bytes) -> bytes:
    return decompress_bytes(method, payload)


# ---------------------------------------------------------------------------
# 3. File format I/O
# ---------------------------------------------------------------------------
def compress_file(
    input_path: str,
    output_path: str | None = None,
    level: str = "balanced",
    canonical: bool = False,
    formatter: str | None = None,
) -> dict:
    with open(input_path, "rb") as f:
        data = f.read()

    if canonical:
        if formatter is None:
            formatter = _find_canonical_formatter()
        data = _apply_canonical(data, formatter)

    original_hash = hashlib.sha256(data).digest()
    method, payload = compress_data(data, level=level)

    filename = os.path.basename(input_path)
    stem, ext = os.path.splitext(filename)
    out_filename = stem + ".czc" if not canonical else stem + ".canonical.czc"
    filename_b = filename.encode("utf-8")

    header = bytearray()
    header += MAGIC
    header.append(VERSION)
    header.append(METHOD_CODE[method] | (0x80 if canonical else 0x00))
    header += struct.pack(">H", len(filename_b))
    header += filename_b
    header += struct.pack(">Q", len(data))
    header += original_hash

    if output_path is None:
        output_path = out_filename
    with open(output_path, "wb") as f:
        f.write(header)
        f.write(payload)

    compressed_size = len(header) + len(payload)
    return {
        "output_path": output_path,
        "method": method,
        "level": level,
        "canonical": canonical,
        "formatter": formatter if canonical else None,
        "original_size": len(data),
        "compressed_size": compressed_size,
        "ratio": (1 - compressed_size / len(data)) * 100 if data else 0.0,
    }


def _find_canonical_formatter() -> str:
    if shutil.which("ruff"):
        return "ruff"
    if _check_black_import():
        return "black"
    raise RuntimeError(
        "Canonical mode requires either `ruff format` or `black` to be installed. "
        "Install with: pip install ruff (recommended) or pip install black"
    )


def _check_black_import() -> bool:
    try:
        import black  # noqa: F811

        return True
    except ImportError:
        return False


def _apply_canonical(data: bytes, formatter: str) -> bytes:
    if formatter == "ruff":
        result = subprocess.run(
            ["ruff", "format", "-"],
            input=data,
            capture_output=True,
        )
        if result.returncode != 0:
            msg = result.stderr.decode("utf-8", errors="replace")[-300:]
            raise RuntimeError(f"ruff format failed (rc={result.returncode}): {msg}")
        return result.stdout
    elif formatter == "black":
        import black

        src = data.decode("utf-8")
        try:
            formatted = black.format_str(src, mode=black.Mode())
        except Exception as e:
            raise RuntimeError(f"black format failed: {e}") from e
        return formatted.encode("utf-8") + b"\n"
    else:
        raise ValueError(f"Unknown formatter: {formatter}")


def decompress_file(input_path: str, output_path: str | None = None) -> dict:
    with open(input_path, "rb") as f:
        blob = f.read()

    if blob[:4] != MAGIC:
        raise ValueError("Not a valid .czc file (MAGIC mismatch)")
    version = blob[4]
    if version != VERSION:
        raise ValueError(f"Unsupported format version: {version}")

    method_byte = blob[5]
    canonical = bool(method_byte & 0x80)
    method_code = method_byte & 0x7F
    method = CODE_METHOD.get(method_code)
    if method is None:
        raise ValueError(f"Unknown compression method code: {method_code}")

    offset = 6
    (filename_len,) = struct.unpack_from(">H", blob, offset)
    offset += 2
    filename = blob[offset : offset + filename_len].decode("utf-8")
    offset += filename_len

    (original_size,) = struct.unpack_from(">Q", blob, offset)
    offset += 8
    original_hash = blob[offset : offset + 32]
    offset += 32

    payload = blob[offset:]
    restored = decompress_bytes(method, payload)

    if len(restored) != original_size:
        raise IOError("Integrity check failed: restored size differs from original")
    if hashlib.sha256(restored).digest() != original_hash:
        raise IOError("Integrity check failed: SHA-256 hash differs from original")

    if output_path is None:
        output_path = filename
    with open(output_path, "wb") as f:
        f.write(restored)

    return {
        "output_path": output_path,
        "method": method,
        "canonical": canonical,
        "restored_size": len(restored),
        "verified": True,
    }


# ---------------------------------------------------------------------------
# 4. Benchmarks (all in-memory — no temp files, unlike v1)
# ---------------------------------------------------------------------------
def _time_fn(fn, *args):
    t0 = time.perf_counter()
    out = fn(*args)
    return out, (time.perf_counter() - t0) * 1000.0


def benchmark_data(data: bytes, best: bool = False):
    """Return rows of (label, size, ratio, enc_ms, dec_ms, ok) for *data*."""
    original_size = len(data)
    rows = []

    def ratio_of(n):
        return (1 - n / original_size) * 100 if original_size else 0.0

    def add_baseline(label, enc_fn, dec_fn):
        try:
            comp, enc_ms = _time_fn(enc_fn, data)
            rest, dec_ms = _time_fn(dec_fn, comp)
            rows.append((label, len(comp), ratio_of(len(comp)), enc_ms, dec_ms, rest == data))
        except Exception:
            rows.append((label, -1, 0.0, 0.0, 0.0, False))

    def add_codezip(label, level):
        try:
            (method, payload), enc_ms = _time_fn(compress_data, data, level)
            rest, dec_ms = _time_fn(decompress_payload, method, payload)
            rows.append(
                (
                    f"{label} [{method}]",
                    len(payload),
                    ratio_of(len(payload)),
                    enc_ms,
                    dec_ms,
                    rest == data,
                )
            )
        except Exception:
            rows.append((label, -1, 0.0, 0.0, 0.0, False))

    add_baseline("gzip -9", lambda d: zlib.compress(d, 9), zlib.decompress)
    add_baseline(
        "bzip2 -9", lambda d: bz2.compress(d, compresslevel=9), bz2.decompress
    )
    add_baseline("lzma std", lzma.compress, lzma.decompress)
    if _HAS_ZSTD:
        add_baseline(
            "zstd -3",
            lambda d: _zstd_compress(d, 3),
            _zstd_decompress,
        )
        add_baseline(
            "zstd -19",
            lambda d: _zstd_compress(d, 19),
            _zstd_decompress,
        )
    add_codezip("codezip fast", "fast")
    add_codezip("codezip balanced", "balanced")
    if best:
        add_baseline(
            "lzma -9e",
            lambda d: lzma.compress(d, preset=_LZMA_PRESET_MAX),
            lzma.decompress,
        )
        add_codezip("codezip max", "max")
    return rows


def run_benchmark(input_path: str, best: bool = False):
    with open(input_path, "rb") as f:
        data = f.read()
    rows = benchmark_data(data, best=best)

    print(f"\nBenchmark: {input_path} ({len(data):,} bytes)\n")
    print(
        f"{'Method':<24} {'Size':>8} {'Ratio':>8} {'Enc(ms)':>10} {'Dec(ms)':>10} {'OK'}"
    )
    print("-" * 74)
    for label, size, ratio, enc_t, dec_t, ok in rows:
        size_s = f"{size:,}" if size >= 0 else "ERR"
        print(
            f"{label:<24} {size_s:>8} {ratio:>7.1f}% {enc_t:>9.1f} {dec_t:>9.1f} {'PASS' if ok else 'FAIL'}"
        )
    print()


# ---------------------------------------------------------------------------
# 5. CLI
# ---------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(
        description="Lossless compression/decompression tool for Python code files (.czc)"
    )
    sub = parser.add_subparsers(dest="command", required=True)

    p_c = sub.add_parser("compress", help="Compress a file into .czc")
    p_c.add_argument("input", help="Path to the source file to compress")
    p_c.add_argument("-o", "--output", help="Output path (default: <input>.czc)")
    p_c.add_argument(
        "--level",
        choices=["fast", "balanced", "max"],
        default="balanced",
        help="Speed/ratio trade-off (default: balanced)",
    )
    p_c.add_argument(
        "--mode",
        choices=["lossless", "canonical"],
        default="lossless",
        help="Compression mode (default: lossless)",
    )
    p_c.add_argument(
        "--formatter",
        choices=["black", "ruff"],
        help="Formatter for canonical mode (auto-detect if omitted)",
    )

    p_d = sub.add_parser("decompress", help="Restore a .czc file to the original")
    p_d.add_argument("input", help="Path to the .czc file to decompress")
    p_d.add_argument("-o", "--output", help="Output path (default: original filename)")

    p_b = sub.add_parser("bench", help="Benchmark against standard compressors")
    p_b.add_argument("input", nargs="+", help="Path(s) to the file(s) to benchmark")
    p_b.add_argument(
        "--best",
        action="store_true",
        help="Also run the slow extremes (lzma -9e, codezip max)",
    )

    args = parser.parse_args()

    if args.command == "compress":
        canonical = args.mode == "canonical"
        result = compress_file(
            args.input,
            args.output,
            level=args.level,
            canonical=canonical,
            formatter=args.formatter,
        )
        mode_label = "Canonical" if canonical else "Lossless"
        print(f"Compression complete -> {result['output_path']}")
        print(f"  Mode: {mode_label}")
        if canonical:
            print(f"  Formatter: {result['formatter']}")
        print(f"  Level: {result['level']}")
        print(f"  Method: {result['method']}")
        print(f"  Original size: {result['original_size']:,} bytes")
        print(f"  Compressed size: {result['compressed_size']:,} bytes")
        print(f"  Ratio: {result['ratio']:.1f}%")
    elif args.command == "decompress":
        result = decompress_file(args.input, args.output)
        mode_label = "Canonical" if result.get("canonical") else "Lossless"
        print(f"Decompression complete -> {result['output_path']}")
        print(f"  Mode: {mode_label}")
        print(f"  Method: {result['method']}")
        print(f"  Size: {result['restored_size']:,} bytes")
        print("  Integrity check: PASSED (SHA-256 match)")
    elif args.command == "bench":
        for path in args.input:
            run_benchmark(path, best=args.best)


if __name__ == "__main__":
    main()