"""Per-process pool of authenticated remote connections for reuse.

A ``browser-cli serve`` connection stays open after its first (encrypted)
command, so the client can send further commands over it without re-running the
TCP/TLS/challenge/auth handshake (~hundreds of ms each). Only encrypted (PQ)
sessions are pooled — plaintext/legacy sessions stay one-shot, matching the
server, which only loops for encrypted sessions.

Connections are checked out exclusively (never shared between threads at once),
returned on success, and dropped on any I/O error or once older than an idle
bound (kept below the server's idle timeout so we don't reuse a connection the
server has already closed).
"""
from __future__ import annotations

import atexit
import json
import socket
import threading
import time

from browser_cli.constants import REMOTE_SESSION_IDLE_TIMEOUT
from browser_cli.framing import frame

# Retire a pooled connection a few seconds before the server would, so we never
# hand back one the server has just timed out and closed.
_MAX_IDLE_SECONDS = max(5, REMOTE_SESSION_IDLE_TIMEOUT - 5)
_MAX_PER_ENDPOINT = 8

class PooledConnection:
  __slots__ = ("sock", "secret", "last_used")

  def __init__(self, sock: socket.socket, secret: bytes) -> None:
    self.sock = sock
    self.secret = secret
    self.last_used = time.monotonic()

_POOL: dict[str, list[PooledConnection]] = {}
_LOCK = threading.Lock()

def _close(sock: socket.socket) -> None:
  try:
    sock.close()
  except OSError:
    pass

def checkout(endpoint: str) -> PooledConnection | None:
  """Take an idle authenticated connection for *endpoint*, or None."""
  now = time.monotonic()
  with _LOCK:
    conns = _POOL.get(endpoint)
    while conns:
      conn = conns.pop()
      if now - conn.last_used <= _MAX_IDLE_SECONDS:
        return conn
      _close(conn.sock)  # too old — assume the server has dropped it
  return None

def checkin(endpoint: str, conn: PooledConnection) -> None:
  """Return a still-healthy connection to the pool for reuse."""
  conn.last_used = time.monotonic()
  with _LOCK:
    bucket = _POOL.setdefault(endpoint, [])
    if len(bucket) >= _MAX_PER_ENDPOINT:
      _close(conn.sock)
      return
    bucket.append(conn)

def discard(conn: PooledConnection) -> None:
  """Drop a connection that errored or is no longer usable."""
  _close(conn.sock)

def close_all() -> None:
  """Close every pooled connection (process exit / test isolation)."""
  with _LOCK:
    for bucket in _POOL.values():
      for conn in bucket:
        _close(conn.sock)
    _POOL.clear()

def session_inner_message(msg: dict) -> dict:
  """Strip auth/transport fields, leaving the command for an established session."""
  keep = {"id", "command", "args", "user_agent", "accept_encoding", "_route", "_suppress_pq_warning"}
  return {k: v for k, v in msg.items() if k in keep}

def send_over(conn: PooledConnection, msg: dict) -> bytes | None:
  """Send one command over an existing encrypted session. Raises on I/O error."""
  from browser_cli.auth import pq_encrypt
  from browser_cli.remote.socket import recv_all
  from browser_cli.remote.transport import _decode_pq_response

  inner = json.dumps(session_inner_message(msg)).encode("utf-8")
  envelope = json.dumps({"encrypted": pq_encrypt(conn.secret, "request", inner)}).encode("utf-8")
  conn.sock.sendall(frame(envelope))
  response = recv_all(conn.sock)
  if not response:
    # EOF — an older server (no session loop) closed after one command. Treat as
    # a transport failure so the caller re-handshakes; never as an app error,
    # which could double-execute a non-idempotent command on retry.
    raise EOFError("remote closed the pooled connection")
  return _decode_pq_response(response, conn.secret)

atexit.register(close_all)