#!/usr/bin/env python3
"""Small operator CLI for the local KaiOS Engmode relay."""

import argparse
import base64
import json
import os
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path


DEFAULT_SERVER = os.environ.get("ENGMODE_RELAY_SERVER", "http://127.0.0.1:9000")
DEFAULT_TOKEN = os.environ.get("ENGMODE_RELAY_TOKEN", "kaios-local-relay")

PROBES = {
    "baseline": (
        "echo '=== identity ==='; "
        "id; id -Z 2>&1; whoami 2>&1; pwd; "
        "echo '=== kernel/build ==='; "
        "uname -a; cat /proc/version 2>&1; "
        "getprop ro.build.fingerprint 2>&1; "
        "getprop ro.build.version.release 2>&1; "
        "getprop ro.build.version.security_patch 2>&1; "
        "echo '=== selinux ==='; "
        "getenforce 2>&1; cat /sys/fs/selinux/enforce 2>&1; "
        "echo '=== mounts ==='; "
        "cat /proc/mounts | grep -E ' / | /data | /system | /sdcard | /storage' 2>&1"
    ),
    "adb": (
        "echo '=== adb props ==='; "
        "getprop init.svc.adbd; getprop ro.adb.secure; getprop ro.secure; "
        "getprop ro.debuggable; getprop service.adb.root; "
        "getprop persist.sys.usb.config; getprop sys.usb.config; getprop sys.usb.state; "
        "echo '=== adb paths ==='; "
        "ls -ldZ /adb_keys /data/misc /data/misc/adb /data/misc/adb/adb_keys 2>&1; "
        "echo '=== adbd process/socket ==='; "
        "ps -AZ 2>&1 | grep -E 'adbd|b2g|system_server' 2>&1; "
        "cat /proc/net/unix | grep -E 'adbd|debugger' 2>&1"
    ),
    "service": (
        "echo '=== service dirs ==='; "
        "ls -ldZ /data/local /data/local/tmp /data/local/service "
        "/data/local/service/api-daemon /data/local/service/updater 2>&1; "
        "echo '=== service files ==='; "
        "ls -lZ /data/local/service/api-daemon /data/local/service/updater 2>&1; "
        "echo '=== processes ==='; "
        "ps -AZ 2>&1 | grep -E 'api|updater|b2g|child|adbd' 2>&1; "
        "echo '=== writable checks ==='; "
        "touch /data/local/service/api-daemon/codex_probe 2>&1; echo api_touch_rc=$?; "
        "ls -lZ /data/local/service/api-daemon/codex_probe 2>&1; "
        "rm /data/local/service/api-daemon/codex_probe 2>&1; echo api_rm_rc=$?; "
        "touch /data/local/service/updater/codex_probe 2>&1; echo updater_touch_rc=$?; "
        "ls -lZ /data/local/service/updater/codex_probe 2>&1; "
        "rm /data/local/service/updater/codex_probe 2>&1; echo updater_rm_rc=$?"
    ),
    "devtools": (
        "echo '=== devtools sockets ==='; "
        "ls -lZ /data/local/debugger-socket /data/local/firefox-debugger-socket 2>&1; "
        "cat /proc/net/unix | grep -E 'debugger|firefox|adbd' 2>&1; "
        "echo '=== tcp listeners raw ==='; "
        "cat /proc/net/tcp 2>&1; echo '--- tcp6 ---'; cat /proc/net/tcp6 2>&1; "
        "echo '=== ip hints ==='; "
        "ip addr 2>&1; ifconfig 2>&1; "
        "echo '=== relevant logcat ==='; "
        "logcat -d 2>&1 | grep -i -E 'devtools|debugger|remote debugger|Starting USB|Starting WiFi|Unable to start|AdbController' | tail -120 2>&1"
    ),
    "b2g-loader": (
        "mkdir -p /sdcard/msc_internal 2>/dev/null; "
        "echo '=== b2g native loader probe ==='; date; "
        "echo '=== identity ==='; id; id -Z 2>&1; cat /proc/self/attr/current 2>&1; pwd; "
        "echo '=== policy-relevant paths ==='; "
        "ls -lZ /system/b2g/b2g /system/b2g/plugin-container /system/b2g/xpcshell "
        "/system/b2g/run-mozilla.sh /system/bin/b2g.sh /system/bin/sh /system/bin/toybox "
        "/data/b2g /data/local/webapps /data/local/service/api-daemon "
        "/data/local/service/api-daemon/remote /data/local/service/updater/updater-daemon "
        "/cache /sdcard /sdcard/msc_internal 2>&1; "
        "echo '=== writable-label test ==='; "
        "for d in /data/b2g /data/local/webapps /data/local/service/api-daemon /data/local/service/api-daemon/remote /cache /sdcard/msc_internal; do "
        "  echo '-- dir:' $d; "
        "  touch $d/codex_b2g_loader_probe 2>&1; echo touch_rc=$?; "
        "  ls -lZ $d/codex_b2g_loader_probe 2>&1; "
        "  rm -f $d/codex_b2g_loader_probe 2>&1; echo rm_rc=$?; "
        "done; "
        "echo '=== executable surface strings ==='; "
        "for f in /system/b2g/plugin-container /system/b2g/b2g /system/b2g/xpcshell /system/b2g/run-mozilla.sh; do "
        "  echo '-- file:' $f; "
        "  file $f 2>&1; "
        "  strings $f 2>/dev/null | grep -i -E 'plugin|xpcshell|gre|profile|omni|gmp|extension|debug|remote|load|\\.so|LD_|MOZ_' | head -80 2>&1; "
        "done; "
        "echo '=== timed execution tests ==='; "
        "run_timed() { "
        "  label=$1; shift; echo '-- run:' $label; "
        "  \"$@\" & pid=$!; "
        "  i=0; while [ $i -lt 20 ]; do kill -0 $pid 2>/dev/null || break; sleep 0.1; i=$((i+1)); done; "
        "  if kill -0 $pid 2>/dev/null; then echo timeout_kill_pid=$pid; kill $pid 2>/dev/null; sleep 0.2; kill -9 $pid 2>/dev/null; fi; "
        "  wait $pid 2>/dev/null; echo run_rc=$?; "
        "}; "
        "run_timed sh-version /system/bin/sh --version; "
        "run_timed plugin-container-noargs /system/b2g/plugin-container; "
        "run_timed plugin-container-help /system/b2g/plugin-container --help; "
        "run_timed plugin-container-version /system/b2g/plugin-container --version; "
        "run_timed xpcshell-direct /system/b2g/xpcshell --help; "
        "echo '=== debugger/profile hints ==='; "
        "ls -lZ /data/local/debugger-socket /data/local/firefox-debugger-socket "
        "/data/b2g/mozilla /data/b2g/mozilla/*.default /data/b2g/mozilla/*.default/prefs.js 2>&1; "
        "cat /proc/net/unix | grep -i -E 'debug|b2g|plugin|gecko|firefox' 2>&1; "
        "echo '=== done ==='"
    ),
}


