// Proxies.sx Peer SDK - Community Edition (Go)
// =====================================
// Open-source community tool. Provided "AS IS", WITHOUT WARRANTY OF ANY KIND,
// express or implied. MIT License. Maintained by community contributors via
// github.com/bolivian-peru; it is NOT a warranted product of Proxies.sx, and
// the operators of Proxies.sx accept no liability for its use. Run it at your
// own discretion. See the Peer Partner Agreement for the terms of participation.
// =====================================
// A drop-in Go peer that implements the COMPLETE relay protocol, wire-identical
// to the canonical Node reference (agents.proxies.sx/peer/reference-sdk.js, v1.5.2):
//   - Registers with the platform (reuses the same device identity across restarts)
//   - Connects to the nearest relay over WebSocket
//   - Heartbeats every 30s
//   - Handles `tunnel_connect` - opens a local TCP socket to the target
//   - Forwards bytes both directions (fast binary frames)
//   - Reports `tunnel_open_failed` so customers fail fast
//   - Honors server-driven nearest-relay routing (`relay_redirect`)
//   - Uses one WebSocket by default, with configurable parallel sockets
//   - Fails over between relays, detects dead sockets client-side, and
//     refreshes its token before it expires (v1.6.0; the JS reference does
//     not do these yet)
//
// This file mirrors the JS reference function-for-function so the two read
// side-by-side. One dependency: gorilla/websocket.
//
// USAGE:
//
//	go run reference-sdk.go -key=psx_...
//	API_KEY=psx_... go run reference-sdk.go
//	go build -o proxies-peer reference-sdk.go && ./proxies-peer -key=psx_...
//
// RESILIENCE KNOBS (env only; the value that restores pre-1.6 behaviour is
// noted so a regression can be bisected without a rebuild):
//
//	RELAY_FALLBACKS=wss://relay.proxies.sx,wss://relay-us.proxies.sx   failover ladder (RELAY_URLS accepted as alias)
//	RELAY_FAILOVER_THRESHOLD=4          dial failures in a row per relay; 0 = never fail over
//	RELAY_MAX_SOCKETS=5                 one below the relay's per-device cap of 6 (-connections is clamped to it)
//	WS_READ_TIMEOUT_S=90                seconds of silence before a socket is declared dead; 0 = never
//	TCP_KEEPALIVE_S=30                  TCP keepalive probe interval; 0 = Go default (15 s), -1 = off
//	TOKEN_REFRESH_LEAD_S=600            refresh the 1 h device JWT this long before `exp`; 0 = never
//	RECONNECT_BACKOFF_MAX_S=300         cap for all reconnect failures; 5 = flat 5-7 s
//
// Each process holds wsConnections WebSockets to one relay and serves customer
// tunnels. Run one process per peer device.
package main

import (
	"bytes"
	"context"
	cryptoRand "crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/gorilla/websocket"
)

// CHANGELOG
//
// 1.6.0 - resilience, motivated by relay-side measurements (2026-09-04):
//   - Relay failover (RELAY_FALLBACKS). Before, a dead relay was redialled forever.
//     Dial failures are counted per relay AND per socket goroutine: the Node
//     reference's single global counter (threshold 4) trips on ONE 5 s relay
//     restart, bounces the fleet to relay-us, gets geo-redirected back, and
//     costs >=8 disconnects per restart.
//   - Client-side dead-socket detection: 90 s read deadline extended on every
//     frame/ping/pong, plus TCP keepalive. The relay's zombie reaper cuts a
//     socket after 120 s without a heartbeat and its ping reaper after ONE
//     missed 30 s pong; without a client-side deadline a half-open socket
//     sat "connected" until the kernel gave up, serving no tunnels.
//   - Proactive token refresh ~10 min before `exp`. The JWT lives 1 h and is
//     verified at the handshake only, so pre-1.6 the first reconnect >1 h
//     after start-up was guaranteed a 4002 + full re-register.
//   - -connections clamped to the relay's real cap (6, was 8). A 7th socket
//     is evicted with 4008 and takes its customer tunnels down with it.
//   - 4008/4019/4029 closes back off exponentially (5 s doubling to 5 min)
//     instead of the flat 5-7 s that fed the cap that closed us. Measured on
//     the US relay: 1,898 4029s in 7 min.
//
// 1.5.2 - canonical known-good build, supersedes the broken third-party v1.5.1
// (tls_dropped regression). Forwards raw bytes (net.DialTimeout, no peer-side
// TLS). See SDK-TUNNEL-FIX-AND-VERSIONS.md.
// 1.6.1 - bounded registration retry with Retry-After and startup jitter;
// atomic checked identity persistence, collision-resistant generated names,
// identity/key binding and version reporting. Relay wire protocol unchanged.
// 1.6.4 - target safety: targets are dialled over tcp4 only and every address
// actually dialled is checked after resolution; private/LAN/special ranges are
// refused with tunnel_open_failed reason "target_blocked". The relay's legacy
// path can forward a hostname its own DNS could not check, and a name can
// resolve differently here (LAN-only names, DNS rebinding). IPv4 only because
// the relay, geo and exit qualification are IPv4. Wire protocol unchanged.
// 1.6.5 - every failed or short-lived connection backs off, including HTTP
// 429/502/503 and dial timeouts. Retry-After is respected (bounded by the
// configured cap); sustained inbound liveness resets the failure streak.
// 1.6.6 - authenticated live siblings prevent needless relay failover; stale
// dial attempts cannot change a newer relay selection. New configurations use
// one socket; explicit connection settings and saved identities are preserved.
// 1.6.7 - only the current relay socket/generation/device may redirect the
// shared selection; stale redirects cannot retire a healthy successor relay.
// 1.6.8 - retain the latest assigned relay in fallback and atomic saved state.
const sdkVersion = "1.6.8"

// ---------------------------------------------------------------------
// CONFIG - flags whose defaults come from environment variables, so the
// same binary works with `-key=...` flags OR API_KEY=... env (matching the
// Node reference, which is env-driven). No config framework: flags + env.
// ---------------------------------------------------------------------
var stateFileMu sync.Mutex // Serialize relay persistence with identity replacement.

var (
	apiBase          string
	relayFlag        string
	apiKey           string
	agentName        string
	wallet           string
	country          string
	carrier          string
	connectionMethod string
	wsConnections    int
	stateFile        string
	verbose          bool
	showVersion      bool
	registerJitter   time.Duration

	// v1.6.0 resilience knobs, env-only (see the header for defaults + why).
	relayList              []string
	relayFailoverThreshold int
	relayMaxSockets        int
	wsReadTimeout          time.Duration
	tcpKeepalive           time.Duration
	tokenRefreshLead       time.Duration
	reconnectBackoffMax    time.Duration
)

