"""Length-prefixed byte framing used by browser-cli transports.

Frame format is shared by local IPC and remote TCP: 4-byte little-endian payload
length followed by raw payload bytes. Native Messaging stdio has the same length
prefix but JSON helpers stay in ``browser_cli.native.host`` because they operate
on file streams, not sockets/asyncio streams.
"""
from __future__ import annotations

import asyncio
import struct
from typing import Protocol

from browser_cli.version_manager import MAX_MSG_BYTES

class RecvSocket(Protocol):
    def recv(self, n: int) -> bytes: ...

class SendSocket(Protocol):
    def sendall(self, data: bytes) -> None: ...

def frame(data: bytes) -> bytes:
    """Return *data* with the browser-cli 4-byte little-endian length prefix."""
    return struct.pack("<I", len(data)) + data

def _message_length(raw_len: bytes, *, label: str) -> int:
    msg_len = struct.unpack("<I", raw_len)[0]
    if msg_len > MAX_MSG_BYTES:
        raise ConnectionError(f"{label} too large ({msg_len} bytes)")
    return msg_len

def recv_exact(sock: RecvSocket, n: int, *, allow_eof: bool = False) -> bytes | None:
    """Read exactly *n* bytes from a blocking socket-like object.

    Returns ``None`` on EOF when ``allow_eof=True``; otherwise raises
    ``ConnectionError``.
    """
    buf = b""
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            if allow_eof:
                return None
            raise ConnectionError("Socket closed before full message received")
        buf += chunk
    return buf

def recv_frame(sock: RecvSocket, *, allow_eof: bool = False, label: str = "Message") -> bytes | None:
    """Read one framed payload from a blocking socket-like object."""
    raw_len = recv_exact(sock, 4, allow_eof=allow_eof)
    if raw_len is None:
        return None
    return recv_exact(sock, _message_length(raw_len, label=label), allow_eof=allow_eof)

def send_frame(sock: SendSocket, data: bytes) -> None:
    """Send one framed payload through a blocking socket-like object."""
    sock.sendall(frame(data))

async def async_recv_exact(reader: asyncio.StreamReader, n: int, *, allow_eof: bool = False) -> bytes | None:
    """Read exactly *n* bytes from an asyncio StreamReader."""
    try:
        return await reader.readexactly(n)
    except asyncio.IncompleteReadError as exc:
        if allow_eof:
            return None
        raise ConnectionError("Socket closed before full message received") from exc

async def async_recv_frame(
    reader: asyncio.StreamReader,
    *,
    allow_eof: bool = False,
    label: str = "Message",
) -> bytes | None:
    """Read one framed payload from an asyncio StreamReader."""
    raw_len = await async_recv_exact(reader, 4, allow_eof=allow_eof)
    if raw_len is None:
        return None
    return await async_recv_exact(reader, _message_length(raw_len, label=label), allow_eof=allow_eof)

async def async_send_frame(writer: asyncio.StreamWriter, data: bytes) -> None:
    """Send one framed payload through an asyncio StreamWriter."""
    writer.write(frame(data))
    await writer.drain()