def request_json(method, url, payload=None):
    data = None
    headers = {}
    if payload is not None:
        data = json.dumps(payload).encode("utf-8")
        headers["Content-Type"] = "application/json"

    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            text = resp.read().decode("utf-8", errors="replace")
            return json.loads(text) if text else {}
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        raise SystemExit(f"HTTP {exc.code} from {url}: {body}") from exc
    except urllib.error.URLError as exc:
        raise SystemExit(f"Could not reach {url}: {exc}") from exc


def cmd_status(args):
    print(json.dumps(request_json("GET", args.server + "/relay/status"), indent=2))


def cmd_queue(args):
    started = time.time()
    payload = {
        "token": args.token,
        "command": args.command,
        "wait_ms": args.wait_ms,
    }
    item = request_json("POST", args.server + "/relay/command", payload)
    print(json.dumps(item, indent=2))
    if args.wait:
        wait_for_result(item.get("id"), args.timeout, started)


def cmd_js(args):
    started = time.time()
    script = Path(args.source).read_text(encoding="utf-8")
    payload = {
        "token": args.token,
        "command": args.name or Path(args.source).name,
        "mode": "js",
        "script": script,
        "wait_ms": args.wait_ms,
    }
    item = request_json("POST", args.server + "/relay/command", payload)
    print(json.dumps(item, indent=2))
    if args.wait:
        wait_for_result(item.get("id"), args.timeout, started)


