import logging

FORMAT = "[%(levelname)s] %(name)s %(filename)s:%(lineno)d: %(message)s"
logging.basicConfig(format=FORMAT)
logging.getLogger().setLevel(logging.DEBUG)

import json
import zipfile as z
from pathlib import Path

from androscalpel import Apk, IdType, IdMethodType, ins, DexString, IdMethod, Code, utils  # type: ignore

# APK_NAME = Path(__file__).parent / "test.apk"
APK_NAME = Path(__file__).parent / "app.apk"
DEX_NAME = "classes.dex"

print(f"[+] Load bytecode ")
with z.ZipFile(APK_NAME) as zipf:
    with zipf.open(DEX_NAME, "r") as dex_f:
        dex = dex_f.read()

apk = Apk()
apk.add_dex_file(dex)

clazz_id = IdType("Lcom/example/testapplication/ui/home/HomeViewModel;")
proto_id = IdMethodType(IdType("Ljava/lang/String;"), [])
method_id = IdMethod("text_gen", proto_id, clazz_id)

clazz = apk.classes[clazz_id]
method = clazz.virtual_methods[method_id]
code = method.code

logging.getLogger().setLevel(logging.ERROR)

print(f"[+] Code of {method_id} ")
for i in code.insns:
    print(f"    {i}")
print("[+] Modify code")

new_insns = []
for i in code.insns:
    if isinstance(i, ins.ConstString):
        if i.lit == "Hello":
            i = ins.ConstString(i.reg, DexString("Degemer Mat"))
        elif i.lit == "Bye":
            i = ins.ConstString(i.reg, DexString("Kenavo"))
    new_insns.append(i)

# This need improving!
code = Code(code.registers_size, code.ins_size, code.outs_size, new_insns)
apk.set_method_code(method_id, code)
# apk.set_method_code(method.descriptor, code)


clazz = apk.classes[clazz_id]
method = clazz.virtual_methods[method_id]
code = method.code
print(f"[+] New code of {method_id} ")
for i in code.insns:
    print(f"    {i}")

# Strip class for debugging
classes = list(
    filter(
        lambda x: x
        not in [
            IdType("Lcom/example/testapplication/ui/home/HomeViewModel;"),
            IdType("Landroidx/navigation/NavDeepLink$Builder;"),
            IdType("Landroidx/constraintlayout/core/widgets/ConstraintWidget$1;"),
        ],
        apk.classes.keys(),
    )
)
for cls in classes:
    apk.remove_class(cls)

print("[+] Recompile")

dex_raw = apk.gen_raw_dex()

new_apk = Apk()
for dex in dex_raw:
    new_apk.add_dex_file(dex)

print(f"{new_apk == apk=}")


# print("[+] Repackage")
#
# utils.replace_dex(
#    APK_NAME,
#    APK_NAME.parent / (APK_NAME.name.removesuffix(".apk") + "-instrumented.apk"),
#    dex_raw,
#    Path().parent / "my-release-key.jks",
#    zipalign=Path.home() / "Android" / "Sdk" / "build-tools" / "34.0.0" / "zipalign",
#    apksigner=Path.home() / "Android" / "Sdk" / "build-tools" / "34.0.0" / "apksigner",
# )


def cmp(a, b):
    for f in dir(a):
        if getattr(getattr(a, f), "__call__", None) is None:
            print(f"{f}: {getattr(a, f) == getattr(b, f)}")