"""
MVGAL Python Bindings
=====================
Communicates with mvgal daemon via Unix domain socket using the IPC protocol.

Usage:
    client = MvgalClient()
    devices = client.query_devices()
    print(devices)
    client.close()
"""

import json
import socket
import struct
from typing import Any, Dict, List, Optional


class MvgalError(Exception):
    """Base exception for MVGAL operations."""
    pass


class MvgalConnectionError(MvgalError):
    """Connection-related errors."""
    pass


class MvgalProtocolError(MvgalError):
    """Protocol-level errors."""
    pass


class MvgalClient:
    """Client for communicating with the mvgal daemon over IPC."""

    DEFAULT_SOCKET_PATH = "/run/mvgal/mvgal.sock"
    IPC_MAGIC = 0x4D564741
    IPC_VERSION = 1

    # Message types
    MSG_QUERY_DEVICES = 0x0001
    MSG_SUBMIT_WORKLOAD = 0x0002
    MSG_ALLOC_MEMORY = 0x0003
    MSG_FREE_MEMORY = 0x0004
    MSG_QUERY_CAPABILITIES = 0x0005
    MSG_SET_STRATEGY = 0x0006
    MSG_SUBSCRIBE_TELEMETRY = 0x0007
    MSG_TELEMETRY = 0x0008

    HEADER_FORMAT = "<IHHII"
    HEADER_SIZE = struct.calcsize(HEADER_FORMAT)

    def __init__(self, socket_path: Optional[str] = None):
        """Connect to the mvgal daemon.

        Args:
            socket_path: Path to the Unix domain socket.
                         Defaults to /run/mvgal/mvgal.sock.

        Raises:
            MvgalConnectionError: If connection fails.
        """
        self._socket_path = socket_path or self.DEFAULT_SOCKET_PATH
        self._seq = 0
        self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        try:
            self._sock.connect(self._socket_path)
            self._sock.settimeout(30.0)
        except (FileNotFoundError, ConnectionRefusedError, OSError) as e:
            raise MvgalConnectionError(
                f"Cannot connect to mvgal at {self._socket_path}: {e}"
            ) from e

    def __enter__(self):
        return self

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

    def close(self) -> None:
        """Close the connection to the daemon."""
        if self._sock:
            try:
                self._sock.close()
            except OSError:
                pass
            self._sock = None  # type: ignore

    # ── Low-level send/receive ──────────────────────────────────────────

    def _send_message(self, msg_type: int, payload: Any = None) -> Dict[str, Any]:
        """Send a message and receive the response.

        Args:
            msg_type: One of the MSG_* constants.
            payload: JSON-serializable payload, or None for empty.

        Returns:
            Parsed JSON response dict.

        Raises:
            MvgalConnectionError: On send/recv failures.
            MvgalProtocolError: On invalid responses.
        """
        if not self._sock:
            raise MvgalConnectionError("Not connected")

        self._seq += 1
        payload_bytes = b""
        if payload is not None:
            payload_bytes = json.dumps(payload, separators=(",", ":")).encode("utf-8")

        # Build 16-byte header
        header = struct.pack(
            self.HEADER_FORMAT,
            self.IPC_MAGIC,           # I - uint32 magic
            self.IPC_VERSION,         # H - uint16 version
            msg_type,                 # H - uint16 msg type
            len(payload_bytes),       # I - uint32 payload length
            self._seq,                # I - uint32 sequence number
        )

        try:
            self._sock.sendall(header)
            if payload_bytes:
                self._sock.sendall(payload_bytes)
        except OSError as e:
            raise MvgalConnectionError(f"Send failed: {e}") from e

        # Read response header (16 bytes)
        try:
            resp_header = self._recv_exact(self.HEADER_SIZE)
        except OSError as e:
            raise MvgalConnectionError(f"Recv header failed: {e}") from e

        magic, version, resp_type, payload_len, seq = struct.unpack(
            self.HEADER_FORMAT, resp_header
        )

        if magic != self.IPC_MAGIC:
            raise MvgalProtocolError(
                f"Bad magic: 0x{magic:08X} (expected 0x{self.IPC_MAGIC:08X})"
            )

        # Read payload
        if payload_len > 0:
            try:
                payload_data = self._recv_exact(payload_len)
            except OSError as e:
                raise MvgalConnectionError(f"Recv payload failed: {e}") from e
            try:
                return json.loads(payload_data.decode("utf-8"))
            except (json.JSONDecodeError, UnicodeDecodeError) as e:
                raise MvgalProtocolError(f"Invalid JSON response: {e}") from e

        return {}

    def _recv_exact(self, size: int) -> bytes:
        """Read exactly `size` bytes from the socket."""
        buf = bytearray()
        while len(buf) < size:
            chunk = self._sock.recv(size - len(buf))
            if not chunk:
                raise OSError("Connection closed by peer")
            buf.extend(chunk)
        return bytes(buf)

    # ── High-level API ──────────────────────────────────────────────────

    def query_devices(self) -> List[Dict[str, Any]]:
        """Query available GPU devices.

        Returns:
            List of device information dicts.
        """
        result = self._send_message(self.MSG_QUERY_DEVICES)
        return result.get("devices", [])

    def query_capabilities(self) -> Dict[str, Any]:
        """Query the daemon's capabilities and version info.

        Returns:
            Dict with version, features, supported strategies, etc.
        """
        return self._send_message(self.MSG_QUERY_CAPABILITIES)

    def submit_workload(self, workload: Dict[str, Any]) -> Dict[str, Any]:
        """Submit a workload for execution.

        Args:
            workload: Dict with at minimum:
                - type: str (e.g. "compute", "graphics")
                - duration_ns: int (estimated duration in nanoseconds)
                - memory_bytes: int (memory required)

        Returns:
            Response dict with workload_id and status.
        """
        return self._send_message(self.MSG_SUBMIT_WORKLOAD, workload)

    def set_strategy(self, strategy: str) -> Dict[str, Any]:
        """Set the workload distribution strategy.

        Args:
            strategy: One of "round_robin", "afr", "sfr", "auto",
                     "compute_offload", "hybrid", "single_gpu", "task",
                     "ai_driven".

        Returns:
            Response dict confirming the change.
        """
        return self._send_message(
            self.MSG_SET_STRATEGY, {"strategy": strategy}
        )

    def allocate_memory(self, size_bytes: int) -> Dict[str, Any]:
        """Allocate memory on the daemon.

        Args:
            size_bytes: Number of bytes to allocate.

        Returns:
            Response dict with allocation details.
        """
        return self._send_message(
            self.MSG_ALLOC_MEMORY, {"size": size_bytes}
        )

    def free_memory(self, allocation_id: str) -> Dict[str, Any]:
        """Free previously allocated memory.

        Args:
            allocation_id: ID from a previous allocate_memory call.

        Returns:
            Response dict confirming the deallocation.
        """
        return self._send_message(
            self.MSG_FREE_MEMORY, {"allocation_id": allocation_id}
        )

    def subscribe_telemetry(self, interval_ms: int = 1000) -> Dict[str, Any]:
        """Subscribe to telemetry updates.

        Args:
            interval_ms: Update interval in milliseconds.

        Returns:
            Response dict confirming subscription.
        """
        return self._send_message(
            self.MSG_SUBSCRIBE_TELEMETRY, {"interval_ms": interval_ms}
        )

    def health_check(self) -> bool:
        """Quick connectivity check.

        Returns:
            True if daemon is reachable and responsive.
        """
        try:
            self._send_message(self.MSG_QUERY_CAPABILITIES)
            return True
        except MvgalError:
            return False
