"""Bounded JSON QMP over private Unix sockets, with reply ID matching."""

import json
import socket
import time
from pathlib import Path


class QMPError(RuntimeError):
    pass


class QMPClient:
    def __init__(self, path: Path, timeout: float = 2.0):
        if timeout <= 0:
            raise ValueError('QMP timeout must be positive')
        self.timeout = timeout
        self.socket = socket.socket(socket.AF_UNIX)
        self.buffer = bytearray()
        self.counter = 0
        deadline = time.monotonic() + timeout
        try:
            self.socket.settimeout(timeout)
            self.socket.connect(str(path))
            greeting = self._receive(deadline)
            # Real QEMU can queue a SPICE event across monitor reconnects.
            # Consume only protocol events, retaining the original deadline.
            while 'event' in greeting:
                if not self._async_event(greeting):
                    raise QMPError('Invalid QMP event before greeting')
                greeting = self._receive(deadline)
            if not isinstance(greeting.get('QMP'), dict):
                raise QMPError('Invalid QMP greeting')
            self._command('qmp_capabilities', None, deadline)
        except (OSError, QMPError) as exc:
            self.close()
            if isinstance(exc, QMPError):
                raise
            raise QMPError(f'QMP connection failed: {exc}') from exc

    @staticmethod
    def _async_event(message: dict) -> bool:
        if not {'event', 'timestamp'} <= message.keys() <= {'event', 'timestamp', 'data'}:
            return False
        if not isinstance(message['event'], str) or not message['event']:
            return False
        if 'data' in message and not isinstance(message['data'], dict):
            return False
        stamp = message['timestamp']
        if (not isinstance(stamp, dict) or set(stamp) != {'seconds', 'microseconds'}
                or any(type(value) is not int for value in stamp.values())):
            return False
        seconds, micros = stamp['seconds'], stamp['microseconds']
        # QMP specifies (-1, -1) when the server cannot retrieve host time.
        return (seconds == -1 and micros == -1) or (seconds >= 0 and 0 <= micros < 1000000)

    def _receive(self, deadline: float) -> dict:
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise QMPError('QMP deadline exceeded')
            if b'\n' in self.buffer:
                line, _, rest = self.buffer.partition(b'\n')
                self.buffer = bytearray(rest)
                try:
                    message = json.loads(line)
                except (ValueError, UnicodeDecodeError) as exc:
                    raise QMPError('Malformed QMP JSON') from exc
                if not isinstance(message, dict):
                    raise QMPError('QMP response must be an object')
                return message
            self.socket.settimeout(remaining)
            try:
                chunk = self.socket.recv(65536)
            except OSError as exc:
                raise QMPError(f'QMP receive failed: {exc}') from exc
            if not chunk:
                raise QMPError('QMP connection closed')
            self.buffer.extend(chunk)
            if len(self.buffer) > 1048576:
                raise QMPError('QMP message exceeds 1 MiB')

    def _command(self, name: str, arguments: dict | None, deadline: float) -> dict:
        self.counter += 1
        request = {'execute': name, 'id': self.counter}
        if arguments is not None:
            request['arguments'] = arguments
        try:
            self.socket.settimeout(max(0.001, deadline - time.monotonic()))
            self.socket.sendall(json.dumps(request).encode() + b'\r\n')
        except OSError as exc:
            raise QMPError(f'QMP send failed: {exc}') from exc
        while True:
            response = self._receive(deadline)
            if response.get('id') != self.counter:
                continue
            if 'error' in response:
                raise QMPError(f'QMP {name}: {response["error"]}')
            if 'return' not in response:
                raise QMPError('QMP reply has no return value')
            return response['return']

    def command(self, name: str, arguments: dict | None = None) -> dict:
        return self._command(name, arguments, time.monotonic() + self.timeout)

    def close(self):
        self.socket.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()