"""Verify a source-bound independent kernel in a private QEMU/UEFI guest."""
import argparse
import hashlib
import json
import re
from pathlib import Path
import shutil
import socket
import subprocess
import sys
import time
from PIL import Image

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / 'Tools' / 'CSharpBoot'))
sys.path.insert(0, str(ROOT / 'Tools'))
from ReadImage import inspect_image
from TestBrowserPipeline import _monaco_frame_stats, _monaco_fourth_line_delta


def digest(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def read_json(path):
    return json.loads(Path(path).read_text(encoding='utf-8-sig'))


class Qmp:
    def __init__(self, port):
        self.connection = socket.create_connection(('127.0.0.1', port), timeout=10)
        self.stream = self.connection.makefile('rwb')
        self.number = 0
        greeting = json.loads(self.stream.readline())
        assert 'QMP' in greeting, greeting
        self.call('qmp_capabilities')

    def call(self, command, arguments=None):
        self.number += 1
        item = {'execute': command, 'id': self.number}
        if arguments is not None:
            item['arguments'] = arguments
        self.stream.write(json.dumps(item).encode() + b'\n')
        self.stream.flush()
        while True:
            line = self.stream.readline()
            if not line:
                raise RuntimeError('QEMU closed QMP unexpectedly')
            result = json.loads(line)
            if result.get('id') == self.number:
                if 'error' in result:
                    raise RuntimeError(result['error'])
                return result.get('return')

    def close(self):
        self.stream.close()
        self.connection.close()


def exercise_desktop(qmp, proof, text, browser_required):
    pointer = [640, 400]
    def healthy():
        log = text()
        assert not any(marker in log for marker in ('XFE_NATIVE_FAIL', 'XFE_NATIVE_EXCEPTION', 'XFE_RUNTIME_FATAL',
                                                     'XFE_ZERO_FAILFAST', 'XFE_INDEPENDENT_APPLICATION_FAILED',
                                                     'XFE_BROWSER_PAINT_FAILED', 'Rendering failed:')), log[-4096:]
        return log
    def wait(marker, start=0, seconds=90):
        end = time.monotonic() + seconds
        while time.monotonic() < end:
            log = healthy()
            if marker in log[start:]:
                return
            time.sleep(.1)
        raise AssertionError('Desktop acknowledgement timeout: ' + marker + '\n' + healthy()[-4096:])
    def screenshot(name):
        path = proof / (name + '.ppm')
        qmp.call('screendump', {'filename': str(path)})
        with Image.open(path) as frame:
            frame.convert('RGB').save(proof / (name + '.png'))
        return digest(proof / (name + '.png'))
    def key(chord):
        qmp.call('human-monitor-command', {'command-line': 'sendkey ' + chord})
        time.sleep(.15)
    def type_text(value):
        special = {':': 'shift-semicolon', '.': 'dot', '/': 'slash', '-': 'minus', '_': 'shift-minus', ' ': 'spc'}
        for char in value:
            key(special.get(char, 'shift-' + char.lower() if char.isupper() else char))
    def move(x, y):
        while pointer != [x, y]:
            dx, dy = max(-80, min(80, x-pointer[0])), max(-80, min(80, y-pointer[1]))
            qmp.call('input-send-event', {'events': [
                {'type': 'rel', 'data': {'axis': 'x', 'value': dx}},
                {'type': 'rel', 'data': {'axis': 'y', 'value': dy}}]})
            pointer[0] += dx; pointer[1] += dy
            time.sleep(.025)
    def click(x, y):
        move(x, y)
        for down in (True, False):
            qmp.call('input-send-event', {'events': [{'type': 'btn', 'data': {'down': down, 'button': 'left'}}]})
            time.sleep(.1)
    if 'XFE_DESKTOP_ACTIVE|Terminal' in healthy():
        start = len(text()); key('alt-f4'); wait('XFE_DESKTOP_ACTIVE|Files', start)
    else:
        wait('XFE_DESKTOP_ACTIVE|Files')
    files_hash = screenshot('Desktop-Files')
    start = len(text()); click(370, 116)
    wait('XFE_DESKTOP_ACTIVE|Editor', start)
    type_text('hello')
    start = len(text()); key('ctrl-s')
    wait('XFE_DESKTOP_SAVED|X:\\Users\\Default\\Documents\\Document1.txt|5', start)
    wait('XFE_DESKTOP_SAVED_TEXT_BEGIN\nhello\nXFE_DESKTOP_SAVED_TEXT_END', start)
    editor_hash = screenshot('Desktop-Editor-Saved')
    assert files_hash != editor_hash, 'Editor input produced no changed guest framebuffer'
    start = len(text()); key('alt-f4'); wait('XFE_DESKTOP_ACTIVE|Files', start)
    start = len(text()); click(220, 170); wait('XFE_DESKTOP_ACTIVE|Editor', start)
    start = len(text()); key('ctrl-s')
    wait('XFE_DESKTOP_SAVED|X:\\Users\\Default\\Documents\\Document1.txt|5', start)
    wait('XFE_DESKTOP_SAVED_TEXT_BEGIN\nhello\nXFE_DESKTOP_SAVED_TEXT_END', start)
    start = len(text()); key('alt-f4'); wait('XFE_DESKTOP_ACTIVE|Files', start)
    start = len(text()); key('ctrl-esc'); wait('XFE_DESKTOP_FLAGS|1|0|', start)
    key('down'); key('down'); key('ret')
    wait('XFE_DESKTOP_ACTIVE|Browser', start)
    wait('JavaScript ready: 12', start)
    browser_hash = screenshot('Browser-JavaScript')
    assert browser_hash not in (files_hash, editor_hash), 'Browser did not present a distinct guest frame'
    details = {'MouseCreatedEditor': True, 'KeyboardSavedRamFile': True, 'SavedRamFileReopened': True, 'AltF4ClosedEditor': True,
               'StartMenuKeyboardNavigation': True, 'JavaScriptDomResult': 'JavaScript ready: 12',
               'Screenshots': {'Files': files_hash, 'EditorSaved': editor_hash, 'Browser': browser_hash}}
    if browser_required:
        start = len(text()); key('ctrl-l'); type_text('about:Monaco'); key('ret')
        wait('Official Monaco 0.52.2 - ready', start, seconds=180)
        wait('XFE_BROWSER_PAINT_END', start, seconds=180)
        wait('XFE_MONACO_VALUE_END', start)
        model_pattern = re.compile(r'XFE_MONACO_VALUE_BEGIN\r?\n(.*?)\r?\nXFE_MONACO_VALUE_END', re.S)
        models = model_pattern.findall(text()[start:])
        assert models and 'Editor unavailable' not in models[-1], 'The actual Monaco model is unavailable'
        details['MonacoInitialModel'] = models[-1]
        details['MonacoRasterPaintCompleted'] = True
        windows = re.findall(r'XFE_WINDOW_STATE\|\d+\|Browser\|(-?\d+)\|(-?\d+)\|(\d+)\|(\d+)\|', text())
        assert windows, 'Actual browser window geometry is missing'
        left, top, width, height = map(int, windows[-1])
        viewport = (left + 1, top + 67, width - 2, height - 97)
        move(left + width - 20, top + 50)
        def paint_frame(name, edited=False, before=None):
            end = time.monotonic() + 60
            attempt = 0
            while time.monotonic() < end:
                healthy(); attempt += 1
                screenshot(name + '-' + str(attempt))
                with Image.open(proof / (name + '-' + str(attempt) + '.png')) as pixels:
                    frame = pixels.copy()
                stats = _monaco_frame_stats(frame, viewport, 1)
                if edited:
                    delta = _monaco_fourth_line_delta(before, frame, stats['Lines'][3]['NativeBounds'])
                    stats['FourthLineDelta'] = delta
                    visible = stats['Lines'][3]['HasTextGlyphDistribution'] and delta['NewInkPixels'] >= 35 and delta['NewInkColumns'] >= 24
                else:
                    visible = stats['ReadyPainted']
                if visible:
                    return frame, stats
                time.sleep(.25)
            raise AssertionError('Actual Monaco glyph/syntax pixels did not appear: ' + repr(stats))
        initial_frame, initial_stats = paint_frame('Browser-Monaco-Ready')
        details['MonacoReadyPaint'] = initial_stats
        start = len(text()); click(viewport[0] + 160, viewport[1] + 70)
        key('ctrl-end'); key('ret'); type_text('xfenativeedit')
        wait('Monaco - edited', start)
        end = time.monotonic() + 30
        while time.monotonic() < end:
            models = model_pattern.findall(healthy()[start:])
            if models and 'xfenativeedit' in models[-1]: break
            time.sleep(.1)
        assert models and 'xfenativeedit' in models[-1] and models[-1] != details['MonacoInitialModel'], 'PS/2 input did not update the actual Monaco model'
        details['MonacoEditedModel'] = models[-1]
        move(left + width - 20, top + 50)
        _, edited_stats = paint_frame('Browser-Monaco-Edited', True, initial_frame)
        details['MonacoEditedPaint'] = edited_stats
        details['MonacoKeyboardEditedModel'] = True
        details['Screenshots']['Monaco'] = screenshot('Browser-Monaco')
        details['OfficialMonacoReady'] = True
    # Keep the actual event loop alive after input and verify it responds again.
    time.sleep(2)
    start = len(text()); key('alt-f4'); wait('XFE_DESKTOP_ACTIVE|Files', start)
    screenshot('Desktop-After-Browser-Close')
    start = len(text()); click(220, 248)
    wait('XFE_INDEPENDENT_CONSOLE_IL_OK', start)
    wait('XFE_DESKTOP_ACTIVE|Terminal', start)
    details['ActualConsoleILSamplePassed'] = True
    details['Screenshots']['ConsoleIL'] = screenshot('Terminal-ConsoleIL')
    # Run Demo is a real terminal toolbar action. Its success marker follows
    # actual shared ANSI output and the real external XFEConsole package call.
    log = text()
    windows = re.findall(r'XFE_WINDOW_STATE\|\d+\|Terminal\|(-?\d+)\|(-?\d+)\|', log)
    assert windows, 'Actual terminal window geometry is missing'
    left, top = map(int, windows[-1])
    start = len(text()); click(left + 80, top + 50)
    wait('XFE_INDEPENDENT_CONSOLE_DEMO_OK', start)
    details['ActualAnsiAndXfeConsoleDemoPassed'] = True
    details['Screenshots']['ConsoleDemo'] = screenshot('Terminal-Console-Demo')
    healthy()
    details['NoRenderFailures'] = True
    return details


def verify_boot(kernel_path, cpus, serial_enabled, runtime_required, timeout, desktop_required=False, browser_required=False):
    build = read_json(str(kernel_path) + '.json')
    assert build['Kind'] == 'XfeIndependentKernelBuild' and build['Passed'] is True
    assert build['ContainsCosmos'] is False and build['KernelSha256'] == digest(kernel_path)
    assert build['Target'] == 'x86_64-none'
    if runtime_required:
        assert build['BootstrapOnly'] is False, 'A bootstrap build cannot pass managed-runtime acceptance'
        assert serial_enabled, 'Runtime checks need independent serial evidence'
    if desktop_required:
        assert not build['BootstrapOnly'] and serial_enabled, 'Actual managed desktop acceptance requires a full build and serial evidence'
        assert not runtime_required, 'Runtime-probe acceptance and managed-application acceptance target separate entry points'
    for relative, expected in build['SourceSha256'].items():
        path = (ROOT / relative).resolve()
        assert path.is_relative_to(ROOT) and digest(path) == expected, f'Source changed: {relative}'
    proof = ROOT / 'build' / 'IndependentKernel' / f'Proof-{time.time_ns()}'
    proof.mkdir(parents=True)
    image = proof / 'System.img'
    result = {'Kind': 'XfeIndependentKernelBootProof', 'Passed': False, 'ContainsCosmos': False,
              'VirtualCpus': cpus, 'ExecutingCpus': 1, 'SerialDevice': serial_enabled,
              'RuntimeAcceptanceRequired': runtime_required, 'DesktopAcceptanceRequired': desktop_required,
              'BrowserAcceptanceRequired': browser_required, 'ManagedRuntimeAcceptancePassed': False,
              'DesktopAcceptancePassed': False, 'BrowserAcceptancePassed': False,
              'PhysicalHardwareAcceptancePassed': False, 'FullReleaseQualified': False,
              'KernelSha256': digest(kernel_path), 'BuildReceipt': str(kernel_path) + '.json',
              'BuildReceiptSha256': digest(str(kernel_path) + '.json'),
              'RamOnly': True, 'StorageEnabled': False, 'ProofDirectory': str(proof)}
    qemu = None
    qmp = None
    try:
        with (proof / 'PackImage.log').open('w', encoding='utf-8') as pack_log:
            subprocess.run(['dotnet', 'run', '--project', str(ROOT / 'Tools/CSharpBoot/PackImage.csproj'),
                        '-c', 'Release', '--no-build', '--', str(kernel_path), str(image),
                        '1280x800', '--independent-kernel'], cwd=ROOT, check=True,
                       stdout=pack_log)
        inspected, files = inspect_image(image)
        assert files['System/Kernel.elf'] == kernel_path.read_bytes()
        assert read_json(str(image) + '.json')['KernelImplementation'] == 'XFESystem independent kernel'
        result.update(ImagePath=str(image), ImageSha256=digest(image), ImageVerification=inspected,
                      BootShimSha256=inspected['Files']['EFI/BOOT/BOOTX64.EFI'])
        sources = dict(build['SourceSha256'])
        sources.update(inspected['BootShimSourceSha256'])
        for harness in ('Tools/IndependentKernel/TestBoot.py', 'Tools/TestBrowserPipeline.py', 'Tools/NativeCSharpEvidence.py'):
            sources[harness] = digest(ROOT / harness)
        snapshot = proof / 'SourceSnapshot.json'
        snapshot.write_text(json.dumps(sources, indent=2), encoding='utf-8')
        result.update(SourceSnapshotPath=str(snapshot), SourceSnapshotSha256=digest(snapshot))
        shutil.copy2(ROOT / '.tools/firmware/edk2-i386-vars.fd', proof / 'FirmwareVars.fd')
        with socket.socket() as reservation:
            reservation.bind(('127.0.0.1', 0))
            port = reservation.getsockname()[1]
        command = [str(ROOT / '.tools/qemu/qemu-system-x86_64.exe'), '-machine', 'q35',
                   '-accel', 'tcg', '-cpu', 'max', '-smp', str(cpus), '-m', '1024M',
                   '-vga', 'std', '-global', 'VGA.vgamem_mb=64', '-display', 'none', '-net', 'none', '-rtc', 'base=utc',
                   '-drive', f'if=pflash,format=raw,readonly=on,file={ROOT / ".tools/firmware/edk2-x86_64-code.fd"}',
                   '-drive', f'if=pflash,format=raw,file={proof / "FirmwareVars.fd"}',
                   '-drive', f'file={image},format=raw,if=ide,snapshot=on',
                   '-serial', f'file:{proof / "Serial.log"}' if serial_enabled else 'none',
                   '-qmp', f'tcp:127.0.0.1:{port},server=on,wait=off', '-no-reboot', '-no-shutdown']
        result['QemuArguments'] = command
        with (proof / 'Qemu.log').open('w', encoding='utf-8') as log:
            qemu = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT,
                                    creationflags=getattr(subprocess, 'CREATE_NO_WINDOW', 0))
            deadline = time.monotonic() + timeout
            while qmp is None:
                assert qemu.poll() is None, 'QEMU exited before QMP was available'
                try:
                    qmp = Qmp(port)
                except (OSError, ConnectionError):
                    assert time.monotonic() < deadline, 'QMP connection timeout'
                    time.sleep(0.1)
            ready = False
            serial = ''
            frame = proof / 'Frame.ppm'
            while time.monotonic() < deadline:
                assert qemu.poll() is None, 'Guest/QEMU exited during boot'
                if serial_enabled:
                    serial = (proof / 'Serial.log').read_text(encoding='utf-8', errors='replace')
                    assert not any(marker in serial for marker in ('XFE_NATIVE_FAIL', 'XFE_NATIVE_EXCEPTION', 'XFE_RUNTIME_FATAL',
                                                                   'XFE_ZERO_FAILFAST', 'XFE_INDEPENDENT_APPLICATION_FAILED')), serial[-4096:]
                    ready = 'XFE_NATIVE_SELFTEST_OK' in serial
                    if runtime_required:
                        ready = ready and 'XFE_DOTNET_RUNTIME_PROBE_OK' in serial
                    if desktop_required:
                        ready = ready and 'XFE_INDEPENDENT_DESKTOP_FIRST_FRAME' in serial and 'XFE_APPLICATION_RUNTIME_CHECK_OK' in serial
                else:
                    qmp.call('screendump', {'filename': str(frame)})
                    with Image.open(frame) as rendered:
                        rgb = rendered.convert('RGB')
                        ready = rgb.size == (1280, 800) and all(rgb.getpixel((x, 4)) == (56, 199, 164)
                                                               for x in (0, 320, 640, 960, 1279))
                if ready:
                    break
                time.sleep(0.2)
            assert ready, 'Independent kernel boot/self-test timeout: ' + serial[-4096:]
            if serial_enabled:
                for marker in ('XFE_CPU_OK', 'XFE_MEMORY_MAP_OK', 'XFE_PAGES_OK', 'XFE_LIBC_OK',
                               'XFE_STDIO_SELFTEST_OK', 'XFE_LOADED_ELF_SELFTEST_OK', 'XFE_THREADS_OK', 'XFE_INPUT_READY',
                               'XFE_DESCRIPTORS_OK', 'XFE_VIRTUAL_MEMORY_OK', 'XFE_CLOCK_OK',
                               'XFE_FRAMEBUFFER_OK', 'XFE_NATIVE_SELFTEST_OK'):
                    assert marker in serial, marker
                result['GuestMarkers'] = [line for line in serial.splitlines() if line.startswith('XFE_')]
            qmp.call('screendump', {'filename': str(frame)})
            with Image.open(frame) as rendered:
                rgb = rendered.convert('RGB')
                assert rgb.size == (1280, 800)
                rgb.save(proof / 'Frame.png')
                result['ScreenshotSha256'] = digest(proof / 'Frame.png')
                result['FramebufferSize'] = list(rgb.size)
            qmp.call('query-status')
            if desktop_required:
                result['DesktopInputProof'] = exercise_desktop(qmp, proof,
                    lambda: (proof / 'Serial.log').read_text(encoding='utf-8', errors='replace'), browser_required)
                result['DesktopAcceptancePassed'] = True
                result['BrowserAcceptancePassed'] = browser_required
            result['ManagedRuntimeAcceptancePassed'] = runtime_required or desktop_required
            result['Passed'] = True
    except Exception as error:
        result['Failure'] = str(error)
    finally:
        if qemu is not None:
            if qemu.poll() is None and qmp is not None:
                try:
                    qmp.call('quit')
                except (OSError, RuntimeError):
                    pass
            try:
                qemu.wait(timeout=10)
            except subprocess.TimeoutExpired:
                qemu.terminate()
                qemu.wait(timeout=10)
                result['ForcedCleanup'] = True
            result['QemuExitCode'] = qemu.returncode
        if qmp is not None:
            qmp.close()
        if image.exists():
            result['ImageUnchanged'] = digest(image) == result.get('ImageSha256')
        result['Passed'] = result['Passed'] and result.get('QemuExitCode') == 0 and result.get('ImageUnchanged') is True
        (proof / 'Result.json').write_text(json.dumps(result, indent=2), encoding='utf-8')
    print(json.dumps(result, ensure_ascii=True))
    return result['Passed']


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('kernel', type=Path)
    parser.add_argument('--virtual-cpus', type=int, choices=(1, 8), default=1)
    parser.add_argument('--no-serial', action='store_true')
    parser.add_argument('--runtime', action='store_true')
    parser.add_argument('--desktop', action='store_true')
    parser.add_argument('--browser', action='store_true')
    parser.add_argument('--timeout', type=int, default=90)
    args = parser.parse_args()
    raise SystemExit(0 if verify_boot(args.kernel.resolve(), args.virtual_cpus, not args.no_serial,
                                     args.runtime, args.timeout, args.desktop or args.browser, args.browser) else 1)