def cmd_probe(args):
    started = time.time()
    command = PROBES[args.name]
    payload = {
        "token": args.token,
        "command": command,
        "wait_ms": args.wait_ms,
    }
    item = request_json("POST", args.server + "/relay/command", payload)
    print(json.dumps(item, indent=2))
    if args.wait:
        wait_for_result(item.get("id"), args.timeout, started)


def cmd_probes(args):
    for name in sorted(PROBES):
        print(name)


def cmd_clear(args):
    payload = {"token": args.token}
    print(json.dumps(request_json("POST", args.server + "/relay/clear", payload), indent=2))


def cmd_active(args):
    payload = {"token": args.token, "client": args.client}
    print(json.dumps(request_json("POST", args.server + "/relay/active", payload), indent=2))


def queue_command(args, command, wait_ms=5000):
    payload = {
        "token": args.token,
        "command": command,
        "wait_ms": wait_ms,
    }
    return request_json("POST", args.server + "/relay/command", payload)


def queue_agent_command(args, command):
    payload = {
        "token": args.token,
        "command": command,
    }
    return request_json("POST", args.server + "/agent/command", payload)


def cmd_push_b64(args):
    source = Path(args.source)
    if not source.is_file():
        raise SystemExit(f"Source file not found: {source}")

    raw = source.read_bytes()
    encoded = base64.b64encode(raw).decode("ascii")
    remote_b64 = args.target + ".b64"
    chunks = [encoded[i : i + args.chunk_size] for i in range(0, len(encoded), args.chunk_size)]

    commands = [f": > {remote_b64}"]
    for chunk in chunks:
        commands.append(f"printf '%s' '{chunk}' >> {remote_b64}")

    decode = (
        f"(base64 -d {remote_b64} > {args.target} || "
        f"toybox base64 -d {remote_b64} > {args.target}) && "
        f"chmod 755 {args.target} && "
        f"ls -lZ {args.target} 2>&1 && "
        f"(sha256sum {args.target} 2>&1 || md5sum {args.target} 2>&1)"
    )
    commands.append(decode)

    print(f"source={source} bytes={len(raw)} base64={len(encoded)} chunks={len(chunks)}")
    print(f"target={args.target}")
    for command in commands:
        item = queue_command(args, command, wait_ms=args.wait_ms)
        print(f"queued #{item.get('id')}: {command[:90]}{'...' if len(command) > 90 else ''}")


