# SPDX-License-Identifier: AGPL-3.0-only
# Copyright (c) 2026 sol pbc

"""Paired-link client primitives.

Raises ``TlsError`` for TLS handshake failures, ``ConnectionError`` or
``OSError`` for transport failures, and ``OSError`` for pre-dial credential file
errors in callers that load identities from disk.
"""

from __future__ import annotations

import asyncio
import contextlib
import dataclasses
import logging
import secrets
import urllib.parse
from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Iterator
from typing import Protocol

import requests
import websockets
from cryptography import x509
from OpenSSL import SSL, crypto
from websockets.asyncio.client import ClientConnection
from websockets.exceptions import ConnectionClosed

from solstone.convey.secure_listener.framing import (
    CONTROL_NONCE_LEN,
    FLAG_CLOSE,
    FLAG_DATA,
    FLAG_OPEN,
    FLAG_PING,
    FLAG_PONG,
    FLAG_RESET,
    FLAG_WINDOW,
    INITIAL_WINDOW,
    MAX_CONCURRENT_STREAMS,
    MAX_PAYLOAD,
    MAX_SEND_CREDIT,
    RECOMMENDED_CHUNK,
    RESET_CANCEL,
    RESET_FLOW_CONTROL_ERROR,
    RESET_INTERNAL_ERROR,
    RESET_PROTOCOL_ERROR,
    Frame,
    FrameDecoder,
    ProtocolError,
    build_close,
    build_data,
    build_open,
    build_ping,
    build_pong,
    build_reset,
    build_window,
    parse_control_nonce,
    parse_reset_reason,
    parse_window_credit,
    validate_flags,
)
from solstone.think.link.tls import TlsError as _TlsError

LOG = logging.getLogger(__name__)
_CONNECT_TIMEOUT_SECONDS = 15
_HTTP_TIMEOUT_SECONDS = 30
_KEEPALIVE_INTERVAL_SECONDS = 20
_KEEPALIVE_TIMEOUT_SECONDS = 45
_BODY_DONE = object()


class StreamResetError(ConnectionError):
    """Raised when the peer sends a RESET frame for an active stream."""


class EncryptedTransport(Protocol):
    async def send(self, data: bytes) -> None: ...

    async def recv(self) -> bytes | None: ...

    async def close(self) -> None: ...


@dataclasses.dataclass(frozen=True)
class BodySource:
    length: int
    chunks: Iterable[bytes]


def _coerce_body_source(body: bytes | BodySource) -> BodySource:
    if isinstance(body, BodySource):
        return body
    return BodySource(len(body), (body,))


def _next_body_chunk(chunks: Iterator[bytes]) -> bytes | object:
    return next(chunks, _BODY_DONE)


@dataclasses.dataclass(frozen=True)
class ClientIdentity:
    private_key_pem: str
    client_cert_pem: str
    ca_chain_pem: str
    fingerprint: str
    home_instance_id: str
    home_label: str
    home_attestation: str
    local_endpoints: tuple[dict[str, object], ...] = ()


@dataclasses.dataclass(frozen=True)
class EnrolledDevice:
    device_token: str
    identity: ClientIdentity


@dataclasses.dataclass
class _TlsClientState:
    conn: SSL.Connection
    handshake_done: bool = False


@dataclasses.dataclass
class _StreamState:
    stream_id: int
    buffered: list[bytes] = dataclasses.field(default_factory=list)
    waiters: list[asyncio.Future[bytes | None]] = dataclasses.field(
        default_factory=list
    )
    closed_waiters: list[asyncio.Future[None]] = dataclasses.field(default_factory=list)
    send_credit: int = INITIAL_WINDOW
    recv_credit: int = INITIAL_WINDOW
    unacked_recv: int = 0
    writer_closed: bool = False
    reader_closed: bool = False
    reset_reason: int | None = None
    credit_event: asyncio.Event = dataclasses.field(default_factory=asyncio.Event)

    def __post_init__(self) -> None:
        self.credit_event.set()


