#!/usr/bin/env python3

import argparse
import os
import re
import shlex
import subprocess

from tools.helper import bcolors, print_color


COMMON_BRIDGE_PROPS = [
    "ro.product.cpu.abilist=x86_64,arm64-v8a,x86,armeabi-v7a,armeabi",
    "ro.product.cpu.abilist64=x86_64,arm64-v8a",
    "ro.product.cpu.abilist32=x86,armeabi-v7a,armeabi",
    "ro.dalvik.vm.isa.arm=x86",
    "ro.dalvik.vm.isa.arm64=x86_64",
    "ro.enable.native.bridge.exec=1",
    "ro.enable.native.bridge.exec64=1",
    "ro.vendor.enable.native.bridge.exec=1",
    "ro.vendor.enable.native.bridge.exec64=1",
]

DEVICE_PROFILES = {
    "pixel-3-xl": [
        "ro.product.model=Pixel 3 XL",
        "ro.product.brand=google",
        "ro.product.name=crosshatch",
        "ro.product.device=crosshatch",
        "ro.product.manufacturer=Google",
        "ro.product.system.model=Pixel 3 XL",
        "ro.product.system.brand=google",
        "ro.product.system.name=crosshatch",
        "ro.product.system.device=crosshatch",
        "ro.product.system.manufacturer=Google",
        "ro.product.vendor.model=Pixel 3 XL",
        "ro.product.vendor.brand=google",
        "ro.product.vendor.name=crosshatch",
        "ro.product.vendor.device=crosshatch",
        "ro.product.vendor.manufacturer=Google",
        "ro.product.product.model=Pixel 3 XL",
        "ro.product.product.brand=google",
        "ro.product.product.name=crosshatch",
        "ro.product.product.device=crosshatch",
        "ro.product.product.manufacturer=Google",
        "ro.build.fingerprint=google/crosshatch/crosshatch:11/RQ3A.211001.001/7641976:user/release-keys",
        "ro.system.build.fingerprint=google/crosshatch/crosshatch:11/RQ3A.211001.001/7641976:user/release-keys",
        "ro.vendor.build.fingerprint=google/crosshatch/crosshatch:11/RQ3A.211001.001/7641976:user/release-keys",
        "ro.product.build.fingerprint=google/crosshatch/crosshatch:11/RQ3A.211001.001/7641976:user/release-keys",
        "ro.build.type=user",
        "ro.build.tags=release-keys",
        "ro.secure=1",
        "ro.debuggable=0",
        "ro.boot.verifiedbootstate=green",
        "ro.boot.vbmeta.device_state=locked",
        "ro.boot.flash.locked=1",
        "ro.boot.veritymode=enforcing",
    ],
    "pixel-7-pro": [
        "ro.product.model=Pixel 7 Pro",
        "ro.product.brand=google",
        "ro.product.name=cheetah",
        "ro.product.device=cheetah",
        "ro.product.manufacturer=Google",
        "ro.product.system.model=Pixel 7 Pro",
        "ro.product.system.brand=google",
        "ro.product.system.name=cheetah",
        "ro.product.system.device=cheetah",
        "ro.product.system.manufacturer=Google",
        "ro.product.vendor.model=Pixel 7 Pro",
        "ro.product.vendor.brand=google",
        "ro.product.vendor.name=cheetah",
        "ro.product.vendor.device=cheetah",
        "ro.product.vendor.manufacturer=Google",
        "ro.product.product.model=Pixel 7 Pro",
        "ro.product.product.brand=google",
        "ro.product.product.name=cheetah",
        "ro.product.product.device=cheetah",
        "ro.product.product.manufacturer=Google",
        "ro.build.fingerprint=google/cheetah/cheetah:13/TQ3A.230805.001/2023080800:user/release-keys",
        "ro.system.build.fingerprint=google/cheetah/cheetah:13/TQ3A.230805.001/2023080800:user/release-keys",
        "ro.vendor.build.fingerprint=google/cheetah/cheetah:13/TQ3A.230805.001/2023080800:user/release-keys",
        "ro.product.build.fingerprint=google/cheetah/cheetah:13/TQ3A.230805.001/2023080800:user/release-keys",
        "ro.build.type=user",
        "ro.build.tags=release-keys",
        "ro.secure=1",
        "ro.debuggable=0",
        "ro.boot.verifiedbootstate=green",
        "ro.boot.vbmeta.device_state=locked",
        "ro.boot.flash.locked=1",
        "ro.boot.veritymode=enforcing",
    ]
}

LAUNCH_PRESETS = {
    "redroid-pro": {
        "props": [
            "androidboot.redroid_width=720",
            "androidboot.redroid_height=1280",
            "androidboot.redroid_dpi=320",
            "androidboot.redroid_fps=30",
            "androidboot.redroid_vm_heap=256m",
            "ro.config.low_ram=true",
            "config.disable_consumerir=true",
        ],
        "device_profile": "pixel-3-xl",
        "memory": "4g",
    }
}


def infer_bridge(image, requested):
    if requested != "auto":
        return requested
    if "houdini" in image:
        return "libhoudini"
    if "ndk" in image:
        return "libndk"
    return "none"


