#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = []
# ///

import argparse
import json
import os
import re
import signal
import socket
import subprocess
import sys
import tempfile
import time
from pathlib import Path

FLASH_SIZE = 16 * 1024 * 1024
APP_OFFSET = 0x10000
PARTITIONS_OFFSET = 0x8000
RELEASED_ADC = 4095
PROJECT_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_PATHS = {
    "XTEINK_QEMU": PROJECT_ROOT / "dist/qemu-native/bin/qemu-system-riscv32",
    "XTEINK_BIOS": PROJECT_ROOT / "dist/qemu-native/share/qemu",
    "XTEINK_BOOTLOADER": PROJECT_ROOT / "scripts/assets/esp32c3-bootloader.bin",
    "XTEINK_PARTITIONS": PROJECT_ROOT / "scripts/assets/xteink-partitions.bin",
}
BUTTONS = {
    "left": ("/machine/adc", "adci[2]", 2242, RELEASED_ADC),
    "right": ("/machine/adc", "adci[2]", 5, RELEASED_ADC),
    "bottom-1": ("/machine/adc", "adci[1]", 3512, RELEASED_ADC),
    "bottom-2": ("/machine/adc", "adci[1]", 2694, RELEASED_ADC),
    "bottom-3": ("/machine/adc", "adci[1]", 1493, RELEASED_ADC),
    "bottom-4": ("/machine/adc", "adci[1]", 5, RELEASED_ADC),
    "power": ("/machine/gpio", "input-level[3]", False, True),
}


def default_state_dir():
    return Path(tempfile.gettempdir()) / f"xteink-emu-{os.getuid()}"


def state_dir(args):
    return Path(args.state_dir).expanduser().resolve()


def require_env(name):
    path = Path(os.environ.get(name, DEFAULT_PATHS[name]))
    if not path.exists():
        raise RuntimeError(f"{path} does not exist; run 'make xteink'")
    return path


def pid_from(state):
    try:
        return int((state / "pid").read_text().strip())
    except (FileNotFoundError, ValueError):
        return None


def process_alive(pid):
    if pid is None:
        return False
    try:
        os.kill(pid, 0)
        return True
    except ProcessLookupError:
        return False
    except PermissionError:
        return True


def qmp_command(state, execute, arguments=None, timeout=5):
    with socket.socket(socket.AF_UNIX) as sock:
        sock.settimeout(timeout)
        sock.connect(str(state / "qmp.sock"))
        stream = sock.makefile("rwb")
        json.loads(stream.readline())

        def send(command):
            stream.write(json.dumps(command).encode() + b"\n")
            stream.flush()
            while True:
                response = json.loads(stream.readline())
                if "event" in response:
                    continue
                if "error" in response:
                    raise RuntimeError(response["error"]["desc"])
                return response.get("return")

        send({"execute": "qmp_capabilities"})
        command = {"execute": execute}
        if arguments is not None:
            command["arguments"] = arguments
        return send(command)


def require_running(state):
    pid = pid_from(state)
    if not process_alive(pid):
        raise RuntimeError(f"xteink-emu is not running in {state}")
    return pid


def build_flash(firmware, output):
    app = firmware.read_bytes()
    if len(app) == FLASH_SIZE:
        output.write_bytes(app)
        return
    if not app or app[0] != 0xE9:
        raise RuntimeError(f"{firmware} is not an ESP32 app image")
    if APP_OFFSET + len(app) > FLASH_SIZE:
        raise RuntimeError(f"{firmware} does not fit in a 16 MiB flash image")

    flash = bytearray(b"\xff" * FLASH_SIZE)
    bootloader = require_env("XTEINK_BOOTLOADER").read_bytes()
    partitions = require_env("XTEINK_PARTITIONS").read_bytes()
    flash[: len(bootloader)] = bootloader
    flash[PARTITIONS_OFFSET : PARTITIONS_OFFSET + len(partitions)] = partitions
    flash[APP_OFFSET : APP_OFFSET + len(app)] = app
    output.write_bytes(flash)