func envOr(key, dflt string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return dflt
}

func setupConfig() {
	home, err := os.UserHomeDir()
	if err != nil {
		home = "."
	}
	defaultState := filepath.Join(home, ".proxies-peer-state.json")
	// Empty sentinel: when neither -name nor AGENT_NAME is set we resolve the
	// name in acquireIdentity() - reuse the one saved in the state file, else
	// generate a random one ONCE and persist it. (v1.3.2) This makes the SDK
	// stable-by-default: omitting -name no longer spawns a new device per launch.
	defaultName := envOr("AGENT_NAME", "")

	flag.StringVar(&apiBase, "api", envOr("API_BASE", "https://api.proxies.sx/v1"), "Platform API base URL")
	flag.StringVar(&relayFlag, "relay", os.Getenv("RELAY_URL"), "Pin a relay URL (empty = let the server pick the nearest)")
	flag.StringVar(&apiKey, "key", os.Getenv("API_KEY"), "Your Proxies.sx API key (required)")
	flag.StringVar(&agentName, "name", defaultName, "Friendly agent name shown in your dashboard")
	flag.StringVar(&wallet, "wallet", os.Getenv("WALLET"), "Payout wallet address (optional)")
	flag.StringVar(&country, "country", envOr("COUNTRY", "US"), "Your ISO-2 country hint (server still verifies the real IP)")
	flag.StringVar(&carrier, "carrier", envOr("CARRIER", "unknown"), "Your ISP / carrier name (informational)")
	flag.StringVar(&connectionMethod, "method", envOr("CONNECTION_METHOD", "agent"), "Connection-type badge (agent/vps/docker)")
	flag.IntVar(&wsConnections, "connections", atoiOr(envOr("WS_CONNECTIONS", "1"), 1), "Parallel relay sockets (default 1; hard maximum 6)")
	flag.StringVar(&stateFile, "state", envOr("PEER_STATE_FILE", defaultState), "Path to the saved identity file")
	flag.BoolVar(&verbose, "verbose", os.Getenv("VERBOSE") == "true", "Log every tunnel open/close")
	flag.BoolVar(&showVersion, "version", false, "Print SDK version and exit without connecting")
	flag.Parse()
	if showVersion {
		fmt.Println("proxies-peer " + sdkVersion)
		os.Exit(0)
	}

	if apiKey == "" {
		log.Fatalln("[FATAL] API key required. Use -key=psx_... or set API_KEY=psx_...")
	}

	// Same ladder name and default as reference-sdk.js: RELAY_FALLBACKS. A
	// set-nothing user on the EU hub gets the US relay as a fallback, which is
	// the exact case (EU relay down, peers stranded) failover exists for.
	// RELAY_URLS is honoured as an alias for anyone who already set it.
	ladder := os.Getenv("RELAY_FALLBACKS")
	if ladder == "" {
		ladder = os.Getenv("RELAY_URLS")
	}
	if ladder == "" {
		ladder = "wss://relay.proxies.sx,wss://relay-us.proxies.sx"
	}
	for _, u := range strings.Split(ladder, ",") {
		if u = strings.TrimSpace(u); u != "" {
			relayList = append(relayList, u)
		}
	}
	relayFailoverThreshold = atoiOr(os.Getenv("RELAY_FAILOVER_THRESHOLD"), 4)
	relayMaxSockets = atoiOr(os.Getenv("RELAY_MAX_SOCKETS"), 5)
	if relayMaxSockets < 1 {
		relayMaxSockets = 1
	}
	if relayMaxSockets > 6 {
		relayMaxSockets = 6
	}
	jitterMS := atoiOr(os.Getenv("REGISTER_JITTER_MS"), 30000)
	if jitterMS < 0 {
		jitterMS = 0
	}
	if jitterMS > 3600000 {
		jitterMS = 3600000
	}
	registerJitter = time.Duration(jitterMS) * time.Millisecond
	wsReadTimeout = secondsEnv("WS_READ_TIMEOUT_S", 90)
	tcpKeepalive = secondsEnv("TCP_KEEPALIVE_S", 30)
	tokenRefreshLead = secondsEnv("TOKEN_REFRESH_LEAD_S", 600)
	reconnectBackoffMax = secondsEnv("RECONNECT_BACKOFF_MAX_S", 300)

	if wsConnections < 1 {
		wsConnections = 1
	} else if wsConnections > relayMaxSockets {
		logf("[WARN] -connections=%d exceeds the relay's per-device cap of %d; using %d (a socket above the cap is evicted with 4008 and its tunnels die with it)",
			wsConnections, relayMaxSockets, relayMaxSockets)
		wsConnections = relayMaxSockets
	}
}

func atoiOr(s string, dflt int) int {
	if n, err := strconv.Atoi(s); err == nil {
		return n
	}
	return dflt
}

func secondsEnv(key string, dflt int) time.Duration {
	return time.Duration(atoiOr(os.Getenv(key), dflt)) * time.Second
}

// ---------------------------------------------------------------------
// Shared, refreshable state. activeRelay can change via relay_redirect;
// the identity (token+deviceID) is refreshed on auth-failure re-register.
// Relay selection and socket health share one lock so failover cannot race ACKs.
// ---------------------------------------------------------------------
type sharedState struct {
	mu              sync.RWMutex
	token           string
	deviceID        string
	activeRelay     string
	assignedRelay   string // Latest server-selected relay; retained in the failover cycle.
	relayGeneration uint64
	sockets         map[*websocket.Conn]*relaySocket
}

// Socket evidence belongs to one relay selection and authenticated device.
type relaySocket struct {
	relay          string
	deviceID       string
	generation     uint64
	acknowledgedAt time.Time
	lastInbound    time.Time
	retiring       bool
}

func (s *sharedState) setIdentity(token, deviceID string) {
	s.mu.Lock()
	s.token, s.deviceID = token, deviceID
	s.mu.Unlock()
}
func (s *sharedState) getIdentity() (string, string) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.token, s.deviceID
}
func (s *sharedState) setRelay(url string) {
	s.mu.Lock()
	s.assignedRelay = url
	s.setRelayLocked(url, time.Now())
	s.mu.Unlock()
}

// Called with mu held by both server redirects and client failover.
func (s *sharedState) setRelayLocked(url string, now time.Time) {
	if s.activeRelay == url {
		return
	}
	s.relayGeneration++
	s.activeRelay = url
	// Retain an already authenticated socket on the selected target. Never
	// promote an unacknowledged socket's late ACK from an old selection.
	for _, socket := range s.sockets {
		if socket.relay == url && socket.deviceID == s.deviceID && socketLive(socket, now) {
			socket.generation = s.relayGeneration
		}
	}
}
func (s *sharedState) getRelay() string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.activeRelay
}

