"""TCP/TLS transport for talking to a remote ``browser-cli serve``.

Owns the wire mechanics of the remote leg: open a socket (TLS on :443),
complete the signed challenge/response handshake with an optional post-quantum
key exchange, frame the request, and read the framed (possibly encrypted)
response. The higher-level "which endpoint / which profile / which key"
decisions stay in :mod:`browser_cli.client`, which re-exports these for
backward compatibility.
"""
from __future__ import annotations

import json
import socket
import struct
import sys

from browser_cli.errors import BrowserNotConnected
from browser_cli.endpoints import _resolve_connect_endpoint
from browser_cli.version_manager import MAX_MSG_BYTES as _MAX_MSG_BYTES
from browser_cli.version_manager import USER_AGENT as _USER_AGENT

_PQ_WARNING = (
    "** WARNING: connection is not using a post-quantum key exchange algorithm.\n"
    "** This session may be vulnerable to store now, decrypt later attacks.\n"
)

def _recv_exact(sock: socket.socket, n: int) -> bytes:
    buf = b""
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            raise ConnectionError("Socket closed before full message received")
        buf += chunk
    return buf

def _recv_all(sock: socket.socket) -> bytes:
    raw_len = _recv_exact(sock, 4)
    msg_len = struct.unpack("<I", raw_len)[0]
    if msg_len > _MAX_MSG_BYTES:
        raise ConnectionError(f"Response too large ({msg_len} bytes)")
    return _recv_exact(sock, msg_len)

def _send_remote(endpoint: str, msg: dict, private_key=None) -> bytes | None:
    connect_ep = _resolve_connect_endpoint(endpoint)
    host, _, port_str = connect_ep.rpartition(":")
    port = int(port_str)
    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
            ctx = ssl.create_default_context()
            sock = ctx.wrap_socket(raw_sock, server_hostname=host)
        else:
            sock = raw_sock
    except Exception:
        raw_sock.close()
        raise
    with sock:

        # receive challenge
        challenge_raw = _recv_all(sock)
        if challenge_raw is None:
            raise BrowserNotConnected(f"No challenge received from {endpoint}")
        try:
            challenge = json.loads(challenge_raw)
            nonce_hex = challenge.get("nonce") if challenge.get("type") == "challenge" else None
        except (json.JSONDecodeError, AttributeError):
            nonce_hex = None

        min_ver = challenge.get("min_client_version") if isinstance(challenge, dict) else None
        if min_ver:
            from browser_cli.version_manager import parse_version
            try:
                client_ver = _USER_AGENT.split("/", 1)[1]
                if parse_version(client_ver) < parse_version(min_ver):
                    raise BrowserNotConnected(
                        f"Client version {client_ver} is too old for this server "
                        f"(requires >= {min_ver}). Run: pip install --upgrade browser-cli"
                    )
            except (IndexError, ValueError):
                pass

        pq_shared_secret = None
        if nonce_hex and private_key is not None:
            from browser_cli.auth import PQ_KEX_ALG, pq_encrypt, pq_kex_client_encapsulate, sign, public_key_hex
            nonce = bytes.fromhex(nonce_hex)
            clean_msg = {k: v for k, v in msg.items() if k not in {"token", "pubkey", "sig", "pq_kex", "encrypted"}}
            kex = challenge.get("pq_kex") if isinstance(challenge, dict) else None
            if isinstance(kex, dict) and kex.get("alg") == PQ_KEX_ALG and kex.get("public_key"):
                ciphertext_hex, pq_shared_secret = pq_kex_client_encapsulate(str(kex["public_key"]))
                clean_msg["pq_kex"] = {"alg": PQ_KEX_ALG, "ciphertext": ciphertext_hex}
            else:
                sys.stderr.write(_PQ_WARNING)
            sig = sign(private_key, nonce, clean_msg, pq_shared_secret)
            msg = {**clean_msg, "pubkey": public_key_hex(private_key), "sig": sig.hex()}
            if pq_shared_secret is not None:
                encrypted = pq_encrypt(pq_shared_secret, "request", json.dumps(clean_msg).encode("utf-8"))
                msg = {
                    "id": clean_msg.get("id"),
                    "user_agent": clean_msg.get("user_agent"),
                    "pubkey": public_key_hex(private_key),
                    "sig": sig.hex(),
                    "pq_kex": clean_msg["pq_kex"],
                    "encrypted": encrypted,
                }
        else:
            sys.stderr.write(_PQ_WARNING)

        payload = json.dumps(msg).encode("utf-8")
        framed = struct.pack("<I", len(payload)) + payload
        sock.sendall(framed)
        response = _recv_all(sock)
        if response is not None and pq_shared_secret is not None:
            try:
                from browser_cli.auth import pq_decrypt
                envelope = json.loads(response)
                if isinstance(envelope, dict) and "encrypted" in envelope:
                    return pq_decrypt(pq_shared_secret, "response", envelope["encrypted"])
            except Exception as e:
                raise BrowserNotConnected(f"Cannot decrypt post-quantum remote response: {e}") from e
        return response