#!/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
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):
    value = os.environ.get(name)
    if not value:
        raise RuntimeError(f"{name} is not set; run xteink-emu through its Nix app")
    return Path(value)


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 start(args):
    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)

    command = [
        str(require_env("XTEINK_QEMU")),
        "-machine", "xteink,variant=x3",
        "-accel", "tcg",
        "-L", str(require_env("XTEINK_BIOS")),
        "-display", "none",
        "-serial", f"file:{state / 'serial.log'}",
        "-monitor", "none",
        "-qmp", f"unix:{state / 'qmp.sock'},server=on,wait=off",
        "-nic", "none",
        "-drive", f"file={flash},if=mtd,format=raw",
        "-drive", f"file={sdcard},if=sd,format=raw",
    ]
    qemu_log = (state / "qemu.log").open("ab")
    process = subprocess.Popen(
        command,
        stdin=subprocess.DEVNULL,
        stdout=qemu_log,
        stderr=subprocess.STDOUT,
        start_new_session=True,
    )
    qemu_log.close()
    (state / "pid").write_text(f"{process.pid}\n")

    deadline = time.monotonic() + 10
    while time.monotonic() < deadline:
        if process.poll() is not None:
            message = (state / "qemu.log").read_text(errors="replace")
            raise RuntimeError(f"QEMU exited during startup\n{message}")
        try:
            qmp_command(state, "query-status")
            print(f"started PID {process.pid}")
            print(f"state: {state}")
            print(f"serial log: {state / 'serial.log'}")
            return
        except (FileNotFoundError, ConnectionRefusedError, socket.timeout):
            time.sleep(0.05)
    process.terminate()
    raise RuntimeError("QEMU did not open its QMP socket")


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]
    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())

    command = commands.add_parser("start")
    add_state(command)
    command.add_argument("--firmware", required=True)
    command.add_argument("--sdcard", required=True)
    command.set_defaults(handler=start)

    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)
    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")
    args.handler(args)


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