#!/usr/bin/env python3

import argparse
import os
import re
import shlex
import socket
import subprocess
import sys
import time

import tools.compat as compat
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 wait_tcp_port(host, port, timeout_sec, interval=1.0):
    """Return True once TCP connect succeeds, False on timeout."""
    deadline = time.monotonic() + timeout_sec
    while time.monotonic() < deadline:
        try:
            with socket.create_connection((host, port), timeout=2.0):
                return True
        except OSError:
            time.sleep(interval)
    return False


def local_post_diagnose(container_runtime, name, port, settle_sec):
    """On the Docker host: container status, logs, adb, short logcat."""
    print_color(
        "\n========== --diagnose (local): waiting {}s ==========".format(settle_sec),
        bcolors.YELLOW,
    )
    time.sleep(settle_sec)
    serial = "127.0.0.1:{}".format(port)
    print_color("--- {} ps -a ---".format(container_runtime), bcolors.GREEN)
    subprocess.run(
        [container_runtime, "ps", "-a", "--filter", "name={}".format(name), "--no-trunc"],
        check=False,
    )
    print_color("--- {} inspect ---".format(container_runtime), bcolors.GREEN)
    subprocess.run(
        [
            container_runtime,
            "inspect",
            "--format",
            "status={{.State.Status}} exit={{.State.ExitCode}} err={{.State.Error}}",
            name,
        ],
        check=False,
    )
    print_color("--- {} logs (tail 120) ---".format(container_runtime), bcolors.GREEN)
    subprocess.run([container_runtime, "logs", "--tail", "120", name], check=False)
    print_color("--- adb ---", bcolors.GREEN)
    subprocess.run(["adb", "connect", serial], check=False)
    subprocess.run(["adb", "devices", "-l"], check=False)
    subprocess.run(["adb", "-s", serial, "shell", "getprop", "sys.boot_completed"], check=False)
    print_color("--- adb logcat (display/GLES, last 80 lines) ---", bcolors.GREEN)
    subprocess.run(
        [
            "adb",
            "-s",
            serial,
            "logcat",
            "-d",
            "-t",
            "80",
            "*:S",
            "SurfaceFlinger:V",
            "EGL:V",
            "OpenGLRenderer:V",
            "redroid:V",
            "AndroidRuntime:E",
        ],
        check=False,
    )
    print_color("========== end --diagnose ==========\n", bcolors.YELLOW)


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=(
            "androidboot.redroid_gpu_mode: default guest (stable on headless). "
            "auto=guest unless REDROID_USE_HOST_GPU=1 and DRI exists, then host. "
            "host=GPU (adds --device per DRI node when present)"
        ),
    )
    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(
        "--rm",
        dest="ephemeral",
        action="store_true",
        help="Ephemeral: delete container when it stops (disables auto-restart; not for production)",
    )
    parser.add_argument(
        "--restart",
        dest="restart_policy",
        default="unless-stopped",
        choices=("no", "always", "unless-stopped", "on-failure"),
        help=(
            "Docker/Podman restart policy when not using --rm (default: unless-stopped = "
            "start after daemon reboot unless you docker stop'd it)"
        ),
    )
    parser.add_argument(
        "--keep-container",
        action="store_true",
        help="Legacy: containers are persistent by default; this flag is a no-op",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print the generated command without executing it",
    )
    parser.add_argument(
        "--soft-display",
        action="store_true",
        help=(
            "Lower software-GLES load: 720x1280, fps=15, dpi=280 (helps black screen / flicker with guest GPU)"
        ),
    )
    parser.add_argument(
        "--extra-prop",
        action="append",
        default=[],
        help="Additional redroid property, repeatable",
    )
    parser.add_argument(
        "--adb-insecure",
        action="store_true",
        help=(
            "Headless ADB: early androidboot.adb.secure=0, ro.adb.secure=0, ro.secure=0, ro.debuggable=1 "
            "(before device profile; see redroid-doc README / #579). Overrides user-build secure props."
        ),
    )
    parser.add_argument(
        "--no-compat-check",
        action="store_true",
        help="Skip compatibility messages and docker/podman probe (still resolves --gpu-mode auto from DRI)",
    )
    parser.add_argument(
        "--wait-tcp",
        type=int,
        default=0,
        metavar="SEC",
        help=(
            "After start, wait up to SEC seconds for 127.0.0.1:PORT to accept TCP (Android adbd). "
            "If it never opens, print container logs. Use 90 when debugging failed adb connect."
        ),
    )
    parser.add_argument(
        "--diagnose",
        action="store_true",
        help="After start: wait, then print docker ps/inspect/logs + adb + logcat (same host as Docker)",
    )
    parser.add_argument(
        "--diagnose-settle",
        type=int,
        default=15,
        metavar="SEC",
        help="Seconds to sleep before --diagnose probes (default 15)",
    )
    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)

    run_report, effective_gpu_mode, dri_device_args = compat.analyze_for_run(
        args.container,
        args.gpu_mode,
        skip_check=args.no_compat_check,
    )
    if not args.no_compat_check:
        compat.emit_compat_report(run_report)
        if run_report.errors:
            sys.exit(1)

    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 args.ephemeral:
        cmd.append("--rm")
    elif args.restart_policy != "no":
        cmd.extend(["--restart", args.restart_policy])
    cmd.extend(dri_device_args)
    cmd.append(image)
    # Init applies read-only props in order; later ro.adb.secure=0 is ignored if adb.secure
    # was already set (redroid-doc #579). Put these immediately after the image name.
    if args.adb_insecure:
        # Must precede device_profile_props (ro.secure=1 / ro.debuggable=0): init keeps first ro.* value.
        # redroid-doc README: ro.secure=0 yields working adb shell without RSA prompt.
        cmd.extend(
            [
                "androidboot.adb.secure=0",
                "ro.adb.secure=0",
                "ro.secure=0",
                "ro.debuggable=1",
            ]
        )
    cmd.append("androidboot.redroid_gpu_mode={}".format(effective_gpu_mode))
    cmd.extend(bridge_props(bridge))
    cmd.extend(preset_props(preset))
    cmd.extend(device_profile_props(device_profile))
    if args.soft_display:
        cmd.extend(
            [
                "androidboot.redroid_width=720",
                "androidboot.redroid_height=1280",
                "androidboot.redroid_fps=15",
                "androidboot.redroid_dpi=280",
            ]
        )
    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 GPU mode: {} (androidboot.redroid_gpu_mode)".format(effective_gpu_mode),
        bcolors.GREEN,
    )
    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)
    if args.soft_display:
        print_color("Soft display: 720x1280 @ 15fps (reduced GLES load)", bcolors.GREEN)
    if args.adb_insecure:
        print_color(
            "ADB insecure: injecting boot props immediately after image name "
            "(check Command: line starts with androidboot.adb.secure=0 …).",
            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)
    if args.ephemeral:
        print_color("Ephemeral (--rm): container is removed when it stops.", bcolors.YELLOW)
    else:
        print_color(
            "Persistent: restart policy is {!r} (survives Docker daemon restart; "
            "use `docker stop {}` to opt out of auto-start).".format(args.restart_policy, name),
            bcolors.GREEN,
        )
    if args.diagnose:
        local_post_diagnose(args.container, name, args.port, args.diagnose_settle)
    if args.wait_tcp > 0:
        print_color(
            "Waiting up to {}s for 127.0.0.1:{} (adbd TCP) ...".format(args.wait_tcp, args.port),
            bcolors.GREEN,
        )
        if not wait_tcp_port("127.0.0.1", args.port, args.wait_tcp):
            print_color(
                "Port {} never opened. Common with androidboot.redroid_gpu_mode=host when EGL/GPU fails inside "
                "the guest — try: python run_redroid.py ... --gpu-mode guest (and drop DRI if needed).".format(
                    args.port
                ),
                bcolors.RED,
            )
            print_color("Last logs from container {!r}:".format(name), bcolors.YELLOW)
            subprocess.run([args.container, "logs", "--tail", "100", name], check=False)
            print(
                "\nIf the container already exited with --rm, logs are gone; re-run without --rm "
                "for a persistent container."
            )
    print("ADB: adb connect 127.0.0.1:{}".format(args.port))
    if effective_gpu_mode == "host":
        print(
            "Note: host GPU mode can prevent full boot; if connect fails, use --gpu-mode guest or "
            "--wait-tcp 90 to diagnose."
        )
    if not args.adb_insecure:
        print(
            "If the device shows as unauthorized, stop the container and re-run with --adb-insecure "
            "(headless instances cannot tap the USB debugging authorization dialog)."
        )
    else:
        print(
            "If adb still shows unauthorized: stop the container, delete host folder "
            "{}/misc/adb (persisted keys from an older boot), then start again. "
            "Also try: adb kill-server && adb start-server".format(data_dir)
        )


if __name__ == "__main__":
    main()