import os
import struct
import sys
import zipfile

RENAME = {"Cif": "if", "Cdo": "do"}


def rewrite_utf8(s: str) -> str:
    if "defpackage/" in s:
        s = s.replace("defpackage/", "")
    for old, new in RENAME.items():
        if old in s:
            s = s.replace(old, new)
    return s


def rewrite_class(data: bytes) -> bytes:
    if data[:4] != b"\xca\xfe\xba\xbe":
        return data
    cp_count = struct.unpack_from(">H", data, 8)[0]
    i = 10
    entries = []
    idx = 1
    while idx < cp_count:
        tag = data[i]
        i += 1
        if tag == 1:
            ln = struct.unpack_from(">H", data, i)[0]
            i += 2
            raw = data[i : i + ln]
            i += ln
            try:
                txt = raw.decode("utf-8")
            except UnicodeDecodeError:
                entries.append(bytes([1]) + struct.pack(">H", ln) + raw)
                idx += 1
                continue
            new = rewrite_utf8(txt).encode("utf-8")
            entries.append(struct.pack(">BH", 1, len(new)) + new)
        elif tag in (5, 6):
            entries.append(bytes([tag]) + data[i : i + 8])
            i += 8
            entries.append(None)
            idx += 2
            continue
        elif tag in (3, 4):
            entries.append(bytes([tag]) + data[i : i + 4])
            i += 4
        elif tag in (7, 8, 16, 19, 20):
            entries.append(bytes([tag]) + data[i : i + 2])
            i += 2
        elif tag == 15:
            entries.append(bytes([tag]) + data[i : i + 3])
            i += 3
        elif tag in (9, 10, 11, 12, 17, 18):
            entries.append(bytes([tag]) + data[i : i + 4])
            i += 4
        else:
            raise ValueError(f"unknown constant pool tag {tag} at offset {i - 1}")
        idx += 1
    out = bytearray(data[:8])
    out += struct.pack(">H", cp_count)
    for e in entries:
        if e is not None:
            out += e
    out += data[i:]
    return bytes(out)


def rename_stem(name: str, sep: str) -> str:
    parts = name.split(sep)
    stem = parts[-1][: -len(".class")] if parts[-1].endswith(".class") else None
    if stem in RENAME:
        parts[-1] = RENAME[stem] + ".class"
        return sep.join(parts)
    return name


def process_dir(root: str) -> tuple[int, int]:
    moved = 0
    touched = 0
    for dp, _, files in os.walk(root):
        for name in files:
            if not name.endswith(".class"):
                continue
            path = os.path.join(dp, name)
            data = open(path, "rb").read()
            new = rewrite_class(data)
            rel = os.path.relpath(path, root)
            prefix = "defpackage" + os.sep
            if rel.startswith(prefix):
                target = os.path.join(root, rename_stem(rel[len(prefix) :], os.sep))
            else:
                target = os.path.join(root, rename_stem(rel, os.sep))
            if target != path:
                os.makedirs(os.path.dirname(target), exist_ok=True)
                out = open(target, "wb")
                out.write(new)
                out.close()
                os.remove(path)
                moved += 1
            elif new != data:
                out = open(path, "wb")
                out.write(new)
                out.close()
                touched += 1
    return moved, touched


def process_jar(path: str) -> tuple[int, int]:
    src = zipfile.ZipFile(path)
    tmp = path + ".tmp"
    dst = zipfile.ZipFile(tmp, "w", zipfile.ZIP_DEFLATED)
    moved = 0
    touched = 0
    for info in src.infolist():
        name = info.filename
        data = src.read(name)
        if name.endswith(".class"):
            new = rewrite_class(data)
            if name.startswith("defpackage/"):
                name = name[len("defpackage/") :]
                moved += 1
            elif new != data:
                touched += 1
            renamed = rename_stem(name, "/")
            if renamed != name:
                name = renamed
                touched += 1
            data = new
        dst.writestr(info, data)
    dst.close()
    src.close()
    os.replace(tmp, path)
    return moved, touched


def main() -> int:
    if len(sys.argv) != 2:
        print("usage: postbuild.py <classes-dir | jar>", file=sys.stderr)
        return 2
    target = sys.argv[1]
    if os.path.isdir(target):
        moved, touched = process_dir(target)
    elif target.endswith(".jar"):
        moved, touched = process_jar(target)
    else:
        print(f"not a dir or jar: {target}", file=sys.stderr)
        return 2
    print(f"postbuild: moved={moved} rewritten={touched} ({target})")
    return 0


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