func (s *sharedState) relaySelection() (string, uint64) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.activeRelay, s.relayGeneration
}

func socketLive(socket *relaySocket, now time.Time) bool {
	freshness := wsReadTimeout
	if freshness <= 0 {
		freshness = 90 * time.Second // Disabling read deadlines must not create immortal health proof.
	}
	return !socket.retiring && !socket.acknowledgedAt.IsZero() && now.Sub(socket.lastInbound) < freshness
}

// Keep the proof check and target change under the same lock as ACK/liveness
// updates. The generation also rejects failures from a previous A -> B -> A trip.
func (s *sharedState) failover(cur, next string, generation uint64, now time.Time) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.activeRelay != cur || s.relayGeneration != generation {
		return false
	}
	for _, socket := range s.sockets {
		if socket.relay == cur && socket.generation == generation && socket.deviceID == s.deviceID && socketLive(socket, now) {
			return false
		}
	}
	s.setRelayLocked(next, now)
	return true
}

func (s *sharedState) addSocket(ws *websocket.Conn, relay, deviceID string, generation uint64) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.sockets == nil {
		s.sockets = make(map[*websocket.Conn]*relaySocket)
	}
	s.sockets[ws] = &relaySocket{relay: relay, deviceID: deviceID, generation: generation}
}

func (s *sharedState) removeSocket(ws *websocket.Conn) {
	s.mu.Lock()
	delete(s.sockets, ws)
	s.mu.Unlock()
}

// Only an ACK for this device authenticates a socket. Pings/data refresh that
// proof, but cannot create it or resurrect an obsolete connection generation.
func (s *sharedState) socketInbound(ws *websocket.Conn, ackDeviceID string, now time.Time) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	socket := s.sockets[ws]
	if socket == nil || socket.retiring || socket.generation != s.relayGeneration || socket.deviceID != s.deviceID {
		return false
	}
	socket.lastInbound = now
	if ackDeviceID != "" && ackDeviceID == socket.deviceID && socket.acknowledgedAt.IsZero() {
		socket.acknowledgedAt = now
	}
	return socket.relay == s.activeRelay && !socket.acknowledgedAt.IsZero() && now.Sub(socket.acknowledgedAt) >= reconnectHealthyAfter
}

// Snapshot before socket I/O: slow close writes must not hold the state lock.
func (s *sharedState) socketSnapshot() map[*websocket.Conn]string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	sockets := make(map[*websocket.Conn]string, len(s.sockets))
	for ws, socket := range s.sockets {
		sockets[ws] = socket.relay
	}
	return sockets
}

var (
	state               = &sharedState{}
	relayPinned         bool
	lastRedirectAt      int64 // unix ms
	reregisterScheduled int32
	ourRelayRe          = regexp.MustCompile(`(?i)^wss://[a-z0-9.-]+\.proxies\.sx(/|$)`)
)

// Registration is the /peer/agents/register (and /token/refresh) response.
type Registration struct {
	DeviceId string `json:"deviceId"`
	// AgentName is persisted (v1.3.2) so a restart reuses the SAME name even
	// when -name was omitted - re-registration then yields the SAME deviceId
	// (SHA-256(apiKey+name)) instead of a fresh random one. This is what stops
	// a fleet ballooning into duplicate devices when the launch command has no
	// -name. Resolved in acquireIdentity().
	AgentName     string          `json:"agentName"`
	APIKeyHash    string          `json:"apiKeyHash,omitempty"`
	Token         string          `json:"token"`
	Jwt           string          `json:"jwt"`
	RefreshToken  string          `json:"refreshToken"`
	Relay         string          `json:"relay"`
	EarningsPerGB json.RawMessage `json:"earningsPerGB"`
}

// WSMessage is the JSON control-plane envelope.
type WSMessage struct {
	Type    string          `json:"type"`
	Payload json.RawMessage `json:"payload,omitempty"`
}

func ts() string                      { return time.Now().Format("2006-01-02 15:04:05") }
func logf(f string, a ...interface{}) { log.Printf("["+ts()+"] "+f, a...) }
func vlogf(f string, a ...interface{}) {
	if verbose {
		log.Printf("["+ts()+"] "+f, a...)
	}
}

func main() {
	setupConfig()
	relayPinned = relayFlag != ""

	ctx, cancel := context.WithCancel(context.Background())
	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
	go func() {
		<-sig
		logf("Stopping. Goodbye.")
		cancel()
		os.Exit(0)
	}()

	// Spread both registration and restart refreshes across the configured window.
	if registerJitter > 0 {
		time.Sleep(time.Duration(rand.Int63n(int64(registerJitter))))
	}
	// 1. Acquire identity (resume saved one, else register).
	token, deviceID, relay, err := acquireIdentity()
	if err != nil {
		log.Fatalf("[FATAL] register failed: %v", err)
	}
	state.setIdentity(token, deviceID)

	// Relay precedence: manual pin > server geo-assignment > built-in default.
	switch {
	case relayPinned:
		state.setRelay(relayFlag)
	case relay != "":
		state.setRelay(relay)
	default:
		state.setRelay(relayList[0])
	}

	logf("[REGISTERED] device=%s relay=%s failover=%s sockets=%d", deviceID, state.getRelay(), strings.Join(relayList, ","), wsConnections)
	logf("Sharing bandwidth as %q. Leave this running. Ctrl+C to stop.", agentName)

	if tokenRefreshLead > 0 {
		go tokenRefreshLoop(ctx)
	}

	// 2. Launch a FIXED set of socket goroutines. Each one reconnects itself;
	//    no new goroutines are ever spawned (this is the fix vs a recursive
	//    re-register that leaks goroutine generations).
	var wg sync.WaitGroup
	for i := 0; i < wsConnections; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			connectLoop(ctx)
		}()
	}
	wg.Wait()
}

