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

from __future__ import annotations

import base64
import json
import logging
import os
import re
import threading
import time
import urllib.parse
import urllib.error
import urllib.request
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:
    header_bytes = json.dumps(header, separators=(",", ":")).encode()
    return 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())
    return header, bytes(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 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

    @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)

    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)

        url = f"{self.local_base_url}{path}"
        req = urllib.request.Request(url, data=data, headers=headers, method=method)
        try:
            with urllib.request.urlopen(req, timeout=300.0) as resp:
                resp_headers = dict(resp.headers)
                content_type = resp.headers.get("Content-Type", "")
                if "text/event-stream" in content_type:
                    self._stream_response(request_id, resp, resp_headers)
                    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)
        except urllib.error.HTTPError as exc:
            self._send_response_frame({
                "request_id": request_id,
                "status": exc.code,
                "headers": {"Content-Type": exc.headers.get("Content-Type", "application/json")},
                "body": exc.read().decode(errors="replace"),
            })
        except Exception as exc:
            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) -> None:
        """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
        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
            event_lines = []
        if event_lines:
            if not self._send_response_frame({
                "request_id": request_id,
                "stream": True,
                "chunk": "".join(event_lines),
                "done": 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