"""Socket helpers for remote TCP/TLS transport."""
from __future__ import annotations

import asyncio
import socket
from contextlib import contextmanager

from browser_cli.endpoints import _resolve_connect_endpoint
from browser_cli.framing import async_recv_exact, async_recv_frame, recv_exact, recv_frame

def recv_exact_bytes(sock: socket.socket, n: int) -> bytes:
  return recv_exact(sock, n) or b""

def recv_all(sock: socket.socket) -> bytes:
  return recv_frame(sock, label="Response") or b""

async def async_recv_exact_bytes(reader: asyncio.StreamReader, n: int) -> bytes:
  return await async_recv_exact(reader, n) or b""

async def async_recv_all(reader: asyncio.StreamReader) -> bytes:
  return await async_recv_frame(reader, label="Response") or b""

def split_endpoint(endpoint: str) -> tuple[str, int]:
  connect_ep = _resolve_connect_endpoint(endpoint)
  host, _, port_str = connect_ep.rpartition(":")
  return host, int(port_str)

def connect_socket(endpoint: str) -> socket.socket:
  """Open and (on :443) TLS-wrap a socket. Caller owns closing it."""
  host, port = split_endpoint(endpoint)
  raw_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  raw_sock.settimeout(30)
  try:
    raw_sock.connect((host, port))
    if port == 443:
      import ssl
      sock = ssl.create_default_context().wrap_socket(raw_sock, server_hostname=host)
    else:
      sock = raw_sock
  except Exception:
    raw_sock.close()
    raise
  return sock

@contextmanager
def open_socket(endpoint: str):
  sock = connect_socket(endpoint)
  with sock:
    yield sock

async def open_async_connection(endpoint: str) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
  host, port = split_endpoint(endpoint)
  ssl_ctx = None
  if port == 443:
    import ssl
    ssl_ctx = ssl.create_default_context()
  return await asyncio.open_connection(host, port, ssl=ssl_ctx, server_hostname=host if ssl_ctx else None)