def newest_result_file():
    files = sorted(
        Path.cwd().glob("relay_result_*.txt"),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    return files[0] if files else None


def read_result_id(path):
    try:
        for line in path.read_text(encoding="utf-8", errors="replace").splitlines()[:10]:
            if line.startswith("id="):
                return line.split("=", 1)[1].strip()
    except OSError:
        return None
    return None


def wait_for_result(result_id, timeout, since=0):
    deadline = time.time() + timeout
    result_id = str(result_id)
    print(f"waiting for result id={result_id} ...", file=sys.stderr)
    seen = set()
    while time.time() < deadline:
        for path in sorted(Path.cwd().glob("relay_result_*.txt"), key=lambda p: p.stat().st_mtime, reverse=True):
            if path.stat().st_mtime < since:
                continue
            if path in seen:
                continue
            if read_result_id(path) == result_id:
                print(path.read_text(encoding="utf-8", errors="replace"))
                return
            seen.add(path)
        time.sleep(1)
    raise SystemExit(f"Timed out waiting for result id={result_id}")


def wait_for_agent_result(result_id, timeout, since=0):
    deadline = time.time() + timeout
    result_id = str(result_id)
    print(f"waiting for agent result id={result_id} ...", file=sys.stderr)
    seen = set()
    while time.time() < deadline:
        for path in sorted(Path.cwd().glob("agent_result_*.txt"), key=lambda p: p.stat().st_mtime, reverse=True):
            if path.stat().st_mtime < since:
                continue
            if path in seen:
                continue
            if read_result_id(path) == result_id:
                print(path.read_text(encoding="utf-8", errors="replace"))
                return
            seen.add(path)
        time.sleep(1)
    raise SystemExit(f"Timed out waiting for agent result id={result_id}")


def cmd_latest(args):
    path = newest_result_file()
    if not path:
        raise SystemExit("No relay_result_*.txt files found")
    if args.path:
        print(path)
    else:
        print(path.read_text(encoding="utf-8", errors="replace"))


def newest_agent_result_file():
    files = sorted(
        Path.cwd().glob("agent_result_*.txt"),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    return files[0] if files else None


def cmd_agent_status(args):
    print(json.dumps(request_json("GET", args.server + "/agent/status"), indent=2))


def cmd_agent_queue(args):
    started = time.time()
    payload = {
        "token": args.token,
        "action": args.action,
        "arg": args.arg or "",
    }
    item = request_json("POST", args.server + "/agent/command", payload)
    print(json.dumps(item, indent=2))
    if args.wait:
        wait_for_agent_result(item.get("id"), args.timeout, started)


def cmd_agent_shell(args):
    started = time.time()
    payload = {
        "token": args.token,
        "command": args.command,
    }
    item = request_json("POST", args.server + "/agent/command", payload)
    print(json.dumps(item, indent=2))
    if args.wait:
        wait_for_agent_result(item.get("id"), args.timeout, started)


def cmd_agent_clear(args):
    payload = {"token": args.token}
    print(json.dumps(request_json("POST", args.server + "/agent/clear", payload), indent=2))


def cmd_agent_push_b64(args):
    source = Path(args.source)
    if not source.is_file():
        raise SystemExit(f"Source file not found: {source}")

    raw = source.read_bytes()
    encoded = base64.b64encode(raw).decode("ascii")
    remote_b64 = args.target + ".b64"
    chunks = [encoded[i : i + args.chunk_size] for i in range(0, len(encoded), args.chunk_size)]

    commands = [f": > {remote_b64}"]
    for chunk in chunks:
        commands.append(f"printf '%s' '{chunk}' >> {remote_b64}")

    decode = (
        f"(base64 -d {remote_b64} > {args.target} || "
        f"toybox base64 -d {remote_b64} > {args.target}) && "
        f"chmod 755 {args.target} && "
        f"ls -lZ {args.target} 2>&1 && "
        f"(sha256sum {args.target} 2>&1 || md5sum {args.target} 2>&1)"
    )
    commands.append(decode)

    print(f"source={source} bytes={len(raw)} base64={len(encoded)} chunks={len(chunks)}")
    print(f"target={args.target}")
    for command in commands:
        item = queue_agent_command(args, command)
        print(f"queued #{item.get('id')}: {command[:90]}{'...' if len(command) > 90 else ''}")


def cmd_agent_latest(args):
    path = newest_agent_result_file()
    if not path:
        raise SystemExit("No agent_result_*.txt files found")
    if args.path:
        print(path)
    else:
        print(path.read_text(encoding="utf-8", errors="replace"))


def cmd_load_snippet(args):
    base = args.public_server.rstrip("/") + "/" + args.script.lstrip("/")
    joiner = "&" if "?" in base else "?"
    snippet = (
        "javascript:(function(){"
        "var s=document.createElement('script');"
        f"s.src='{base}{joiner}t='+Date.now();"
        "document.body.appendChild(s)"
        "})()"
    )
    print(snippet)


def build_parser():
    parser = argparse.ArgumentParser(description="Control the local KaiOS Engmode relay")
    parser.add_argument("--server", default=DEFAULT_SERVER, help=f"relay server URL (default: {DEFAULT_SERVER})")
    parser.add_argument("--token", default=DEFAULT_TOKEN, help="relay token")
    sub = parser.add_subparsers(dest="cmd", required=True)

    status = sub.add_parser("status", help="show relay queue status")
    status.set_defaults(func=cmd_status)

    queue = sub.add_parser("queue", help="queue a shell command for the phone")
    queue.add_argument("command", help="shell command to run through Engmode")
    queue.add_argument("--wait-ms", type=int, default=8000, help="phone-side output read timeout")
    queue.add_argument("--wait", action="store_true", help="wait for and print the matching result file")
    queue.add_argument("--timeout", type=int, default=60, help="seconds to wait with --wait")
    queue.set_defaults(func=cmd_queue)

    js = sub.add_parser("js", help="queue a JavaScript probe for the phone relay context")
    js.add_argument("source", help="local JavaScript file to run in the relay page")
    js.add_argument("--name", help="short label shown in relay results")
    js.add_argument("--wait-ms", type=int, default=12000, help="phone-side JavaScript timeout")
    js.add_argument("--wait", action="store_true", help="wait for and print the matching result file")
    js.add_argument("--timeout", type=int, default=60, help="seconds to wait with --wait")
    js.set_defaults(func=cmd_js)

    probe = sub.add_parser("probe", help="queue a preset diagnostic probe")
    probe.add_argument("name", choices=sorted(PROBES), help="probe name")
    probe.add_argument("--wait-ms", type=int, default=12000, help="phone-side output read timeout")
    probe.add_argument("--wait", action="store_true", help="wait for and print the matching result file")
    probe.add_argument("--timeout", type=int, default=75, help="seconds to wait with --wait")
    probe.set_defaults(func=cmd_probe)

    probes = sub.add_parser("probes", help="list preset diagnostic probes")
    probes.set_defaults(func=cmd_probes)

    clear = sub.add_parser("clear", help="clear queued relay commands")
    clear.set_defaults(func=cmd_clear)

    active = sub.add_parser("active", help="set or clear the active relay client")
    active.add_argument("client", nargs="?", help="client id to use; omit to clear")
    active.set_defaults(func=cmd_active)

    push = sub.add_parser("push-b64", help="queue base64 chunks to copy a local file to the phone")
    push.add_argument("source", help="local file to transfer")
    push.add_argument("target", help="target path on the phone")
    push.add_argument("--chunk-size", type=int, default=700, help="base64 characters per queued command")
    push.add_argument("--wait-ms", type=int, default=3000, help="phone-side output read timeout per chunk")
    push.set_defaults(func=cmd_push_b64)

    latest = sub.add_parser("latest", help="print newest relay result")
    latest.add_argument("--path", action="store_true", help="print only the newest result path")
    latest.set_defaults(func=cmd_latest)

    snippet = sub.add_parser("snippet", help="print the javascript: loader snippet")
    snippet.add_argument("--public-server", default="http://192.168.40.252:9000")
    snippet.add_argument("--script", default="install_api_agent.js", help="server-hosted JavaScript file to load")
    snippet.set_defaults(func=cmd_load_snippet)

    agent_status = sub.add_parser("agent-status", help="show native updater-agent status")
    agent_status.set_defaults(func=cmd_agent_status)

    agent_queue = sub.add_parser("agent-queue", help="queue a fixed native updater-agent action")
    agent_queue.add_argument("action", choices=["identity", "service_tree", "remote_tree", "restore_updater", "stop"])
    agent_queue.add_argument("--arg", default="")
    agent_queue.add_argument("--wait", action="store_true", help="wait for and print the matching native-agent result")
    agent_queue.add_argument("--timeout", type=int, default=75, help="seconds to wait with --wait")
    agent_queue.set_defaults(func=cmd_agent_queue)

    agent_shell = sub.add_parser("agent-shell", help="queue a shell command for the native updater-agent")
    agent_shell.add_argument("command", help="shell command to run through the native updater-agent")
    agent_shell.add_argument("--wait", action="store_true", help="wait for and print the matching native-agent result")
    agent_shell.add_argument("--timeout", type=int, default=90, help="seconds to wait with --wait")
    agent_shell.set_defaults(func=cmd_agent_shell)

    agent_clear = sub.add_parser("agent-clear", help="clear queued native updater-agent commands")
    agent_clear.set_defaults(func=cmd_agent_clear)

    agent_push = sub.add_parser("agent-push-b64", help="queue base64 chunks to copy a local file through the native updater-agent")
    agent_push.add_argument("source", help="local file to transfer")
    agent_push.add_argument("target", help="target path on the phone")
    agent_push.add_argument("--chunk-size", type=int, default=700, help="base64 characters per queued command")
    agent_push.set_defaults(func=cmd_agent_push_b64)

    agent_latest = sub.add_parser("agent-latest", help="print newest native updater-agent result")
    agent_latest.add_argument("--path", action="store_true", help="print only the newest result path")
    agent_latest.set_defaults(func=cmd_agent_latest)

    return parser


def main():
    args = build_parser().parse_args()
    args.func(args)


if __name__ == "__main__":
    main()