"""Outbound relay bridge for NAT-safe node HTTP requests."""

from __future__ import annotations

import base64
import http.client
import json
import logging
import os
import re
import threading
import time
import urllib.parse
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass

log = logging.getLogger(__name__)

DEFAULT_MAX_CONCURRENCY = 8

# Activation tensors ride the relay as one WebSocket frame per hop, so the
# websockets default of 1 MiB rejects any real prefill (close code 1009).
DEFAULT_WS_MAX_BYTES = 256 * 1024 * 1024


def ws_max_size() -> int | None:
    """Max inbound WebSocket frame size; MESHNET_WS_MAX_BYTES<=0 means unlimited."""
    raw = os.environ.get("MESHNET_WS_MAX_BYTES", "").strip()
    if not raw:
        return DEFAULT_WS_MAX_BYTES
    try:
        value = int(raw)
    except ValueError:
        return DEFAULT_WS_MAX_BYTES
    return None if value <= 0 else value


# Binary relay frame: JSON header + raw body in one WebSocket binary message,
# so activation bodies travel as bytes instead of base64 inside JSON. Same wire
# format as meshnet_relay.server — duplicated because node and relay ship as
# independent distributions.
BINARY_FRAME_MAGIC = b"MRF1"


def encode_binary_frame(header: dict, body: bytes) -> bytes:
    """Build one request-owned binary frame without base64 expansion.

    ``join`` makes one owned output frame rather than creating intermediate
    concatenation frames.  The layout is intentionally unchanged because the
    relay ships an independent copy of this codec.
    """
    header_bytes = json.dumps(header, separators=(",", ":")).encode()
    return b"".join((BINARY_FRAME_MAGIC, len(header_bytes).to_bytes(4, "big"), header_bytes, body))


def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]:
    if len(frame) < 8 or frame[:4] != BINARY_FRAME_MAGIC:
        raise ValueError("not a meshnet binary relay frame")
    header_len = int.from_bytes(frame[4:8], "big")
    header = json.loads(frame[8:8 + header_len].decode())
    # The slice is a request-owned body.  It cannot retain the enclosing relay
    # frame after callers finish processing it.
    return header, frame[8 + header_len:]


@dataclass(frozen=True)
class RelayBridgeInfo:
    peer_id: str
    relay_addr: str


def _make_envelope(topic: str, payload: dict, peer_id: str) -> dict:
    return {
        "topic": topic,
        "version": 1,
        "from_peer": peer_id,
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "msg_id": f"{peer_id}-{time.time_ns():x}",
        "ttl": 1,
        "payload": payload,
    }


def _max_concurrency_from_env() -> int:
    try:
        value = int(os.environ.get("MESHNET_RELAY_CONCURRENCY", DEFAULT_MAX_CONCURRENCY))
    except (TypeError, ValueError):
        return DEFAULT_MAX_CONCURRENCY
    return max(1, value)


class _LoopbackHttpClientPool:
    """Bounded worker-local HTTP/1.1 clients for relay loopback forwarding."""

    def __init__(self, base_url: str, timeout: float = 300.0) -> None:
        parsed = urllib.parse.urlsplit(base_url.rstrip("/"))
        if parsed.scheme not in {"http", "https"} or not parsed.hostname:
            raise ValueError(f"invalid local bridge URL: {base_url!r}")
        self._scheme = parsed.scheme
        self._host = parsed.hostname
        self._port = parsed.port
        self._base_path = parsed.path.rstrip("/")
        self._timeout = timeout
        self._local = threading.local()
        self._lock = threading.Lock()
        self._clients: set[http.client.HTTPConnection] = set()

    def _connection(self) -> http.client.HTTPConnection:
        connection = getattr(self._local, "connection", None)
        if connection is None:
            kind = http.client.HTTPSConnection if self._scheme == "https" else http.client.HTTPConnection
            connection = kind(self._host, self._port, timeout=self._timeout)
            self._local.connection = connection
            with self._lock:
                self._clients.add(connection)
        return connection

    def request(self, method: str, path: str, body: bytes, headers: dict):
        request_path = f"{self._base_path}{path if path.startswith('/') else '/' + path}"
        connection = self._connection()
        try:
            connection.request(method, request_path, body=body, headers=headers)
            return connection.getresponse()
        except Exception:
            self.discard()
            raise

    def discard(self) -> None:
        connection = getattr(self._local, "connection", None)
        if connection is None:
            return
        try:
            connection.close()
        finally:
            self._local.connection = None
            with self._lock:
                self._clients.discard(connection)

    def close(self) -> None:
        with self._lock:
            clients = tuple(self._clients)
            self._clients.clear()
        for connection in clients:
            connection.close()
        self._local.connection = None


