"""Quick rebuild after a change in the runtime Java only (src/lihont/two).

Takes an APK that py -m lihont.upgrade built earlier, swaps the mod's runtime classes inside its
classes35.dex for a fresh compile of src/lihont/two, and signs the result. The smali patches of TikTok's
own classes are not touched, so this is only right while lihont/*_patch.py and lihont/upgrade.py are
unchanged since that APK was built; after any such change run the full build.

usage (from the repository root):
    py -B scripts/quick_runtime.py ComplectAPK/lihont-2.0-luna-beta5.apk -o ComplectAPK/lihont-2.0-luna-beta6.apk
"""

import argparse
import os
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
os.chdir(ROOT)

from lihont import dexpatch, manifest as manifest_module  # noqa: E402
from lihont.apkzip import Apk  # noqa: E402
from lihont.axml import Axml  # noqa: E402
from lihont.toolchain import Toolchain  # noqa: E402

RUNTIME = "Llihont/two/"
DEX = "classes35.dex"


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("built", help="an APK built by lihont.upgrade")
    parser.add_argument("-o", "--out", required=True)
    parser.add_argument("--work", default="work/quick")
    args = parser.parse_args()

    source, output, room = Path(args.built).resolve(), Path(args.out).resolve(), Path(args.work).resolve()
    if source == output:
        raise SystemExit("Write to a new file: the built APK stays as it is")
    room.mkdir(parents=True, exist_ok=True)
    tools = Toolchain("tools")
    apk = Apk(str(source))
    try:
        api = manifest_module.min_sdk(Axml.parse(apk.read("AndroidManifest.xml")))
        runtime = Path(tools.compile_dex("src/lihont/two", str(room / "runtime"), api)).read_bytes()

        compiler = room / "compiler"
        compiler.mkdir(parents=True, exist_ok=True)
        subprocess.run(["javac", "--release", "17", "-cp", tools.smali, "-d", str(compiler),
                        "src/RefreshClasses.java"], check=True)
        host_dex = apk.read(DEX)
        host, fresh, merged = room / "host.dex", room / "fresh.dex", room / "merged.dex"
        host.write_bytes(host_dex)
        fresh.write_bytes(runtime)
        subprocess.run(["java", "-Xmx1g", "-cp", os.pathsep.join([str(compiler), tools.smali]),
                        "RefreshClasses", str(host), str(fresh), str(merged), str(api)], check=True)
        result = merged.read_bytes()

        if dexpatch.dex_format(host_dex) != dexpatch.dex_format(result):
            raise SystemExit("The refresh changed the DEX format")
        before = {m for m in dexpatch.defined_methods(host_dex) if not m.startswith(RUNTIME)}
        after = dexpatch.defined_methods(result)
        if not before.issubset(after):
            raise SystemExit("The refresh lost classes that are not part of the runtime")
        if not any(m.startswith("Llihont/two/LihonTRegion;->") for m in after):
            raise SystemExit("LihonTRegion is missing after the refresh")

        apk.replace(DEX, result)
        apk.drop_signature()
        output.parent.mkdir(parents=True, exist_ok=True)
        apk.write(str(output))
    finally:
        apk.close()
    tools.sign(str(output))
    print("%s  %d bytes" % (output.name, output.stat().st_size))


if __name__ == "__main__":
    main()