class _DialerStream:
    def __init__(self, mux: _DialerMultiplexer, state: _StreamState) -> None:
        self._mux = mux
        self._state = state

    @property
    def id(self) -> int:
        return self._state.stream_id

    async def write(self, data: bytes) -> None:
        if self._state.writer_closed:
            raise ConnectionError(f"stream {self._state.stream_id} writer is closed")
        view = memoryview(data)
        while view:
            if self._state.reset_reason is not None:
                raise StreamResetError(
                    f"stream {self._state.stream_id} reset: {self._state.reset_reason}"
                )
            if self._state.writer_closed:
                raise ConnectionError(
                    f"stream {self._state.stream_id} writer is closed"
                )
            chunk_len = min(
                len(view),
                RECOMMENDED_CHUNK,
                MAX_PAYLOAD,
                self._state.send_credit,
            )
            if chunk_len <= 0:
                self._state.credit_event.clear()
                await self._state.credit_event.wait()
                continue
            chunk = bytes(view[:chunk_len])
            view = view[chunk_len:]
            self._state.send_credit -= chunk_len
            await self._mux._emit(build_data(self._state.stream_id, chunk))

    async def wait_for_send_credit(self) -> None:
        while self._state.send_credit <= 0:
            if self._state.reset_reason is not None:
                raise StreamResetError(
                    f"stream {self._state.stream_id} reset: {self._state.reset_reason}"
                )
            if self._state.writer_closed:
                raise ConnectionError(
                    f"stream {self._state.stream_id} writer is closed"
                )
            self._state.credit_event.clear()
            await self._state.credit_event.wait()

    async def close(self) -> None:
        if self._state.writer_closed:
            return
        self._state.writer_closed = True
        await self._mux._emit(build_close(self._state.stream_id))
        if self._state.reader_closed:
            self._mux._forget(self._state.stream_id)

    async def reset(self, reason: int = RESET_INTERNAL_ERROR) -> None:
        if self._state.writer_closed and self._state.reader_closed:
            return
        self._state.writer_closed = True
        self._state.reader_closed = True
        self._state.reset_reason = reason
        await self._mux._emit(build_reset(self._state.stream_id, reason))
        self._mux._close_stream(self._state, forget=True)

    async def cancel(self) -> None:
        await self.reset(RESET_CANCEL)

    async def read(self) -> AsyncIterator[bytes]:
        while True:
            if self._state.buffered:
                chunk = self._state.buffered.pop(0)
                await self._mux._note_drained(self._state, len(chunk))
                yield chunk
                continue
            if self._state.reader_closed:
                if self._state.reset_reason is not None:
                    raise StreamResetError(
                        f"stream {self._state.stream_id} reset: {self._state.reset_reason}"
                    )
                return
            fut: asyncio.Future[bytes | None] = (
                asyncio.get_running_loop().create_future()
            )
            self._state.waiters.append(fut)
            chunk = await fut
            if chunk is None:
                if self._state.reset_reason is not None:
                    raise StreamResetError(
                        f"stream {self._state.stream_id} reset: {self._state.reset_reason}"
                    )
                return
            await self._mux._note_drained(self._state, len(chunk))
            yield chunk

    async def read_all(self) -> bytes:
        parts = bytearray()
        async for chunk in self.read():
            parts.extend(chunk)
        return bytes(parts)

    @property
    def closed(self) -> asyncio.Future[None]:
        fut: asyncio.Future[None] = asyncio.get_running_loop().create_future()
        if self._state.reader_closed and self._state.writer_closed:
            fut.set_result(None)
            return fut
        self._state.closed_waiters.append(fut)
        return fut