// connectLoop runs ONE socket for the life of the process: connect, serve,
// reconnect on close with backoff. On an auth-failure close it triggers a
// single shared re-register (which updates the token every socket reads).
//
// Dial failures are counted per relay, per goroutine (v1.6.0): four sockets
// each failing once during a 5 s relay restart must not read as four
// failures. One socket has to fail relayFailoverThreshold times in a row
// (with backoff between attempts) before this process rotates to the next
// RELAY_URLS entry, provided no authenticated sibling remains live. The counter
// resets on any successful dial or relay selection change.
func connectLoop(ctx context.Context) {
	var failRelay string
	var failGeneration uint64
	var dialFails int // consecutive dial failures against failRelay
	var failures int  // reset only after a sustained live connection, not an HTTP upgrade
	for {
		select {
		case <-ctx.Done():
			return
		default:
		}
		token, deviceID := state.getIdentity()
		relay, generation := state.relaySelection()
		if relay != failRelay || generation != failGeneration {
			failRelay, failGeneration, dialFails = relay, generation, 0
		}
		err := connect(ctx, token, deviceID, relay, generation, func() { failures = 0 })
		if errors.Is(err, errDial) {
			dialFails++
			if relayFailoverThreshold > 0 && dialFails >= relayFailoverThreshold {
				failoverRelay(relay, generation, dialFails)
				dialFails = 0
			}
		} else {
			dialFails = 0
		}
		if isAuthClose(err) {
			scheduleReregister()
		}

		failures++
		retryAfter := ""
		var rejected *relayDialError
		if errors.As(err, &rejected) {
			retryAfter = rejected.retryAfter
		}
		base := reconnectDelay(failures, retryAfter, time.Now(), reconnectBackoffMax)
		// Spread retries across a growing window at fleet scale. Keep
		// the server's Retry-After as a minimum and preserve the existing 2s
		// startup spread when an operator explicitly chooses a 5s cap.
		spread := base / 2
		if spread > reconnectBackoffMax-base {
			spread = reconnectBackoffMax - base
		}
		if spread < 2*time.Second {
			spread = 2 * time.Second
		}
		wait := base + time.Duration(rand.Int63n(int64(spread)))
		logf("[BACKOFF] %v - reconnecting in %s", err, wait)
		select {
		case <-ctx.Done():
			return
		case <-time.After(wait):
		}
	}
}

// reconnectDelay also covers failed HTTP upgrades and rapid accepted/closed
// sockets. A 101 response alone must not reset retries during relay overload.
func reconnectDelay(failures int, retryAfter string, now time.Time, cap time.Duration) time.Duration {
	if cap < 5*time.Second {
		cap = 5 * time.Second
	}
	delay := 5 * time.Second
	for i := 1; i < failures && delay < cap; i++ {
		if delay > cap/2 {
			delay = cap
			break
		}
		delay *= 2
	}
	// Share the registration path's checked seconds/HTTP-date parser.
	if requested := registrationRetryDelay(0, retryAfter, now); requested > delay {
		delay = requested
	}
	if delay > cap {
		delay = cap
	}
	return delay
}

// Keep the latest server assignment in the existing bounded failover cycle.
// Otherwise a temporarily unavailable spare outside RELAY_FALLBACKS is lost
// forever once the peer falls back to the primary relays. Healthy connections
// still shield the selection, and a pinned -relay never fails over.
func failoverRelay(cur string, generation uint64, fails int) {
	if relayPinned {
		return
	}
	state.mu.RLock()
	assigned := state.assignedRelay
	state.mu.RUnlock()
	candidates := append([]string(nil), relayList...)
	if assigned != "" {
		found := false
		for _, relay := range candidates {
			found = found || relay == assigned
		}
		if !found {
			candidates = append(candidates, assigned)
		}
	}
	if len(candidates) == 0 {
		return
	}
	next := candidates[0]
	for i, r := range candidates {
		if r == cur {
			next = candidates[(i+1)%len(candidates)]
			break
		}
	}
	if next == cur {
		logf("[FAILOVER] %s unreachable %d times in a row and RELAY_FALLBACKS has no alternative; retrying it", cur, fails)
		return
	}
	if state.failover(cur, next, generation, time.Now()) {
		logf("[FAILOVER] %s unreachable %d times in a row -> %s", cur, fails, next)
	}
}

func isAuthClose(err error) bool {
	if err == nil {
		return false
	}
	return websocket.IsCloseError(err, 4001, 4002) ||
		strings.Contains(err.Error(), "4001") || strings.Contains(err.Error(), "4002")
}

// errDial marks a connect() error as "never got a socket" (vs a socket that
// was accepted and later closed); only those count toward relay failover.
var errDial = errors.New("dial failed")

type relayDialError struct {
	status     int
	retryAfter string
}

func (e *relayDialError) Error() string { return fmt.Sprintf("dial failed: HTTP %d", e.status) }
func (e *relayDialError) Unwrap() error { return errDial }

// Injectable for the isolated liveness tests; not a device configuration knob.
var reconnectHealthyAfter = time.Minute