class RelayHttpBridge:
    """Connect outbound to a relay and proxy relay HTTP requests to localhost.

    Requests are dispatched on a bounded worker pool (US-037) so a node that is
    the head of one route can still serve per-token ``/forward`` hops of another.
    Streaming responses (``text/event-stream``) are forwarded as multiple chunk
    frames sharing one ``request_id`` (US-036); everything else stays a single
    frame for backward compatibility.
    """

    def __init__(
        self,
        relay_url: str,
        peer_id: str,
        local_base_url: str,
        advertised_addr: str,
        reconnect_interval: float = 3.0,
        max_concurrency: int | None = None,
    ) -> None:
        self.relay_url = relay_url.rstrip("/")
        self.peer_id = peer_id
        self.local_base_url = local_base_url.rstrip("/")
        self.advertised_addr = advertised_addr
        self.reconnect_interval = reconnect_interval
        self.max_concurrency = max(1, max_concurrency) if max_concurrency else _max_concurrency_from_env()
        self._running = False
        self._thread: threading.Thread | None = None
        self._connected = threading.Event()
        self._executor: ThreadPoolExecutor | None = None
        self._send_lock = threading.Lock()
        self._decode_log_lock = threading.Lock()
        self._decode_steps: dict[str, int] = {}
        self._ws = None
        self._loopback_clients = _LoopbackHttpClientPool(self.local_base_url)

    @property
    def relay_addr(self) -> str:
        base = self.relay_url
        if base.endswith("/ws"):
            base = base[:-3]
        return f"{base}/rpc/{self.peer_id}"

    def start(self) -> RelayBridgeInfo:
        self._running = True
        self._executor = ThreadPoolExecutor(
            max_workers=self.max_concurrency, thread_name_prefix="relay-http-worker"
        )
        self._thread = threading.Thread(target=self._run, daemon=True, name="relay-http-bridge")
        self._thread.start()
        return RelayBridgeInfo(peer_id=self.peer_id, relay_addr=self.relay_addr)

    def wait_connected(self, timeout: float = 5.0) -> bool:
        return self._connected.wait(timeout)

    def stop(self) -> None:
        self._running = False
        if self._thread:
            self._thread.join(timeout=3.0)
        if self._executor is not None:
            self._executor.shutdown(wait=False)
        self._loopback_clients.close()

    def _run(self) -> None:
        import websockets.sync.client as wsc  # type: ignore[import]

        while self._running:
            try:
                with wsc.connect(
                    self.relay_url, open_timeout=5, max_size=ws_max_size(), compression=None,
                ) as ws:
                    self._ws = ws
                    self._connected.set()
                    ws.send(json.dumps(_make_envelope(
                        "peer-register",
                        {"peer_id": self.peer_id, "addr": self.advertised_addr},
                        self.peer_id,
                    )))
                    while self._running:
                        try:
                            raw = ws.recv(timeout=1)
                        except TimeoutError:
                            continue
                        if isinstance(raw, (bytes, bytearray)):
                            try:
                                payload, body = decode_binary_frame(bytes(raw))
                            except (ValueError, json.JSONDecodeError):
                                continue
                            if payload.get("target_peer") not in {None, self.peer_id}:
                                continue
                            if self._executor is None:
                                break
                            self._executor.submit(self._process_request, payload, body)
                            continue
                        try:
                            envelope = json.loads(raw)
                        except (TypeError, json.JSONDecodeError):
                            continue
                        if envelope.get("topic") != "relay-http-request":
                            continue
                        payload = envelope.get("payload", {})
                        if payload.get("target_peer") not in {None, self.peer_id}:
                            continue
                        if self._executor is None:
                            break
                        self._executor.submit(self._process_request, payload)
            except Exception as exc:
                self._connected.clear()
                self._ws = None
                if self._running:
                    log.debug("relay bridge disconnected: %s", exc)
                    time.sleep(self.reconnect_interval)
        self._ws = None

    def _send_response_frame(self, payload: dict) -> bool:
        """Send one relay-http-response frame; False if the socket is gone.

        The lock is held per frame so concurrent workers interleave whole
        frames on the shared WebSocket, never torn ones.
        """
        ws = self._ws
        if ws is None:
            return False
        message = json.dumps(_make_envelope("relay-http-response", payload, self.peer_id))
        try:
            with self._send_lock:
                ws.send(message)
            return True
        except Exception as exc:
            log.debug("relay bridge send failed (request orphaned): %s", exc)
            return False

    def _send_binary_response_frame(self, header: dict, body: bytes) -> bool:
        """Send one binary response frame; False if the socket is gone."""
        ws = self._ws
        if ws is None:
            return False
        frame = encode_binary_frame(header, body)
        try:
            with self._send_lock:
                ws.send(frame)
            return True
        except Exception as exc:
            log.debug("relay bridge binary send failed (request orphaned): %s", exc)
            return False

    def _process_request(self, payload: dict, binary_body: bytes | None = None) -> None:
        request_id = str(payload.get("request_id") or "")
        method = str(payload.get("method") or "POST").upper()
        path = str(payload.get("path") or "/")
        headers = payload.get("headers") if isinstance(payload.get("headers"), dict) else {}
        binary_mode = binary_body is not None

        session = str(headers.get("X-Meshnet-Session") or "")
        cache_mode = headers.get("X-Meshnet-Cache")
        req_suffix = f" request_id={request_id}" if request_id else ""
        if path == "/forward" and cache_mode == "decode" and session:
            with self._decode_log_lock:
                steps = self._decode_steps.get(session, 0) + 1
                self._decode_steps[session] = steps
            if steps == 1 or steps % 32 == 0:
                print(
                    f"  [node] relay {method} {path} session={session[:8]} steps={steps}{req_suffix}",
                    flush=True,
                )
        else:
            session_suffix = f" session={session[:8]}" if session else ""
            print(f"  [node] relay {method} {path}{session_suffix}{req_suffix}", flush=True)

        if binary_mode:
            data = binary_body
        else:
            # Legacy JSON request: body_base64 carries binary data, text "body"
            # covers non-binary requests.
            body_b64 = payload.get("body_base64")
            if body_b64:
                data = base64.b64decode(body_b64)
            else:
                body_text = payload.get("body") or ""
                data = body_text.encode() if isinstance(body_text, str) else bytes(body_text)

        try:
            resp = self._loopback_clients.request(method, path, data, headers)
            try:
                resp_headers = dict(resp.headers)
                content_type = resp.headers.get("Content-Type", "")
                if "text/event-stream" in content_type:
                    if not self._stream_response(request_id, resp, resp_headers):
                        self._loopback_clients.discard()
                    return
                resp_bytes = resp.read()
                # Forward all X-Meshnet-* headers so the caller can reconstruct the activation.
                if binary_mode:
                    self._send_binary_response_frame({
                        "request_id": request_id,
                        "status": resp.status,
                        "headers": resp_headers,
                    }, resp_bytes)
                    return
                is_binary = "octet-stream" in content_type
                result: dict = {
                    "request_id": request_id,
                    "status": resp.status,
                    "headers": resp_headers,
                }
                if is_binary:
                    result["body_base64"] = base64.b64encode(resp_bytes).decode()
                else:
                    result["body"] = resp_bytes.decode(errors="replace")
                self._send_response_frame(result)
            finally:
                resp.close()
        except http.client.HTTPException as exc:
            self._loopback_clients.discard()
            self._send_response_frame({
                "request_id": request_id,
                "status": 503,
                "headers": {"Content-Type": "application/json"},
                "body": json.dumps({"error": f"relay bridge local request failed: {exc}"}),
            })
        except Exception as exc:
            self._loopback_clients.discard()
            self._send_response_frame({
                "request_id": request_id,
                "status": 503,
                "headers": {"Content-Type": "application/json"},
                "body": json.dumps({"error": f"relay bridge local request failed: {exc}"}),
            })

    def _stream_response(self, request_id: str, resp, resp_headers: dict) -> bool:
        """Forward an SSE response as chunk frames, one per complete SSE event.

        Frame order: header frame (status + headers), chunk frames, done frame.
        A receiver that sees no ``stream`` key treats the frame as a complete
        legacy response, so non-streaming peers are unaffected.
        """
        sent = self._send_response_frame({
            "request_id": request_id,
            "status": resp.status,
            "headers": resp_headers,
            "stream": True,
            "done": False,
        })
        if not sent:
            return False
        event_lines: list[str] = []
        for raw_line in resp:
            line = raw_line.decode(errors="replace")
            event_lines.append(line)
            if line.strip():
                continue
            # Blank line terminates one SSE event — flush it as a frame.
            if not self._send_response_frame({
                "request_id": request_id,
                "stream": True,
                "chunk": "".join(event_lines),
                "done": False,
            }):
                return False
            event_lines = []
        if event_lines:
            if not self._send_response_frame({
                "request_id": request_id,
                "stream": True,
                "chunk": "".join(event_lines),
                "done": False,
            }):
                return False
        return self._send_response_frame({
            "request_id": request_id,
            "stream": True,
            "done": True,
        })


def _peer_id_suffix(value: str) -> str:
    """Return a relay-safe suffix for a human node name or numeric instance id."""
    suffix = re.sub(r"[^A-Za-z0-9_.-]+", "-", value.strip()).strip("-._")
    return suffix[:32]


def peer_id_from_wallet(
    wallet_address: str,
    *,
    node_name: str | None = None,
    advertised_addr: str | None = None,
) -> str:
    """Build a per-node relay peer id from the wallet plus node identity.

    Multiple nodes can legitimately share one wallet for payouts, but the relay
    registry is keyed by peer_id.  Using only the wallet prefix makes those
    nodes overwrite each other at the relay.  Prefer the operator-provided node
    name; if absent, use the advertised endpoint port as the stable integer
    instance suffix (7001, 7002, ... for local multi-node runs).
    """
    wallet_prefix = wallet_address[:16] if len(wallet_address) >= 16 else wallet_address
    suffix = _peer_id_suffix(node_name or "") if node_name else ""
    if not suffix and advertised_addr:
        parsed = urllib.parse.urlparse(advertised_addr)
        if parsed.port is not None:
            suffix = str(parsed.port)
    return f"{wallet_prefix}-{suffix}" if suffix else wallet_prefix