"""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


@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._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) 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
                        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 _process_request(self, payload: dict) -> 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 {}

        req_suffix = f" request_id={request_id}" if request_id else ""
        print(f"  [node] relay {method} {path}{req_suffix}", flush=True)

        # body_base64 carries binary data (e.g. bfloat16 activation tensors) safely.
        # Fallback to text "body" for backward-compat with 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.
                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