class _DialerMultiplexer:
    def __init__(
        self,
        send_frame: Callable[[bytes], Awaitable[None]],
        on_inbound: Callable[[], None] | None = None,
    ) -> None:
        self._decoder = FrameDecoder()
        self._send_frame = send_frame
        self._on_inbound = on_inbound
        self._streams: dict[int, _StreamState] = {}
        self._next_local_id = 1
        self._closed = False

    async def open_stream(self, initial: bytes = b"") -> _DialerStream:
        if self._closed:
            raise ConnectionError("mux is closed")
        if len(self._streams) >= MAX_CONCURRENT_STREAMS:
            raise ConnectionError("concurrent stream cap reached")
        if len(initial) > MAX_PAYLOAD:
            raise ValueError(f"initial payload exceeds framing max {MAX_PAYLOAD}")
        stream_id = self._next_local_id
        self._next_local_id += 2
        state = _StreamState(stream_id=stream_id)
        if initial:
            state.send_credit -= len(initial)
        self._streams[stream_id] = state
        await self._emit(build_open(stream_id, initial))
        return _DialerStream(self, state)

    async def feed(self, plaintext: bytes) -> None:
        if self._closed or not plaintext:
            return
        if self._on_inbound is not None:
            self._on_inbound()
        self._decoder.feed(plaintext)
        while True:
            try:
                frame = self._decoder.next()
            except ProtocolError:
                self.close()
                return
            if frame is None:
                return
            await self._dispatch(frame)
            if self._closed:
                return

    def close(self) -> None:
        if self._closed:
            return
        self._closed = True
        for state in list(self._streams.values()):
            if state.reset_reason is None:
                state.reset_reason = RESET_INTERNAL_ERROR
            self._close_stream(state, forget=True)

    async def _dispatch(self, frame: Frame) -> None:
        if frame.stream_id == 0:
            await self._dispatch_control(frame)
            return
        if frame.flags & (FLAG_PING | FLAG_PONG):
            await self._reject_stream(frame)
            return
        try:
            validate_flags(frame.flags)
        except ProtocolError:
            await self._reject_stream(frame)
            return

        if frame.flags & FLAG_OPEN:
            await self._emit(build_reset(frame.stream_id, RESET_PROTOCOL_ERROR))
            return

        state = self._streams.get(frame.stream_id)
        if state is None:
            if frame.flags & (FLAG_DATA | FLAG_WINDOW):
                await self._emit(build_reset(frame.stream_id, RESET_PROTOCOL_ERROR))
            return

        if frame.flags & FLAG_DATA:
            if len(frame.payload) > state.recv_credit:
                await self._emit(build_reset(frame.stream_id, RESET_FLOW_CONTROL_ERROR))
                state.reset_reason = RESET_FLOW_CONTROL_ERROR
                self._close_stream(state, forget=True)
                return
            state.recv_credit -= len(frame.payload)
            if state.waiters:
                waiter = state.waiters.pop(0)
                if not waiter.done():
                    waiter.set_result(frame.payload)
            else:
                state.buffered.append(frame.payload)

        if frame.flags & FLAG_CLOSE:
            state.reader_closed = True
            while state.waiters:
                waiter = state.waiters.pop(0)
                if not waiter.done():
                    waiter.set_result(None)
            if state.writer_closed:
                self._forget(frame.stream_id)
            self._resolve_closed(state)

        if frame.flags & FLAG_WINDOW:
            try:
                credit = parse_window_credit(frame)
            except ProtocolError:
                await self._emit(build_reset(frame.stream_id, RESET_PROTOCOL_ERROR))
                state.reset_reason = RESET_PROTOCOL_ERROR
                self._close_stream(state, forget=True)
                return
            if state.send_credit + credit > MAX_SEND_CREDIT:
                await self._emit(build_reset(frame.stream_id, RESET_FLOW_CONTROL_ERROR))
                state.reset_reason = RESET_FLOW_CONTROL_ERROR
                self._close_stream(state, forget=True)
                return
            state.send_credit += credit
            state.credit_event.set()

        if frame.flags & FLAG_RESET:
            try:
                state.reset_reason = parse_reset_reason(frame)
            except ProtocolError:
                state.reset_reason = RESET_PROTOCOL_ERROR
            self._close_stream(state, forget=True)

    async def _dispatch_control(self, frame: Frame) -> None:
        is_ping = bool(frame.flags & FLAG_PING)
        is_pong = bool(frame.flags & FLAG_PONG)
        if is_ping == is_pong:
            self.close()
            return
        if frame.flags & ~(FLAG_PING | FLAG_PONG):
            self.close()
            return
        try:
            nonce = parse_control_nonce(frame)
        except ProtocolError:
            self.close()
            return
        if is_ping:
            await self._emit(build_pong(nonce))

    async def _reject_stream(self, frame: Frame) -> None:
        state = self._streams.get(frame.stream_id)
        await self._emit(build_reset(frame.stream_id, RESET_PROTOCOL_ERROR))
        if state is not None:
            state.reset_reason = RESET_PROTOCOL_ERROR
            self._close_stream(state, forget=True)

    async def _emit(self, frame: Frame) -> None:
        if self._closed:
            return
        await self._send_frame(frame.encode())

    async def _note_drained(self, state: _StreamState, n: int) -> None:
        if n <= 0 or state.reader_closed:
            return
        state.unacked_recv += n
        if state.unacked_recv < INITIAL_WINDOW // 2:
            return
        grant = state.unacked_recv
        state.recv_credit += grant
        state.unacked_recv = 0
        await self._emit(build_window(state.stream_id, grant))

    def _close_stream(self, state: _StreamState, *, forget: bool) -> None:
        state.writer_closed = True
        state.reader_closed = True
        while state.waiters:
            waiter = state.waiters.pop(0)
            if not waiter.done():
                waiter.set_result(None)
        state.credit_event.set()
        self._resolve_closed(state)
        if forget:
            self._forget(state.stream_id)

    def _resolve_closed(self, state: _StreamState) -> None:
        while state.closed_waiters:
            waiter = state.closed_waiters.pop(0)
            if not waiter.done():
                waiter.set_result(None)

    def _forget(self, stream_id: int) -> None:
        self._streams.pop(stream_id, None)