def bridge_props(mode):
    if mode == "none":
        return []
    props = list(COMMON_BRIDGE_PROPS)
    if mode == "libhoudini":
        props.append("ro.dalvik.vm.native.bridge=libhoudini.so")
    elif mode == "libnb":
        props.append("ro.dalvik.vm.native.bridge=libnb.so")
    elif mode == "libndk":
        props.append("ro.dalvik.vm.native.bridge=libndk_translation.so")
        props.append("ro.ndk_translation.version=0.2.3")
    else:
        raise ValueError("Unsupported bridge mode: {}".format(mode))
    return props


def device_profile_props(profile):
    if profile == "none":
        return []
    try:
        return list(DEVICE_PROFILES[profile])
    except KeyError as exc:
        raise ValueError("Unsupported device profile: {}".format(profile)) from exc


def preset_props(preset):
    if preset == "none":
        return []
    try:
        return list(LAUNCH_PRESETS[preset]["props"])
    except KeyError as exc:
        raise ValueError("Unsupported preset: {}".format(preset)) from exc


def resolve_device_profile(preset, requested):
    if requested != "none":
        return requested
    if preset == "none":
        return "none"
    return LAUNCH_PRESETS[preset].get("device_profile", "none")


def resolve_memory(preset, requested):
    if requested:
        return requested
    if preset == "none":
        return None
    return LAUNCH_PRESETS[preset].get("memory")


def default_name(image):
    name = image.split("/")[-1].replace(":", "_")
    name = re.sub(r"[^a-zA-Z0-9_.-]+", "-", name).strip("-")
    return name or "redroid"


def quote_cmd(args):
    return " ".join(shlex.quote(part) for part in args)


def main():
    parser = argparse.ArgumentParser(
        description="Launch a redroid container with a tested native bridge configuration."
    )
    parser.add_argument(
        "--image",
        required=True,
        help="Image name, for example redroid/redroid:13.0.0_mindthegapps_houdini_magisk",
    )
    parser.add_argument(
        "-c",
        "--container",
        default="docker",
        choices=["docker", "podman"],
        help="Container runtime",
    )
    parser.add_argument(
        "--name",
        help="Container name. Defaults to a slug derived from the image tag.",
    )
    parser.add_argument(
        "--data-dir",
        default="~/data-redroid",
        help="Host directory mapped to /data",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=5555,
        help="Host TCP port mapped to guest ADB 5555",
    )
    parser.add_argument(
        "--bridge",
        choices=["auto", "none", "libhoudini", "libnb", "libndk"],
        default="auto",
        help="Native bridge runtime property to inject",
    )
    parser.add_argument(
        "--preset",
        choices=["none"] + sorted(LAUNCH_PRESETS.keys()),
        default="none",
        help="Inject a tested launch preset",
    )
    parser.add_argument(
        "--gpu-mode",
        default="guest",
        help="Value for androidboot.redroid_gpu_mode",
    )
    parser.add_argument(
        "--device-profile",
        choices=["none"] + sorted(DEVICE_PROFILES.keys()),
        default="none",
        help="Inject a tested device identity profile",
    )
    parser.add_argument(
        "--memory",
        help="Optional container memory limit, for example 4g",
    )
    parser.add_argument(
        "--replace",
        action="store_true",
        help="Remove an existing container with the same name before running",
    )
    parser.add_argument(
        "--keep-container",
        action="store_true",
        help="Do not pass --rm to the runtime",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print the generated command without executing it",
    )
    parser.add_argument(
        "--extra-prop",
        action="append",
        default=[],
        help="Additional redroid property, repeatable",
    )
    args = parser.parse_args()

    image = args.image
    name = args.name or default_name(image)
    data_dir = os.path.abspath(os.path.expanduser(args.data_dir))
    bridge = infer_bridge(image, args.bridge)
    preset = args.preset
    device_profile = resolve_device_profile(preset, args.device_profile)
    memory = resolve_memory(preset, args.memory)

    os.makedirs(data_dir, exist_ok=True)

    cmd = [
        args.container,
        "run",
        "-d",
        "--privileged",
        "--name",
        name,
        "-v",
        "{}:/data".format(data_dir),
        "-p",
        "{}:5555".format(args.port),
    ]
    if memory:
        cmd.extend(["--memory", memory])
    if not args.keep_container:
        cmd.append("--rm")
    cmd.extend(
        [
            image,
            "androidboot.redroid_gpu_mode={}".format(args.gpu_mode),
        ]
    )
    cmd.extend(bridge_props(bridge))
    cmd.extend(preset_props(preset))
    cmd.extend(device_profile_props(device_profile))
    cmd.extend(args.extra_prop)

    if args.replace and not args.dry_run:
        subprocess.run(
            [args.container, "rm", "-f", name],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            check=False,
        )

    print_color("Selected native bridge: {}".format(bridge), bcolors.GREEN)
    if preset != "none":
        print_color("Selected launch preset: {}".format(preset), bcolors.GREEN)
    if device_profile != "none":
        print_color("Selected device profile: {}".format(device_profile), bcolors.GREEN)
    print("Command:\n{}\n".format(quote_cmd(cmd)))

    if args.dry_run:
        return

    subprocess.run(cmd, check=True)
    print_color("Container started as {}".format(name), bcolors.GREEN)
    print("ADB: adb connect 127.0.0.1:{}".format(args.port))


if __name__ == "__main__":
    main()