def run(args, interactive):
    state = state_dir(args)
    old_pid = pid_from(state)
    if process_alive(old_pid):
        raise RuntimeError(f"xteink-emu is already running as PID {old_pid} in {state}")

    firmware = Path(args.firmware).expanduser().resolve()
    sdcard = Path(args.sdcard).expanduser().resolve()
    if not firmware.is_file():
        raise RuntimeError(f"firmware not found: {firmware}")
    if not sdcard.is_file():
        raise RuntimeError(f"SD card image not found: {sdcard}")

    state.mkdir(parents=True, exist_ok=True, mode=0o700)
    for name in ("pid", "qmp.sock", "serial.log", "qemu.log", "wait.offset"):
        (state / name).unlink(missing_ok=True)
    (state / "serial.log").touch()
    (state / "wait.offset").write_text("0\n")
    flash = state / "flash.bin"
    build_flash(firmware, flash)
    efuse = bytearray(1024)
    efuse[0x18:0x1E] = bytes.fromhex("563412005452")
    efuse[0x26] = 0x0C
    (state / "efuse.bin").write_bytes(efuse)

    command = [
        str(require_env("XTEINK_QEMU")),
        "-machine", "xteink,variant=x3",
        "-accel", "tcg",
        "-L", str(require_env("XTEINK_BIOS")),
        "-display", "gtk" if interactive else "none",
        "-serial", "stdio",
        "-monitor", "none",
        "-qmp", f"unix:{state / 'qmp.sock'},server=on,wait=off",
        "-nic", "user,model=esp32c3_wifi,mac=52:54:00:12:34:56,net=192.168.4.0/24",
        "-drive", f"file={flash},if=mtd,format=raw",
        "-drive", f"file={state / 'efuse.bin'},if=none,format=raw,id=efuse",
        "-global", "driver=nvram.esp32c3.efuse,property=drive,value=efuse",
        "-drive", f"file={sdcard},if=sd,format=raw",
    ]
    if interactive:
        process = subprocess.Popen(command)
    else:
        with (state / "serial.log").open("ab") as serial_log:
            process = subprocess.Popen(
                command,
                stdin=subprocess.DEVNULL,
                stdout=serial_log,
                stderr=subprocess.STDOUT,
                start_new_session=True,
            )
    (state / "pid").write_text(f"{process.pid}\n")

    deadline = time.monotonic() + 10
    while time.monotonic() < deadline:
        if process.poll() is not None:
            message = (state / "serial.log").read_text(errors="replace")
            raise RuntimeError(f"QEMU exited during startup\n{message}")
        try:
            qmp_command(state, "query-status")
            if args.power_hold_ms:
                path, prop, pressed, released = BUTTONS["power"]
                qmp_command(state, "qom-set", {"path": path, "property": prop, "value": pressed})
                try:
                    time.sleep(args.power_hold_ms / 1000)
                finally:
                    qmp_command(state, "qom-set", {"path": path, "property": prop, "value": released})
            break
        except (FileNotFoundError, ConnectionRefusedError, socket.timeout):
            time.sleep(0.05)
    else:
        process.terminate()
        raise RuntimeError("QEMU did not open its QMP socket")

    print(f"started PID {process.pid}")
    print(f"state: {state}")
    if not interactive:
        print(f"serial log: {state / 'serial.log'}")
        return

    print("keys: Left/Right, 1-4 bottom buttons, P power; close the window to stop")
    try:
        returncode = process.wait()
    except KeyboardInterrupt:
        if process.poll() is None:
            try:
                qmp_command(state, "quit")
            except OSError:
                process.terminate()
        returncode = process.wait()
    finally:
        (state / "pid").unlink(missing_ok=True)
        (state / "qmp.sock").unlink(missing_ok=True)
    if returncode:
        raise RuntimeError(f"QEMU exited with status {returncode}")


def start(args):
    run(args, interactive=False)


def interactive(args):
    run(args, interactive=True)


def stop(args):
    state = state_dir(args)
    pid = require_running(state)
    try:
        qmp_command(state, "quit")
    except (ConnectionError, FileNotFoundError):
        os.kill(pid, signal.SIGTERM)

    deadline = time.monotonic() + 5
    while process_alive(pid) and time.monotonic() < deadline:
        time.sleep(0.05)
    if process_alive(pid):
        os.kill(pid, signal.SIGKILL)
    (state / "pid").unlink(missing_ok=True)
    (state / "qmp.sock").unlink(missing_ok=True)
    print(f"stopped PID {pid}")


