#!/usr/bin/env python3
"""Replace compatible FFmpeg APK libraries without changing first-party payloads."""

import argparse
import os
from pathlib import Path
import tempfile
import zipfile

LIBRARIES = ("avcodec", "avfilter", "avformat", "avutil", "swresample", "swscale")
MACHINES = {"arm64-v8a": 183, "x86_64": 62}


def replace(apk, libraries, abis, output):
    if output.exists():
        raise FileExistsError(f"Refusing to overwrite: {output}")
    replacements = {}
    for abi in abis:
        for library in LIBRARIES:
            name = f"lib{library}.so"
            data = (libraries / abi / name).read_bytes()
            if (len(data) < 20 or data[:6] != b"\x7fELF\x02\x01"
                    or int.from_bytes(data[18:20], "little") != MACHINES[abi]):
                raise ValueError(f"Not a compatible 64-bit ELF architecture: {abi}/{name}")
            replacements[f"lib/{abi}/{name}"] = data
    with zipfile.ZipFile(apk) as original:
        names = original.namelist()
        if len(names) != len(set(names)):
            raise ValueError("APK contains duplicate entries")
        missing = replacements.keys() - set(names)
        if missing:
            raise ValueError(f"APK does not contain selected libraries: {sorted(missing)}")
        output.parent.mkdir(parents=True, exist_ok=True)
        with tempfile.NamedTemporaryFile(dir=output.parent, suffix=".apk", delete=False) as staged:
            temporary = Path(staged.name)
        try:
            with zipfile.ZipFile(temporary, "w") as result:
                for entry in original.infolist():
                    upper = entry.filename.upper()
                    if upper.startswith("META-INF/") and (
                        upper == "META-INF/MANIFEST.MF"
                        or upper.endswith((".SF", ".RSA", ".DSA", ".EC"))
                        or Path(upper).name.startswith("SIG-")
                    ):
                        continue
                    data = replacements.get(entry.filename)
                    result.writestr(entry, original.read(entry.filename) if data is None else data)
            with zipfile.ZipFile(temporary) as result:
                bad = result.testzip()
                if bad:
                    raise ValueError(f"APK integrity check failed: {bad}")
            os.link(temporary, output)
        finally:
            temporary.unlink(missing_ok=True)
    print(f"Created {output}. Align and sign it with YOUR key before installation.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--apk", type=Path, required=True)
    parser.add_argument("--libraries", type=Path, required=True)
    parser.add_argument("--abi", action="append", choices=MACHINES, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    replace(args.apk, args.libraries, args.abi, args.output)