// connect opens one WebSocket, serves tunnels, and returns when it closes.
func connect(ctx context.Context, token, deviceID, url string, generation uint64, onHealthy func()) error {
	connectionState := state
	dialer := websocket.Dialer{
		Subprotocols:     []string{"token." + token}, // JWT in Sec-WebSocket-Protocol, never the URL
		HandshakeTimeout: 10 * time.Second,
		// TCP keepalive so a NAT that silently drops the mapping surfaces as
		// a read error instead of a socket that looks open forever.
		NetDialContext: (&net.Dialer{KeepAlive: tcpKeepalive}).DialContext,
	}
	ws, resp, err := dialer.DialContext(ctx, url, nil)
	if err != nil {
		if resp != nil {
			if resp.Body != nil {
				resp.Body.Close()
			}
			return &relayDialError{status: resp.StatusCode, retryAfter: resp.Header.Get("Retry-After")}
		}
		return fmt.Errorf("%w: %v", errDial, err)
	}
	defer ws.Close()
	logf("[CONNECTED] device=%s relay=%s", deviceID, url)
	connectionState.addSocket(ws, url, deviceID, generation)
	defer connectionState.removeSocket(ws)

	// Client-side liveness. The relay pings every 30 s and acks every
	// heartbeat, so a healthy socket delivers a frame at least every 30 s;
	// wsReadTimeout of silence means the path is dead even if TCP has not
	// noticed. Pings and pongs count as life too, hence the handlers.
	healthy := false
	extendDeadline := func(ackDeviceID string) {
		now := time.Now()
		if connectionState.socketInbound(ws, ackDeviceID, now) && !healthy {
			healthy = true
			if onHealthy != nil {
				onHealthy()
			}
		}
		if wsReadTimeout > 0 {
			_ = ws.SetReadDeadline(now.Add(wsReadTimeout))
		}
	}
	extendDeadline("")
	ws.SetPingHandler(func(data string) error {
		extendDeadline("")
		err := ws.WriteControl(websocket.PongMessage, []byte(data), time.Now().Add(time.Second))
		if err == websocket.ErrCloseSent {
			return nil
		}
		// Same tolerance as gorilla's default handler: a pong that could not be
		// written within 1 s because tunnel frames hold the write lock is a
		// temporary error, not a dead socket — closing here killed every tunnel
		// on a saturated uplink.
		if e, ok := err.(net.Error); ok && e.Temporary() {
			return nil
		}
		return err
	})
	ws.SetPongHandler(func(string) error {
		extendDeadline("")
		return nil
	})

	var writeMu sync.Mutex // gorilla forbids concurrent writes; one mutex per socket
	send := func(m WSMessage) error {
		b, _ := json.Marshal(m)
		writeMu.Lock()
		defer writeMu.Unlock()
		return ws.WriteMessage(websocket.TextMessage, b)
	}
	sendBinary := func(frame []byte) error {
		writeMu.Lock()
		defer writeMu.Unlock()
		return ws.WriteMessage(websocket.BinaryMessage, frame)
	}

	tunnels := sync.Map{} // sessionId -> net.Conn
	defer tunnels.Range(func(_, v interface{}) bool {
		if c, ok := v.(net.Conn); ok {
			c.Close()
		}
		return true
	})

	// Announce capabilities + initial state.
	_ = send(WSMessage{Type: "device_info", Payload: json.RawMessage(fmt.Sprintf(
		`{"country":%q,"carrier":%q,"currentIp":"auto","protocol":"binary-v1","supportsRelayRedirect":true,"sdkVersion":%q}`,
		country, carrier, sdkVersion))})

	// Heartbeat every 30s; stop it when this socket closes.
	hbCtx, hbCancel := context.WithCancel(ctx)
	defer hbCancel()
	go func() {
		t := time.NewTicker(30 * time.Second)
		defer t.Stop()
		for {
			select {
			case <-hbCtx.Done():
				ws.Close()
				return
			case <-t.C:
				_ = send(WSMessage{Type: "heartbeat"})
				// A control frame, not JSON: it does not count toward the
				// relay's 100 msg/min cap, and its pong feeds the read deadline.
				_ = ws.WriteControl(websocket.PingMessage, nil, time.Now().Add(time.Second))
			}
		}
	}()

	for {
		msgType, raw, err := ws.ReadMessage()
		if err != nil {
			logf("[CLOSED] %v", err)
			return err
		}
		extendDeadline("")

		// Hot path: binary tunnel frames. [0x01 data | 0x03 close][sidLen][sid][payload]
		if msgType == websocket.BinaryMessage && len(raw) >= 2 {
			action, sidLen := raw[0], int(raw[1])
			if len(raw) < 2+sidLen {
				continue
			}
			sid := string(raw[2 : 2+sidLen])
			payload := raw[2+sidLen:]
			if v, ok := tunnels.Load(sid); ok {
				sock := v.(net.Conn)
				if action == 0x01 {
					_, _ = sock.Write(payload)
				} else if action == 0x03 {
					sock.(*targetConn).finish()
				}
			}
			continue
		}

		var msg WSMessage
		if json.Unmarshal(raw, &msg) != nil {
			continue
		}
		switch msg.Type {
		case "connected":
			var ack struct {
				DeviceID string `json:"deviceId"`
			}
			if json.Unmarshal(msg.Payload, &ack) == nil && ack.DeviceID == deviceID {
				extendDeadline(ack.DeviceID)
				logf("[ACK] relay confirmed connection")
			}
		case "heartbeat_ack":
			// no-op
		case "relay_redirect":
			var p struct{ Relay, Reason string }
			if json.Unmarshal(msg.Payload, &p) == nil {
				connectionState.handleRedirect(ws, p.Relay, p.Reason)
			}
		case "tunnel_connect":
			var p struct {
				SessionId string      `json:"sessionId"`
				Host      string      `json:"host"`
				Port      interface{} `json:"port"`
			}
			if json.Unmarshal(msg.Payload, &p) == nil {
				go handleTunnelConnect(send, sendBinary, &tunnels, p.SessionId, p.Host, portToString(p.Port))
			}
		case "tunnel_close":
			var p struct {
				SessionId string `json:"sessionId"`
			}
			if json.Unmarshal(msg.Payload, &p) == nil {
				if v, ok := tunnels.Load(p.SessionId); ok {
					v.(*targetConn).finish()
				}
			}
		case "rotate_ip_request":
			// Desktop/server peers can't toggle a carrier IP; report unsupported.
			var p struct {
				RequestId string `json:"requestId"`
			}
			if json.Unmarshal(msg.Payload, &p) == nil {
				_ = send(WSMessage{Type: "rotation_complete", Payload: json.RawMessage(
					fmt.Sprintf(`{"requestId":%q,"success":false,"error":"not_supported_on_server_peer"}`, p.RequestId))})
			}
		default:
			// Unknown control messages are ignored, never fatal.
		}
	}
}

// portToString accepts a JSON number or string port and returns "443" etc.
func portToString(p interface{}) string {
	switch v := p.(type) {
	case float64:
		return strconv.Itoa(int(v))
	case string:
		return v
	default:
		return fmt.Sprintf("%v", v)
	}
}

const targetQueueLimit = 1 << 20
const totalTargetQueueLimit = 16 << 20
const targetWriteTimeout = 15 * time.Second

var queuedTargetBytes int64
var errTargetBackpressure = errors.New("target write queue full")
var errTargetBlocked = errors.New("target is not a public IPv4 address")
var dialTarget = dialPublicIPv4

// dialPublicIPv4 dials over tcp4 and vets each resolved address in Control,
// i.e. the address really dialled, never the name. A blocked address fails
// the dial with errTargetBlocked before any packet leaves this host.
func dialPublicIPv4(network, address string, timeout time.Duration) (net.Conn, error) {
	d := net.Dialer{Timeout: timeout, Control: func(_, addr string, _ syscall.RawConn) error {
		host, _, err := net.SplitHostPort(addr)
		if err != nil {
			return err
		}
		if ip := net.ParseIP(host).To4(); ip == nil || blockedV4(ip) {
			return errTargetBlocked
		}
		return nil
	}}
	return d.Dial(network, address)
}

// blockedV4 refuses the same IPv4 ranges as the relay:
// relay-server/src/security/ip-blocking.ts blockedV4() (this, private,
// loopback, CGNAT, link-local, benchmark, multicast/reserved). Keep in step.
func blockedV4(ip net.IP) bool {
	a, b := ip[0], ip[1]
	return a == 0 || a == 10 || a == 127 || a >= 224 ||
		(a == 100 && b >= 64 && b <= 127) ||
		(a == 169 && b == 254) || (a == 172 && b >= 16 && b <= 31) ||
		(a == 192 && b == 168) || (a == 198 && (b == 18 || b == 19))
}

// A target that stops reading must not block the WebSocket reader (including
// control PONGs and unrelated tunnels). Bound queued + in-flight bytes both per
// target and across every WebSocket in this process; close only the slow target.
type targetConn struct {
	net.Conn
	mu           sync.Mutex
	queue        []byte
	wake         chan struct{}
	done         chan struct{}
	closed       bool
	finishing    bool
	finishBy     time.Time
	pendingBytes int
}

