# ABOUTME: greetd IPC protocol implementation — communicates via Unix socket.
# ABOUTME: Uses length-prefixed JSON encoding as specified by the greetd IPC protocol.

import json
import struct
from typing import Any

MAX_PAYLOAD_SIZE = 65536


def _recvall(sock: Any, n: int) -> bytes:
    """Receive exactly n bytes from socket, looping on partial reads."""
    buf = bytearray()
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            raise ConnectionError("Connection closed while reading data")
        buf.extend(chunk)
    return bytes(buf)


def send_message(sock: Any, msg: dict) -> None:
    """Send a length-prefixed JSON message to the greetd socket."""
    payload = json.dumps(msg).encode("utf-8")
    header = struct.pack("!I", len(payload))
    sock.sendall(header + payload)


def recv_message(sock: Any) -> dict:
    """Receive a length-prefixed JSON message from the greetd socket."""
    header = _recvall(sock, 4)
    length = struct.unpack("!I", header)[0]

    if length > MAX_PAYLOAD_SIZE:
        raise ConnectionError(f"Payload too large: {length} bytes (max {MAX_PAYLOAD_SIZE})")

    payload = _recvall(sock, length)
    return json.loads(payload.decode("utf-8"))


def create_session(sock: Any, username: str) -> dict:
    """Send a create_session request to greetd and return the response."""
    send_message(sock, {"type": "create_session", "username": username})
    return recv_message(sock)


def post_auth_response(sock: Any, response: str | None) -> dict:
    """Send an authentication response (e.g. password) to greetd."""
    send_message(sock, {"type": "post_auth_message_response", "response": response})
    return recv_message(sock)


def start_session(sock: Any, cmd: list[str]) -> dict:
    """Send a start_session request to launch the user's session."""
    send_message(sock, {"type": "start_session", "cmd": cmd})
    return recv_message(sock)


def cancel_session(sock: Any) -> dict:
    """Cancel the current authentication session."""
    send_message(sock, {"type": "cancel_session"})
    return recv_message(sock)