class _WsEncryptedTransport:
    def __init__(self, ws: ClientConnection) -> None:
        self._ws = ws

    async def send(self, data: bytes) -> None:
        await self._ws.send(data)

    async def recv(self) -> bytes | None:
        try:
            message = await self._ws.recv()
        except ConnectionClosed:
            return None
        return message if isinstance(message, bytes) else message.encode("utf-8")

    async def close(self) -> None:
        await self._ws.close()


class _TcpEncryptedTransport:
    def __init__(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        self._reader = reader
        self._writer = writer

    async def send(self, data: bytes) -> None:
        self._writer.write(data)
        await self._writer.drain()

    async def recv(self) -> bytes | None:
        chunk = await self._reader.read(65536)
        return chunk if chunk else None

    async def close(self) -> None:
        self._writer.close()
        with contextlib.suppress(Exception):
            await self._writer.wait_closed()


class TunnelSession:
    def __init__(
        self,
        *,
        transport: EncryptedTransport,
        tls: _TlsClientState,
        identity: ClientIdentity | None = None,
        keepalive_interval: float = _KEEPALIVE_INTERVAL_SECONDS,
        keepalive_timeout: float = _KEEPALIVE_TIMEOUT_SECONDS,
    ) -> None:
        self._transport = transport
        self._tls = tls
        self._identity = identity
        self._keepalive_interval = keepalive_interval
        self._keepalive_timeout = keepalive_timeout
        self._tls_lock = asyncio.Lock()
        self._mux = _DialerMultiplexer(
            self._send_plaintext,
            on_inbound=self._record_activity,
        )
        self._closed = asyncio.Event()
        self._last_activity = asyncio.get_running_loop().time()
        self._dead = False
        task_name = (
            f"link-client-{identity.home_instance_id}"
            if identity is not None
            else "link-pair"
        )
        self._reader_task = asyncio.create_task(
            self._read_transport(),
            name=task_name,
        )
        self._keepalive_task: asyncio.Task[None] | None = asyncio.create_task(
            self._keepalive_loop(),
            name=f"{task_name}-keepalive",
        )

    async def __aenter__(self) -> TunnelSession:
        return self

    async def __aexit__(self, *_exc: object) -> None:
        await self.close()

    def peer_certificate(self) -> x509.Certificate | None:
        """The leaf certificate the TLS peer presented, post-handshake.

        Used during pairing to bind the live TLS endpoint to the pinned CA so a
        relay that merely echoes the real CA chain in the response body cannot
        masquerade as the home. Returns None if no peer certificate is set.
        """
        cert = self._tls.conn.get_peer_certificate()
        if cert is None:
            return None
        return cert.to_cryptography()

    def peer_certificate_chain(self) -> tuple[x509.Certificate, ...]:
        """The presented TLS certificate chain, post-handshake."""

        chain = self._tls.conn.get_peer_cert_chain()
        if chain is None:
            return ()
        return tuple(cert.to_cryptography() for cert in chain)

    async def request(
        self,
        method: str,
        path: str,
        *,
        headers: dict[str, str] | None = None,
        body: bytes | BodySource = b"",
    ) -> tuple[int, dict[str, str], bytes]:
        stream = await self._open_request_stream(
            method,
            path,
            headers=headers,
            body=body,
        )
        try:
            response = await stream.read_all()
        except asyncio.CancelledError:
            with contextlib.suppress(Exception):
                await stream.cancel()
            raise
        return _parse_http_response(response)

    async def stream_request(
        self,
        method: str,
        path: str,
        *,
        headers: dict[str, str] | None = None,
        body: bytes | BodySource = b"",
    ) -> tuple[int, dict[str, str], bytes, _DialerStream]:
        stream = await self._open_request_stream(
            method,
            path,
            headers=headers,
            body=body,
        )
        buffered = bytearray()
        try:
            async for chunk in stream.read():
                buffered.extend(chunk)
                split = buffered.find(b"\r\n\r\n")
                if split < 0:
                    continue
                status, response_headers = _parse_http_head(bytes(buffered[:split]))
                return status, response_headers, bytes(buffered[split + 4 :]), stream
        except asyncio.CancelledError:
            with contextlib.suppress(Exception):
                await stream.cancel()
            raise
        raise ValueError("response missing header terminator")

    async def _open_request_stream(
        self,
        method: str,
        path: str,
        *,
        headers: dict[str, str] | None,
        body: bytes | BodySource,
    ) -> _DialerStream:
        source = _coerce_body_source(body)
        head = _http_head_bytes(
            method,
            path,
            headers=headers,
            content_length=source.length,
        )
        stream = await self._mux.open_stream(head)
        try:
            chunks = iter(source.chunks)
            while True:
                await stream.wait_for_send_credit()
                chunk = await asyncio.to_thread(_next_body_chunk, chunks)
                if chunk is _BODY_DONE:
                    break
                if not isinstance(chunk, bytes):
                    raise TypeError("body source yielded a non-bytes chunk")
                await stream.write(chunk)
            await stream.close()
        except asyncio.CancelledError:
            with contextlib.suppress(Exception):
                await stream.cancel()
            raise
        except Exception:
            with contextlib.suppress(Exception):
                await stream.reset()
            raise
        return stream

    async def close(self) -> None:
        if self._closed.is_set():
            return
        await self._cancel_keepalive()
        self._mux.close()
        await self._transport.close()
        await self._reader_task
        self._closed.set()

    @property
    def is_alive(self) -> bool:
        return not (self._closed.is_set() or self._mux._closed or self._dead)

    def failure_reason(self) -> str | None:
        if self._dead:
            return "liveness_failed"
        if self._closed.is_set():
            return "session_closed"
        return None

    async def _read_transport(self) -> None:
        try:
            while True:
                inbound = await self._transport.recv()
                if inbound is None:
                    return
                async with self._tls_lock:
                    outbound, plaintext = _drive_tls_client(self._tls, inbound=inbound)
                if outbound:
                    await self._transport.send(outbound)
                if plaintext:
                    await self._mux.feed(plaintext)
                    if self._mux._closed:
                        return
        finally:
            await self._cancel_keepalive()
            self._mux.close()
            # Close the transport here too: on a peer-initiated close (e.g. the
            # home journal restarting), recv() returns None and the reader exits
            # with `_closed` unset. A later close() then early-returns on the set
            # flag and never reaches `_transport.close()`, leaving the socket in
            # CLOSE-WAIT and leaking a fd per reconnect. Closing here (idempotent
            # with close()) retires the dead socket on every reader exit.
            with contextlib.suppress(Exception):
                await self._transport.close()
            self._closed.set()

    async def _send_plaintext(self, plaintext: bytes) -> None:
        async with self._tls_lock:
            outbound, _ = _drive_tls_client(self._tls, plaintext_out=plaintext)
        if outbound:
            await self._transport.send(outbound)

    def _record_activity(self) -> None:
        self._last_activity = asyncio.get_running_loop().time()

    async def _keepalive_loop(self) -> None:
        try:
            while True:
                await asyncio.sleep(self._keepalive_interval)
                if self._closed.is_set():
                    return
                idle = asyncio.get_running_loop().time() - self._last_activity
                if idle > self._keepalive_timeout:
                    self._dead = True
                    await self._transport.close()
                    return
                nonce = secrets.token_bytes(CONTROL_NONCE_LEN)
                await self._send_plaintext(build_ping(nonce).encode())
        except asyncio.CancelledError:
            raise
        except (OSError, _TlsError):
            self._dead = True

    async def _cancel_keepalive(self) -> None:
        task = self._keepalive_task
        if task is None or task.done():
            return
        task.cancel()
        with contextlib.suppress(asyncio.CancelledError):
            await task


class Client:
    @staticmethod
    def enroll_device(relay_url: str, identity: ClientIdentity) -> EnrolledDevice:
        endpoint = f"{relay_url.rstrip('/')}/enroll/device"
        LOG.info("client %s: enrolling device token", identity.fingerprint)
        payload = _post_json(
            endpoint,
            {
                "instance_id": identity.home_instance_id,
                "client_cert": identity.client_cert_pem,
                "home_attestation": identity.home_attestation,
            },
        )
        device_token = _required_str(payload, "device_token")
        LOG.info("client %s: enroll complete", identity.fingerprint)
        return EnrolledDevice(device_token=device_token, identity=identity)

    @staticmethod
    async def dial(relay_url: str, enrolled: EnrolledDevice) -> TunnelSession:
        identity = enrolled.identity
        url = (
            _to_ws(relay_url.rstrip("/"))
            + "/session/dial?"
            + urllib.parse.urlencode(
                {
                    "instance": identity.home_instance_id,
                    "token": enrolled.device_token,
                }
            )
        )
        LOG.info("client %s: dialing %s", identity.fingerprint, _redact_url(url))
        ws = await websockets.connect(url, max_size=None)
        return await _open_tunnel_session(_WsEncryptedTransport(ws), identity)

    @staticmethod
    async def dial_direct(
        host: str,
        enrolled: EnrolledDevice,
        *,
        port: int = 7657,
    ) -> TunnelSession:
        identity = enrolled.identity
        LOG.info("client %s: dialing direct %s:%d", identity.fingerprint, host, port)
        reader, writer = await asyncio.open_connection(host, port)
        return await _open_tunnel_session(
            _TcpEncryptedTransport(reader, writer),
            identity,
        )


async def _open_tunnel_session(
    transport: EncryptedTransport,
    identity: ClientIdentity,
) -> TunnelSession:
    tls = _new_tls_client(_build_tls_client_ctx(identity))
    pending_plaintext = await _drive_client_handshake(transport, tls)
    session = TunnelSession(
        transport=transport,
        tls=tls,
        identity=identity,
    )
    if pending_plaintext:
        await session._mux.feed(bytes(pending_plaintext))
    return session


async def _open_pairing_session(transport: EncryptedTransport) -> TunnelSession:
    tls = _new_tls_client(_build_no_cert_client_ctx())
    pending_plaintext = await _drive_client_handshake(transport, tls)
    session = TunnelSession(transport=transport, tls=tls)
    if pending_plaintext:
        await session._mux.feed(bytes(pending_plaintext))
    return session


async def _drive_client_handshake(
    transport: EncryptedTransport,
    tls: _TlsClientState,
) -> bytearray:
    try:
        pending_plaintext = bytearray()
        outbound, plaintext = _drive_tls_client(tls)
        if outbound:
            await transport.send(outbound)
        pending_plaintext.extend(plaintext)
        while not tls.handshake_done:
            inbound = await asyncio.wait_for(
                transport.recv(),
                timeout=_CONNECT_TIMEOUT_SECONDS,
            )
            if inbound is None:
                raise _TlsError("transport closed during TLS handshake")
            outbound, plaintext = _drive_tls_client(tls, inbound=inbound)
            if outbound:
                await transport.send(outbound)
            pending_plaintext.extend(plaintext)
    except Exception:
        await transport.close()
        raise
    return pending_plaintext


def _build_tls_client_ctx(identity: ClientIdentity) -> SSL.Context:
    ctx = SSL.Context(SSL.TLS_METHOD)
    ctx.set_min_proto_version(SSL.TLS1_3_VERSION)
    ctx.set_max_proto_version(SSL.TLS1_3_VERSION)
    ctx.use_certificate(
        crypto.load_certificate(
            crypto.FILETYPE_PEM,
            identity.client_cert_pem.encode("ascii"),
        )
    )
    ctx.use_privatekey(
        crypto.load_privatekey(
            crypto.FILETYPE_PEM,
            identity.private_key_pem.encode("ascii"),
        )
    )
    store = ctx.get_cert_store()
    assert store is not None, "client TLS context must expose a cert store"
    for cert_pem in _split_pem_chain(identity.ca_chain_pem):
        store.add_cert(
            crypto.X509.from_cryptography(x509.load_pem_x509_certificate(cert_pem)),
        )
    ctx.set_verify(SSL.VERIFY_PEER, _verify_server_cert)
    ctx.check_privatekey()
    return ctx


def _build_no_cert_client_ctx() -> SSL.Context:
    ctx = SSL.Context(SSL.TLS_METHOD)
    ctx.set_min_proto_version(SSL.TLS1_3_VERSION)
    ctx.set_max_proto_version(SSL.TLS1_3_VERSION)
    ctx.set_verify(SSL.VERIFY_PEER, _accept_any_server_cert)
    return ctx


def _accept_any_server_cert(*_args: object) -> bool:
    return True


def _verify_server_cert(
    _conn: SSL.Connection,
    _cert: crypto.X509,
    _errno: int,
    _depth: int,
    preverify_ok: int,
) -> bool:
    return bool(preverify_ok)


def _new_tls_client(ctx: SSL.Context) -> _TlsClientState:
    conn = SSL.Connection(ctx, None)
    conn.set_connect_state()
    return _TlsClientState(conn=conn)


def _drive_tls_client(
    state: _TlsClientState,
    *,
    inbound: bytes = b"",
    plaintext_out: bytes = b"",
) -> tuple[bytes, bytes]:
    if inbound:
        state.conn.bio_write(inbound)
    if plaintext_out:
        try:
            # sendall, not send: pyOpenSSL's Connection.send() writes at most one
            # TLS record (16 KiB) per call and returns the partial count. Calling
            # send() once and ignoring the return value silently truncated any
            # frame larger than one record on the wire, so the peer's frame
            # decoder blocked forever waiting for the dropped tail (manifested as
            # ~30s link-tunnel timeouts on observer-segment uploads >16 KiB).
            # sendall() loops SSL_write internally until the whole buffer is
            # encrypted into the write BIO.
            state.conn.sendall(plaintext_out)
        except SSL.WantReadError:
            pass
        except SSL.Error as exc:
            raise _TlsError(f"send failed: {exc}") from exc

    if not state.handshake_done:
        try:
            state.conn.do_handshake()
            state.handshake_done = True
        except SSL.WantReadError:
            pass
        except SSL.Error as exc:
            raise _TlsError(f"handshake failed: {exc}") from exc

    plaintext_in = bytearray()
    if state.handshake_done:
        while True:
            try:
                chunk = state.conn.recv(16 * 1024)
            except SSL.WantReadError:
                break
            except SSL.ZeroReturnError:
                break
            except SSL.Error as exc:
                raise _TlsError(f"recv failed: {exc}") from exc
            if not chunk:
                break
            plaintext_in.extend(chunk)

    outbound = bytearray()
    while True:
        try:
            chunk = state.conn.bio_read(16 * 1024)
        except SSL.WantReadError:
            break
        if not chunk:
            break
        outbound.extend(chunk)
    return bytes(outbound), bytes(plaintext_in)


def _http_head_bytes(
    method: str,
    path: str,
    *,
    headers: dict[str, str] | None,
    content_length: int,
) -> bytes:
    normalized_headers = {"host": "spl.local"}
    if headers:
        normalized_headers.update({k.lower(): v for k, v in headers.items()})
    normalized_headers["content-length"] = str(content_length)
    head = (
        f"{method} {path} HTTP/1.1\r\n"
        + "".join(f"{name}: {value}\r\n" for name, value in normalized_headers.items())
        + "\r\n"
    )
    return head.encode("ascii")


def _parse_http_response(raw: bytes) -> tuple[int, dict[str, str], bytes]:
    split = raw.find(b"\r\n\r\n")
    if split < 0:
        raise ValueError("response missing header terminator")
    body = raw[split + 4 :]
    status, headers = _parse_http_head(raw[:split])
    if headers.get("transfer-encoding", "").lower() == "chunked":
        return status, headers, _dechunk(body)
    content_length = headers.get("content-length")
    if content_length is None:
        return status, headers, body
    return status, headers, body[: int(content_length)]


def _parse_http_head(raw: bytes) -> tuple[int, dict[str, str]]:
    head = raw.decode("latin-1")
    lines = head.split("\r\n")
    if not lines:
        raise ValueError("response missing status line")
    parts = lines[0].split(" ", 2)
    if len(parts) < 2 or not parts[1].isdigit():
        raise ValueError(f"bad status line: {lines[0]!r}")
    status = int(parts[1])
    headers: dict[str, str] = {}
    for line in lines[1:]:
        if not line or ":" not in line:
            continue
        name, value = line.split(":", 1)
        headers[name.strip().lower()] = value.strip()
    return status, headers


def _dechunk(raw: bytes) -> bytes:
    out = bytearray()
    index = 0
    while index < len(raw):
        line_end = raw.find(b"\r\n", index)
        if line_end < 0:
            raise ValueError("chunked response missing size terminator")
        size_text = raw[index:line_end].decode("ascii").split(";", 1)[0].strip()
        size = int(size_text, 16)
        index = line_end + 2
        if size == 0:
            return bytes(out)
        out.extend(raw[index : index + size])
        index += size + 2
    return bytes(out)


def _post_json(url: str, payload: dict[str, object]) -> dict[str, object]:
    response = requests.post(url, json=payload, timeout=_HTTP_TIMEOUT_SECONDS)
    if not response.ok:
        raise RuntimeError(
            f"POST {url} failed: HTTP {response.status_code}: {response.text}"
        )
    parsed = response.json()
    if not isinstance(parsed, dict):
        raise RuntimeError(f"unexpected JSON response from {url}")
    return parsed


def _required_str(payload: dict[str, object], key: str) -> str:
    value = payload.get(key)
    if not isinstance(value, str) or not value:
        raise RuntimeError(f"missing string field: {key}")
    return value


def _split_pem_chain(pem_bundle: str) -> list[bytes]:
    marker = "-----END CERTIFICATE-----"
    certs: list[bytes] = []
    for chunk in pem_bundle.split(marker):
        chunk = chunk.strip()
        if not chunk:
            continue
        certs.append(f"{chunk}\n{marker}\n".encode("ascii"))
    return certs


def _to_ws(url: str) -> str:
    if url.startswith("https://"):
        return "wss://" + url[len("https://") :]
    if url.startswith("http://"):
        return "ws://" + url[len("http://") :]
    return url


def _redact_url(url: str) -> str:
    parsed = urllib.parse.urlparse(url)
    query = urllib.parse.parse_qs(parsed.query)
    if "token" in query:
        query["token"] = ["<redacted>"]
    return urllib.parse.urlunparse(
        parsed._replace(query=urllib.parse.urlencode(query, doseq=True))
    )


__all__ = [
    "BodySource",
    "Client",
    "ClientIdentity",
    "EncryptedTransport",
    "EnrolledDevice",
    "StreamResetError",
    "TunnelSession",
    "_http_head_bytes",
]