func newTargetConn(conn net.Conn) *targetConn {
	t := &targetConn{Conn: conn, wake: make(chan struct{}, 1), done: make(chan struct{})}
	go t.writeLoop()
	return t
}

func (t *targetConn) Write(p []byte) (int, error) {
	t.mu.Lock()
	if t.closed || t.finishing {
		t.mu.Unlock()
		return 0, net.ErrClosed
	}
	if len(p) > targetQueueLimit-t.pendingBytes {
		t.mu.Unlock()
		t.Close()
		return 0, errTargetBackpressure
	}
	if atomic.AddInt64(&queuedTargetBytes, int64(len(p))) > totalTargetQueueLimit {
		atomic.AddInt64(&queuedTargetBytes, -int64(len(p)))
		t.mu.Unlock()
		t.Close()
		return 0, errTargetBackpressure
	}
	t.pendingBytes += len(p)
	t.queue = append(t.queue, p...)
	select {
	case t.wake <- struct{}{}:
	default:
	}
	t.mu.Unlock()
	return len(p), nil
}

func (t *targetConn) release(p []byte) {
	t.mu.Lock()
	t.pendingBytes -= len(p)
	atomic.AddInt64(&queuedTargetBytes, -int64(len(p)))
	t.mu.Unlock()
}

func (t *targetConn) Close() error {
	t.mu.Lock()
	if t.closed {
		t.mu.Unlock()
		return nil
	}
	t.closed = true
	close(t.done)
	t.mu.Unlock()
	return t.Conn.Close()
}

// Preserve frame order at a relay-requested FIN: preceding queued data drains
// before close. A lost WebSocket or write failure still calls Close immediately.
func (t *targetConn) finish() {
	t.mu.Lock()
	if t.closed || t.finishing {
		t.mu.Unlock()
		return
	}
	t.finishing = true
	t.finishBy = time.Now().Add(targetWriteTimeout)
	empty := t.pendingBytes == 0
	t.mu.Unlock()
	if empty {
		t.Close()
	}
}

func (t *targetConn) writeLoop() {
	defer func() {
		t.mu.Lock()
		t.pendingBytes -= len(t.queue)
		atomic.AddInt64(&queuedTargetBytes, -int64(len(t.queue)))
		t.queue = nil
		t.mu.Unlock()
	}()
	for {
		select {
		case <-t.done:
			return
		case <-t.wake:
			t.mu.Lock()
			p := t.queue
			t.queue = nil
			deadline := time.Now().Add(targetWriteTimeout)
			if t.finishing && t.finishBy.Before(deadline) {
				deadline = t.finishBy
			}
			t.mu.Unlock()
			if len(p) == 0 {
				continue
			}
			err := t.Conn.SetWriteDeadline(deadline)
			if err == nil {
				var n int
				n, err = t.Conn.Write(p)
				if err == nil && n != len(p) {
					err = io.ErrShortWrite
				}
			}
			t.release(p)
			t.mu.Lock()
			finished := t.finishing && t.pendingBytes == 0
			t.mu.Unlock()
			if err != nil || finished {
				t.Close()
				return
			}
		}
	}
}

// handleTunnelConnect opens a raw TCP socket to the target and pipes bytes.
// For HTTPS (443) this is a PLAIN TCP relay - TLS stays end-to-end between the
// customer and the target. Do NOT terminate TLS here.
func handleTunnelConnect(send func(WSMessage) error, sendBinary func([]byte) error, tunnels *sync.Map, sid, host, port string) {
	vlogf("[TUNNEL_CONNECT] %s %s:%s", sid, host, port)
	start := time.Now()
	rawSock, err := dialTarget("tcp4", net.JoinHostPort(host, port), 30*time.Second)
	if err != nil {
		vlogf("[TUNNEL_FAIL] %s %v", sid, err)
		_ = send(WSMessage{Type: "tunnel_open_failed", Payload: json.RawMessage(fmt.Sprintf(
			`{"sessionId":%q,"reason":%q,"detail":%s,"durationMs":%d}`,
			sid, failReason(err), strconv.Quote(err.Error()), time.Since(start).Milliseconds()))})
		return
	}
	if tcp, ok := rawSock.(*net.TCPConn); ok {
		_ = tcp.SetNoDelay(true)
	}
	sock := newTargetConn(rawSock)
	tunnels.Store(sid, sock)
	vlogf("[TCP_OPEN_OK] %s", sid)
	defer func() {
		sock.Close()
		tunnels.Delete(sid)
		_ = send(WSMessage{Type: "tunnel_closed", Payload: json.RawMessage(fmt.Sprintf(`{"sessionId":%q}`, sid))})
	}()
	if err := send(WSMessage{Type: "tunnel_connected", Payload: json.RawMessage(fmt.Sprintf(`{"sessionId":%q}`, sid))}); err != nil {
		return // The WebSocket may have closed while the target dial was pending.
	}

	// Forward target->peer back to the relay until the socket closes.
	sidBytes := []byte(sid)
	buf := make([]byte, 32*1024)
	for {
		n, err := sock.Read(buf)
		if n > 0 {
			frame := make([]byte, 2+len(sidBytes)+n)
			frame[0] = 0x01
			frame[1] = byte(len(sidBytes))
			copy(frame[2:], sidBytes)
			copy(frame[2+len(sidBytes):], buf[:n])
			if sendBinary(frame) != nil {
				return
			}
		}
		if err != nil {
			return
		}
	}
}

func failReason(err error) string {
	if errors.Is(err, errTargetBlocked) {
		return "target_blocked"
	}
	if ne, ok := err.(net.Error); ok && ne.Timeout() {
		return "timeout"
	}
	s := err.Error()
	switch {
	case strings.Contains(s, "refused"):
		return "connection_refused"
	case strings.Contains(s, "no such host"):
		return "dns_failed"
	case strings.Contains(s, "timeout"):
		return "timeout"
	default:
		return "unknown"
	}
}

// handleRedirect accepts routing instructions only from the current socket
// selection. Current pre-ACK redirects remain valid: some relays redirect before
// sending connected. An obsolete socket must not retire a healthy successor.
func (s *sharedState) handleRedirect(source *websocket.Conn, target, reason string) {
	old, obsolete, accepted := s.redirectSockets(source, target)
	if !accepted {
		return
	}
	logf("[REDIRECT] %s -> %s (%s)", old, target, reason)
	for _, ws := range obsolete {
		_ = ws.WriteControl(websocket.CloseMessage,
			websocket.FormatCloseMessage(4100, "relay_redirect"), time.Now().Add(time.Second))
		ws.Close()
	}
}

