"""Event-driven serial/QMP support for isolated ARM integration tests."""
import json
import base64
import binascii
import hashlib
import os
from pathlib import Path
import re
import selectors
import socket
import subprocess
import time

PROJECT = Path(__file__).resolve().parents[1]

class VM:
    def __init__(self, work, name, image, initramfs=None, extra=(), system_usb=False):
        self.work = Path(work)
        self.name = name
        self.qmp_path = self.work/(name+'.qmp')
        self.log = (self.work/(name+'.log')).open('wb')
        self.errors = (self.work/(name+'.stderr.log')).open('wb')
        self.data = bytearray()
        self.position = 0
        command = [str(PROJECT/'tools/in-void'), '--isolated-network', 'qemu-system-aarch64',
                   '-machine', 'virt', '-cpu', 'max', '-accel', 'tcg', '-m', '1024', '-smp', '2',
                   '-nodefaults', '-display', 'none', '-serial', 'stdio', '-nic', 'none', '-no-reboot',
                   '-qmp', f'unix:{self.qmp_path},server=on,wait=off',
                   '-kernel', str(PROJECT/'out/kernel/boot/kernel_2712.img'),
                   '-initrd', str(initramfs or PROJECT/'out/fds-initramfs.img'),
                   '-append', 'console=ttyAMA0 rdinit=/init ro quiet loglevel=3',
                   '-drive', f'file={image},if=none,id=system,format=raw,readonly=on',
                   *(['-device', 'qemu-xhci,id=xhci,addr=05.0', '-device', 'usb-storage,id=systemusb,drive=system,bus=xhci.0,port=1'] if system_usb else ['-device', 'virtio-blk-pci,drive=system']), *extra]
        self.child = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=self.errors)
        self.selector = selectors.DefaultSelector()
        self.selector.register(self.child.stdout, selectors.EVENT_READ)
    def expect(self, pattern, timeout=120):
        deadline = time.monotonic()+timeout
        expression = re.compile(pattern, re.MULTILINE | re.DOTALL)
        while True:
            match = expression.search(self.data, self.position)
            if match:
                self.position = match.end()
                return match
            remaining = deadline-time.monotonic()
            if remaining <= 0: raise AssertionError(f'{self.name}: missing {pattern!r}; inspect serial log')
            for key, _ in self.selector.select(remaining):
                chunk = os.read(key.fd, 65536)
                if not chunk: raise AssertionError(f'{self.name}: VM exited before expected output; inspect {self.errors.name}')
                self.data.extend(chunk)
                self.log.write(chunk)
                self.log.flush()
                assert len(self.data) < 4*1024*1024, 'Unexpected serial output flood'
                assert b'Kernel panic' not in self.data and b'FDS_STAGE0_ERROR' not in self.data, 'Guest boot failed'
    def send(self, command):
        self.child.stdin.write(command.encode()+b'\n')
        self.child.stdin.flush()
    def capture(self, command, ok=True, timeout=180):
        """Execute once; retry only the checked serial transfer after log noise."""
        self.send('( '+command+' ) >/tmp/fds-test-response 2>&1; printf "%s" "$?" >/tmp/fds-test-status; printf "\\nVM_SAVED\\n"')
        self.expect(rb'^VM_SAVED\r?\n', timeout=timeout)
        for _ in range(20):
            self.send('printf "\\nVM_BEGIN\\n"; base64 -w0 /tmp/fds-test-response; printf "\\n"; sha256sum /tmp/fds-test-response; cat /tmp/fds-test-status; printf "\\nVM_END\\n"')
            frame=self.expect(rb'^VM_BEGIN\r?\n(.*?)\r?\nVM_END\r?\n').group(1).decode().replace('\r','')
            try:
                encoded,checksum,status=frame.split('\n')
                payload=base64.b64decode(encoded, validate=True)
                if hashlib.sha256(payload).hexdigest()!=checksum.split()[0]:continue
            except (ValueError, binascii.Error):continue
            assert (status=='0')==ok, (command,status,payload)
            return payload.decode().strip()
        raise AssertionError('Repeated serial transfer corruption')
    def qmp(self, execute, arguments=None):
        with socket.socket(socket.AF_UNIX) as control:
            control.settimeout(10)
            control.connect(str(self.qmp_path))
            with control.makefile('rwb') as protocol:
                # A device-deletion event can arrive while a fresh monitor
                # connection is being greeted. It is not a command response.
                while True:
                    greeting = json.loads(protocol.readline())
                    if 'QMP' in greeting: break
                    assert 'event' in greeting, greeting
                def request(name, args):
                    protocol.write((json.dumps({'execute': name, 'arguments': args})+'\n').encode())
                    protocol.flush()
                    while True:
                        result = json.loads(protocol.readline())
                        if 'event' in result: continue
                        assert 'return' in result, result
                        return result['return']
                request('qmp_capabilities', {})
                return request(execute, arguments or {})
    def close(self):
        try:
            if self.child.poll() is None:
                self.qmp('quit')
                self.child.wait(timeout=10)
        finally:
            if self.child.poll() is None:
                self.child.kill()
                self.child.wait()
            self.selector.close()
            self.log.close()
            self.errors.close()
    def __enter__(self): return self
    def __exit__(self, *_): self.close()