def wait_for_log(args):
    state = state_dir(args)
    require_running(state)
    log_path = state / "serial.log"
    offset_path = state / "wait.offset"
    try:
        pattern = re.compile(args.match)
    except re.error as error:
        raise RuntimeError(f"invalid regex: {error}") from error
    try:
        offset = int(offset_path.read_text().strip())
    except (FileNotFoundError, ValueError):
        offset = 0

    deadline = time.monotonic() + args.timeout
    with log_path.open("r", errors="replace") as log:
        log.seek(offset)
        while True:
            line = log.readline()
            if line:
                if pattern.search(line):
                    offset_path.write_text(f"{log.tell()}\n")
                    print(line, end="" if line.endswith("\n") else "\n")
                    return
                continue
            if time.monotonic() >= deadline:
                raise RuntimeError(f"timed out after {args.timeout:g}s waiting for /{args.match}/")
            if not process_alive(pid_from(state)):
                raise RuntimeError("QEMU exited while waiting for serial output")
            time.sleep(0.05)


def button(args):
    state = state_dir(args)
    require_running(state)
    path, prop, pressed, released = BUTTONS[args.name]
    if args.down:
        qmp_command(state, "qom-set", {"path": path, "property": prop, "value": pressed})
        return
    if args.up:
        qmp_command(state, "qom-set", {"path": path, "property": prop, "value": released})
        return
    qmp_command(state, "qom-set", {"path": path, "property": prop, "value": pressed})
    try:
        time.sleep(args.hold_ms / 1000)
    finally:
        qmp_command(state, "qom-set", {"path": path, "property": prop, "value": released})


def screenshot(args):
    state = state_dir(args)
    require_running(state)
    output = Path(args.output).expanduser().resolve()
    output.parent.mkdir(parents=True, exist_ok=True)
    qmp_command(state, "screendump", {"filename": str(output), "format": "png"}, timeout=30)
    print(output)


def parser():
    root = argparse.ArgumentParser(prog="xteink-emu")
    commands = root.add_subparsers(dest="command", required=True)

    def add_state(command):
        command.add_argument("--state-dir", default=default_state_dir())

    def add_run_arguments(command):
        add_state(command)
        command.add_argument("--firmware", required=True)
        command.add_argument("--sdcard", required=True)
        command.add_argument("--power-hold-ms", type=int, default=2000)

    command = commands.add_parser("start")
    add_run_arguments(command)
    command.set_defaults(handler=start)

    command = commands.add_parser("interactive")
    add_run_arguments(command)
    command.set_defaults(handler=interactive)

    command = commands.add_parser("stop")
    add_state(command)
    command.set_defaults(handler=stop)

    command = commands.add_parser("wait")
    add_state(command)
    command.add_argument("--timeout", type=float, required=True, help="seconds")
    command.add_argument("--match", required=True, help="regular expression")
    command.set_defaults(handler=wait_for_log)

    command = commands.add_parser("button")
    add_state(command)
    command.add_argument("name", choices=BUTTONS)
    command.add_argument("--hold-ms", type=int, default=100)
    button_mode = command.add_mutually_exclusive_group()
    button_mode.add_argument("--down", action="store_true", help="press and hold")
    button_mode.add_argument("--up", action="store_true", help="release")
    command.set_defaults(handler=button)

    command = commands.add_parser("screenshot")
    add_state(command)
    command.add_argument("--output", required=True)
    command.set_defaults(handler=screenshot)
    return root


def main():
    args = parser().parse_args()
    if getattr(args, "timeout", 0) < 0:
        raise RuntimeError("--timeout must not be negative")
    if getattr(args, "hold_ms", 0) < 0:
        raise RuntimeError("--hold-ms must not be negative")
    if getattr(args, "power_hold_ms", 0) < 0:
        raise RuntimeError("--power-hold-ms must not be negative")
    args.handler(args)


if __name__ == "__main__":
    try:
        main()
    except (RuntimeError, OSError) as error:
        print(f"xteink-emu: {error}", file=sys.stderr)
        sys.exit(1)