// Commit selection and socket retirement together; perform close I/O afterward.
func (s *sharedState) redirectSockets(source *websocket.Conn, target string) (string, []*websocket.Conn, bool) {
	if relayPinned || !ourRelayRe.MatchString(target) {
		return "", nil, false
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	socket := s.sockets[source]
	if socket == nil || socket.retiring || socket.relay != s.activeRelay ||
		socket.generation != s.relayGeneration || socket.deviceID != s.deviceID || target == s.activeRelay {
		return "", nil, false
	}
	now := time.Now()
	previous := atomic.LoadInt64(&lastRedirectAt)
	if now.UnixMilli()-previous < 60_000 || !atomic.CompareAndSwapInt64(&lastRedirectAt, previous, now.UnixMilli()) {
		return "", nil, false
	}
	old := s.activeRelay
	s.assignedRelay = target
	if atomic.LoadInt32(&reregisterScheduled) == 0 {
		if err := persistAssignedRelay(s.deviceID, target); err != nil {
			logf("[STATE] could not persist relay assignment: %v", err)
		}
	} // An active identity renewal persists the latest assignment when it completes.
	s.setRelayLocked(target, now)
	var obsolete []*websocket.Conn
	for ws, socket := range s.sockets {
		if socket.relay != target {
			// A concurrent failover cannot adopt proof from a socket already
			// committed to close, even while its close write is still pending.
			socket.retiring = true
			obsolete = append(obsolete, ws)
		}
	}
	return old, obsolete, true
}

// scheduleReregister refreshes the shared identity ONCE on auth failure, so N
// sockets closing on an expired token trigger one re-register, not N.
func scheduleReregister() {
	if !atomic.CompareAndSwapInt32(&reregisterScheduled, 0, 1) {
		return
	}
	go func() {
		time.Sleep(5 * time.Second)
		defer func() {
			state.mu.Lock()
			if !relayPinned && state.assignedRelay != "" {
				if err := persistAssignedRelay(state.deviceID, state.assignedRelay); err != nil {
					logf("[STATE] could not persist relay assignment after identity renewal: %v", err)
				}
			}
			atomic.StoreInt32(&reregisterScheduled, 0)
			state.mu.Unlock()
		}()
		_, startedGeneration := state.relaySelection()
		token, deviceID, relay, err := acquireIdentity()
		if err != nil {
			logf("[REREGISTER] failed: %v (will retry on next reconnect)", err)
			return
		}
		state.mu.Lock()
		sameDevice := state.deviceID == deviceID
		state.token, state.deviceID = token, deviceID
		// A redirect accepted during the request is newer than its saved relay.
		if !relayPinned && relay != "" && (!sameDevice || state.relayGeneration == startedGeneration) {
			state.assignedRelay = relay
			state.setRelayLocked(relay, time.Now())
		}
		state.mu.Unlock()
		logf("[REREGISTER] refreshed identity device=%s", deviceID)
	}()
}

// tokenRefreshLoop keeps the shared token fresh. The device JWT lives 1 h and
// the relay verifies it at the handshake only, so without this the first
// reconnect more than an hour after start-up was guaranteed a 4002 and a full
// re-register. The refresh token lives in the state file (7-day server TTL);
// when a refresh fails the 4002 path still re-registers exactly as before.
func tokenRefreshLoop(ctx context.Context) {
	var lastToken string
	var due time.Time
	for {
		token, deviceID := state.getIdentity()
		if token != lastToken {
			lastToken = token
			due = tokenExpiry(token).Add(-tokenRefreshLead)
			if floor := time.Now().Add(time.Minute); due.Before(floor) {
				due = floor // never hammer the API, whatever the lead is set to
			}
		}
		if wait := time.Until(due); wait > 0 {
			select {
			case <-ctx.Done():
				return
			case <-time.After(wait):
			}
			continue // re-read: a re-register may have replaced the token meanwhile
		}
		if saved, err := loadState(); err == nil && saved != nil && saved.DeviceId == deviceID && saved.RefreshToken != "" {
			if t := refreshSavedToken(deviceID, saved.RefreshToken); t != "" {
				state.setIdentity(t, deviceID)
				logf("[TOKEN] refreshed before expiry; next refresh at %s", tokenExpiry(t).Add(-tokenRefreshLead).Format("15:04:05"))
				continue
			}
		}
		// A refresh that fails is not retried in a loop: the backend caps refresh
		// tokens at 7 days from REGISTRATION, so past that every retry is a 401
		// and a fleet would hammer /token/refresh once a minute per device
		// forever. Re-registering (the existing path) issues a fresh pair.
		logf("[TOKEN] proactive refresh failed; re-registering to obtain a fresh token pair")
		scheduleReregister()
		due = time.Now().Add(tokenRefreshLead)
	}
}

// tokenExpiry reads the JWT `exp` claim; an unreadable token is assumed to
// live the server's documented 1 h from now.
func tokenExpiry(token string) time.Time {
	if parts := strings.Split(token, "."); len(parts) == 3 {
		if b, err := base64.RawURLEncoding.DecodeString(strings.TrimRight(parts[1], "=")); err == nil {
			var c struct {
				Exp int64 `json:"exp"`
			}
			if json.Unmarshal(b, &c) == nil && c.Exp > 0 {
				return time.Unix(c.Exp, 0)
			}
		}
	}
	return time.Now().Add(time.Hour)
}

// ---------------------------------------------------------------------
// REGISTER / IDENTITY
// ---------------------------------------------------------------------
func acquireIdentity() (token, deviceID, relay string, err error) {
	saved, err := loadState()
	if err != nil {
		return "", "", "", fmt.Errorf("read identity state: %w", err)
	}
	keyHash := fmt.Sprintf("%x", sha256.Sum256([]byte(apiKey)))
	if saved != nil {
		if saved.APIKeyHash != "" && saved.APIKeyHash != keyHash {
			return "", "", "", errors.New("state belongs to another API key; restore that key or use a separate state path for a new identity")
		}
		if agentName != "" && saved.AgentName != "" && agentName != saved.AgentName {
			return "", "", "", errors.New("name differs from saved identity; restore the original name or use a separate state path")
		}
	}

	// Resolve a STABLE name when the user didn't pass -name / AGENT_NAME.
	// Reuse the name saved in the state file so a restart re-registers with the
	// SAME deviceId (SHA-256(apiKey+name)) instead of a fresh random one - even
	// if the refresh below fails (expired token, or the device was cleaned up
	// server-side). Generate-and-persist only on a truly first run. This is the
	// single most important guard against duplicate-device churn.
	if agentName == "" {
		if saved != nil && saved.AgentName != "" {
			agentName = saved.AgentName
			logf("[IDENTITY] reusing saved name %q (no -name given)", agentName)
		} else {
			var id [16]byte
			if _, e := cryptoRand.Read(id[:]); e != nil {
				return "", "", "", e
			}
			agentName = fmt.Sprintf("go-%x", id[:])
			logf("[IDENTITY] generated name %q and persisting it - pass -name to choose your own", agentName)
		}
	}

	// Legacy state is re-registered once with the same saved name to bind it to
	// this key; never reuse a different account's refresh token after a key change.
	if saved != nil && saved.APIKeyHash == keyHash && saved.DeviceId != "" && saved.RefreshToken != "" {
		if t := refreshSavedToken(saved.DeviceId, saved.RefreshToken); t != "" {
			logf("[RESUMED] device=%s (reused saved identity)", saved.DeviceId)
			return t, saved.DeviceId, saved.Relay, nil
		}
	}

	r, err := register()
	if err != nil {
		return "", "", "", err
	}
	r.AgentName = agentName // persist the resolved name → future launches are stable
	r.APIKeyHash = keyHash
	token = r.Jwt
	if token == "" {
		token = r.Token
	}
	if r.DeviceId == "" || token == "" || r.RefreshToken == "" {
		return "", "", "", errors.New("registration response missing device ID or tokens")
	}
	if err := saveState(r); err != nil {
		return "", "", "", fmt.Errorf("persist identity state: %w", err)
	}
	return token, r.DeviceId, r.Relay, nil
}

var sharedHTTP = &http.Client{Timeout: 15 * time.Second}

type registrationError struct {
	status     int
	retryAfter string
}

func (e *registrationError) Error() string { return fmt.Sprintf("registration HTTP %d", e.status) }

// Injectable clock/sleeper keep retry tests fast; production uses wall time.
var registrationSleep = time.Sleep
var registrationNow = time.Now
var registrationJitter = func() time.Duration { return time.Duration(rand.Int63n(int64(time.Second))) }

func registrationRetryDelay(attempt int, retryAfter string, now time.Time) time.Duration {
	delay := time.Second << uint(attempt)
	if delay > time.Minute {
		delay = time.Minute
	}
	if n, err := strconv.ParseInt(strings.TrimSpace(retryAfter), 10, 64); err == nil && n >= 0 && n <= int64((time.Duration(1<<63-1)-time.Second)/time.Second) {
		if d := time.Duration(n) * time.Second; d > delay {
			delay = d
		}
	} else if at, err := http.ParseTime(retryAfter); err == nil {
		if d := at.Sub(now); d > delay {
			delay = d
		}
	}
	return delay
}

func register() (*Registration, error) {
	for attempt := 0; ; attempt++ {
		r, err := registerOnce()
		if err == nil {
			return r, nil
		}
		var responseErr *registrationError
		if errors.As(err, &responseErr) && responseErr.status != 429 && responseErr.status < 500 {
			return nil, err
		}
		if attempt >= 4 {
			return nil, err
		}
		retryAfter := ""
		if responseErr != nil {
			retryAfter = responseErr.retryAfter
		}
		delay := registrationRetryDelay(attempt, retryAfter, registrationNow()) + registrationJitter()
		logf("[REGISTER_RETRY] temporary failure; retry %d/4 in %s", attempt+1, delay)
		registrationSleep(delay)
	}
}

func registerOnce() (*Registration, error) {
	body := map[string]interface{}{
		"name":             agentName,
		"type":             "custom",
		"connectionMethod": connectionMethod,
		"sdkVersion":       sdkVersion, // reference sends this; do not omit
		"apiKey":           apiKey,
	}
	if wallet != "" {
		body["walletAddress"] = wallet // only when set: empty fails validation
	}
	b, _ := json.Marshal(body)
	resp, err := sharedHTTP.Post(apiBase+"/peer/agents/register", "application/json", bytes.NewReader(b))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	raw, readErr := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if readErr != nil {
		return nil, readErr
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, &registrationError{status: resp.StatusCode, retryAfter: resp.Header.Get("Retry-After")}
	}
	var r Registration
	if err := json.Unmarshal(raw, &r); err != nil {
		return nil, err
	}
	return &r, nil
}

func refreshSavedToken(deviceID, refreshToken string) string {
	b, _ := json.Marshal(map[string]string{"refreshToken": refreshToken})
	resp, err := sharedHTTP.Post(fmt.Sprintf("%s/peer/token/%s/refresh", apiBase, deviceID), "application/json", bytes.NewReader(b))
	if err != nil {
		return ""
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return ""
	}
	var m map[string]string
	if json.NewDecoder(resp.Body).Decode(&m) == nil {
		if t := m["token"]; t != "" {
			return t
		}
		if t := m["jwt"]; t != "" {
			return t
		}
	}
	return ""
}

func loadState() (*Registration, error) {
	stateFileMu.Lock()
	defer stateFileMu.Unlock()
	return loadStateLocked()
}

func loadStateLocked() (*Registration, error) {
	data, err := os.ReadFile(stateFile)
	if os.IsNotExist(err) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	var r Registration
	if err := json.Unmarshal(data, &r); err != nil {
		return nil, err
	}
	if r.DeviceId == "" || (r.AgentName == "" && agentName == "") {
		return nil, errors.New("incomplete saved identity; restore state or set aside the file and use the original stable name")
	}
	return &r, nil
}

func saveState(r *Registration) error {
	stateFileMu.Lock()
	defer stateFileMu.Unlock()
	return saveStateLocked(r)
}

// Called only after redirect authority is checked, with sharedState.mu held.
// Update the existing identity; never write a stale credential snapshot.
func persistAssignedRelay(deviceID, relay string) error {
	stateFileMu.Lock()
	defer stateFileMu.Unlock()
	saved, err := loadStateLocked()
	if err != nil {
		return err
	}
	keyHash := fmt.Sprintf("%x", sha256.Sum256([]byte(apiKey)))
	if saved == nil || saved.DeviceId != deviceID || saved.APIKeyHash != keyHash || saved.RefreshToken == "" {
		return errors.New("saved identity does not match the authenticated redirect")
	}
	saved.Relay = relay
	return saveStateLocked(saved)
}

func saveStateLocked(r *Registration) error {
	data, err := json.Marshal(r)
	if err != nil {
		return err
	}
	dir := filepath.Dir(stateFile)
	if err := os.MkdirAll(dir, 0o700); err != nil {
		return err
	}
	file, err := os.CreateTemp(dir, ".peer-state-*")
	if err != nil {
		return err
	}
	defer os.Remove(file.Name())
	if _, err = file.Write(data); err != nil {
		file.Close()
		return err
	}
	if err = file.Sync(); err != nil {
		file.Close()
		return err
	}
	if err = file.Close(); err != nil {
		return err
	}
	return os.Rename(file.Name(), stateFile)
}