import argparse
import adbutils
import sys
from pathlib import Path
import os
import tomllib
import tomli_w
import subprocess

def rip(config, d):
    # packages = [x.removeprefix("package:") for x in d.shell("pm list packages -e").splitlines() if "package:" in x]
    # if args.package not in packages:
    #     sys.exit("Could not find your package!")

    package = config["package"]

    version_name = [x.strip().removeprefix("versionName=") for x in d.shell(f"pm dump {package}").splitlines() if x.strip().startswith("versionName")]
    if len(version_name) == 0:
        sys.exit(f"Could not find package {package}")
    
    if len(version_name) > 1:
        sys.exit(f"Found multiple version names: {", ".join(version_name)}")

    version_name = version_name[0]

    print(f"Ripping {package} v{version_name}")

    files = [x.removeprefix("package:") for x in d.shell(f"pm path {package}").splitlines()]
    
    dump_path = Path(version_name) / "_original"
    os.makedirs(dump_path, exist_ok=True)

    for file in files:
        file_base = file.split("/")[-1]
        print(f"Downloading {file_base} ... ", end="", flush=True)
        d.sync.pull(file, dump_path / file_base)
        print(f"✅")


def decompile(config):
    apk_path = Path("_original")
    if not os.path.isdir(apk_path):
        sys.exit("Could not find original APKs in this folder. Did you run 'are rip'?")

    apks = os.listdir(apk_path)
    if not "base.apk" in apks:
        sys.exit("Could not find base.apk")
    
    subprocess.call(["apktool.bat", "d", apk_path / "base.apk", "-o", "src/base"])

def build(config):
    source_dir = Path("src")
    if not os.path.isdir(source_dir):
        sys.exit("Could not find decompiled sources")
    
    out_dir = Path("_dist")
    sources = os.listdir(source_dir)
    for source in sources:
        subprocess.call(["apktool.bat", "b", source_dir / source, "-o", out_dir / f"{source}.apk"])

def sign(config):
    originals_path = Path("_original")
    unsigned_path = Path("_dist")

    if not os.path.isdir(originals_path):
        sys.exit("Could not find original APKs")
    
    if not os.path.isdir(unsigned_path):
        sys.exit("Could not find unsigned dist APKs. Did you run 'are build' yet?")
    
    originals_apks = os.listdir(originals_path)
    unsigned_apks = os.listdir(unsigned_path)
    all_apks = list(set(originals_apks).union(set(unsigned_apks)))


    apk_paths = []
    for apk in all_apks:
        if apk in unsigned_apks:
            apk_paths.append(unsigned_path / apk)
        else:
            apk_paths.append(originals_path / apk)

    command_apks = [ii for i in zip(["-a" for x in range(len(apk_paths))], apk_paths) for ii in i ]
    print(*command_apks)
    subprocess.run(["uber-apk-signer.bat", "--allowResign", *command_apks, "-o", Path("_signed")])

def install(config, d):
    signed_path = Path("_signed")
    signed_apks = [signed_path / x for x in os.listdir(signed_path) if x.endswith(".apk")]

    subprocess.run(["adb", "install-multiple", "-f", *signed_apks])

def init(args):
    config_path = Path("are.toml")

    if os.path.exists(config_path.absolute()):
        sys.exit("A are.toml config file already exists here")
    
    project_data = {
        "package": args.package
    }

    with open(config_path.absolute(), "wb+") as f:
        tomli_w.dump(project_data, f)
    
    print("Project initialized!")

def find_config():
    project_path = Path(".").absolute()
    for i in range(2):
        config_path = project_path / "are.toml"
        if os.path.isfile(config_path):
            with open(config_path, "rb") as f:
                return (project_path, tomllib.load(f))
        else:
            project_path = project_path.parent

    sys.exit("Could not find project file 'are.toml'. Please run 'are init'")


def check():
    try:
        apktool = subprocess.run(["apktool.bat", "--version"], stdout=subprocess.DEVNULL)
    except Exception as e:
        print(e)
        sys.exit("apktool not found")

    try:
        apktool = subprocess.run(["uber-apk-signer.bat", "--version"], stdout=subprocess.DEVNULL)
    except Exception as e:
        print(e)
        sys.exit("uber-apk-signer not found")

    print("✅ All good!")

def main():
    parser = argparse.ArgumentParser(
        prog='are',
        description='Android Reverse Engineering Shortcuts',
    )

    subparsers = parser.add_subparsers(dest='command', required=True)

    # -- check command
    check_parser = subparsers.add_parser('check')

    # -- init command
    init_parser = subparsers.add_parser('init')
    init_parser.add_argument("package", help="The name of the package")

    # -- rip command
    rip_parser = subparsers.add_parser('rip')

    # -- decompile
    decompile_parser = subparsers.add_parser("decompile")

    # -- build
    build_parser = subparsers.add_parser("build")


    # -- sign command
    sign_parser = subparsers.add_parser('sign')

    # -- install command
    install_parser = subparsers.add_parser('install')


    args = parser.parse_args()



    if args.command == "check":
        check()
        return

    config = None
    path = None
    if args.command == "init":
        init(args)
    else:
        path, config = find_config()


    if args.command == "decompile":
        decompile(config)
        return
    elif args.command == "build":
        build(config)
        return
    elif args.command == "sign":
        sign(config)
        return



    adb = adbutils.AdbClient(host="127.0.0.1", port=5037)
    d = adb.device()

    if args.command == "rip":
        rip(config, d)
    
    elif args.command == "install":
        install(config, d)



if __name__ == "__main__":
    main()