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.WARNING)

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;"),
#             IdType("Landroidx/appcompat/app/ActionBar;"),
#             IdType("Landroidx/constraintlayout/core/state/WidgetFrame;"),
#             IdType("Landroidx/appcompat/app/AppCompatViewInflater;"),
#         ],
#         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("[+] 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",
)

last_id = None

MAX_REQ = 1


def cmp(a, b, req=0):
    if req > MAX_REQ:
        return
    if type(a) == dict:
        cmp_dict(a, b, req)
    elif type(a) == list:
        cmp_list(a, b, req)
    else:
        cmp_other(a, b, req)


def nice_bool(b) -> str:
    if b:
        return "\033[32mTrue\033[0m"
    else:
        return "\033[31mFalse\033[0m"


def cmp_other(a, b, req=0):
    ident = "  " * req
    for f in dir(a):
        if getattr(getattr(a, f), "__call__", None) is None and (
            len(f) < 2 or f[:2] != "__"
        ):
            eq = getattr(a, f) == getattr(b, f)
            print(f"{f'{ident}{f}: ':<150}{nice_bool(eq)}")
            if not eq:
                if "descriptor" in dir(a):
                    global last_id
                    last_id = a.descriptor
                cmp(getattr(a, f), getattr(b, f), req + 1)


def cmp_dict(a, b, req=0):
    ident = "  " * req
    keys_a = set(a.keys())
    keys_b = set(b.keys())
    if keys_a != keys_b:
        print(f"{ident}a.keys() != b.keys()")
    tot = 0
    nb_failed = 0
    for key in keys_a & keys_b:
        eq = a[key] == b[key]
        tot += 1
        if not eq:
            nb_failed += 1
            print(f"{f'{ident}{str(key)}: ':<150}{nice_bool(eq)}")
            global last_id
            last_id = key
            cmp(a[key], b[key], req + 1)
    print(f"\033[32m{tot-nb_failed}\033[0m + \033[31m{nb_failed}\033[0m = {tot}")


def cmp_list(a, b, req=0):
    ident = "  " * req
    la = len(a)
    lb = len(b)
    if la != lb:
        print(f"{ident}len(a) != len(b)")
    for i in range(min(la, lb)):
        eq = a[i] == b[i]
        print(f"{f'{ident}{str(i)}: ':<150}{nice_bool(eq)}")
        if not eq:
            cmp(a[i], b[i], req + 1)


# apk_eq = new_apk == apk
# print(f"[+] apk are equals: {nice_bool(apk_eq)}")
# if not apk_eq:
#    cmp(new_apk, apk)

# Landroidx/constraintlayout/core/widgets/ConstraintWidget$1;.<clinit>()V
# mid = IdMethod(
#    "<clinit>",
#    IdMethodType(
#        IdType.void(),
#        [],
#    ),
#    IdType("Landroidx/constraintlayout/core/widgets/ConstraintWidget$1;"),
# )
# m = apk.classes[mid.class_].direct_methods[mid]
# nm = new_apk.classes[mid.class_].direct_methods[mid]


# mid = IdMethod(
#    "setValue",
#    IdMethodType(
#        IdType("Z"),
#        [
#            IdType("Ljava/lang/String;"),
#            IdType("Landroidx/constraintlayout/core/parser/CLElement;"),
#        ],
#    ),
#    IdType("Landroidx/constraintlayout/core/state/WidgetFrame;"),
# )
#
# m = apk.classes[mid.class_].virtual_methods[mid]
# nm = new_apk.classes[mid.class_].virtual_methods[mid]
# c = m.code
# nc = nm.code
# cc = c.with_normalized_labels()
# ncc = nc.with_normalized_labels()