Mining/pkg/node/peer.go

package node

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/Snider/Poindexter"
	"github.com/adrg/xdg"
)

// Peer represents a known remote node.
type Peer struct {
	ID        string    `json:"id"`
	Name      string    `json:"name"`
	PublicKey string    `json:"publicKey"`
	Address   string    `json:"address"` // host:port for WebSocket connection
	Role      NodeRole  `json:"role"`
	AddedAt   time.Time `json:"addedAt"`
	LastSeen  time.Time `json:"lastSeen"`

	// Poindexter metrics (updated dynamically)
	PingMS float64 `json:"pingMs"` // Latency in milliseconds
	Hops   int     `json:"hops"`   // Network hop count
	GeoKM  float64 `json:"geoKm"`  // Geographic distance in kilometers
	Score  float64 `json:"score"`  // Reliability score 0-100

	// Connection state (not persisted)
	Connected bool `json:"-"`
}

// saveDebounceInterval is the minimum time between disk writes.
const saveDebounceInterval = 5 * time.Second

// PeerRegistry manages known peers with KD-tree based selection.
type PeerRegistry struct {
	peers  map[string]*Peer
	kdTree *poindexter.KDTree[string] // KD-tree with peer ID as payload
	path   string
	mu     sync.RWMutex

	// Debounce disk writes
	dirty        bool          // Whether there are unsaved changes
	saveTimer    *time.Timer   // Timer for debounced save
	saveMu       sync.Mutex    // Protects dirty and saveTimer
	stopChan     chan struct{} // Signal to stop background save
	saveStopOnce sync.Once     // Ensure stopChan is closed only once
}

// Dimension weights for peer selection
// Lower ping, hops, geo are better; higher score is better
var (
	pingWeight  = 1.0
	hopsWeight  = 0.7
	geoWeight   = 0.2
	scoreWeight = 1.2
)

// NewPeerRegistry creates a new PeerRegistry, loading existing peers if available.
func NewPeerRegistry() (*PeerRegistry, error) {
	peersPath, err := xdg.ConfigFile("lethean-desktop/peers.json")
	if err != nil {
		return nil, fmt.Errorf("failed to get peers path: %w", err)
	}

	return NewPeerRegistryWithPath(peersPath)
}

// NewPeerRegistryWithPath creates a new PeerRegistry with a custom path.
// This is primarily useful for testing to avoid xdg path caching issues.
func NewPeerRegistryWithPath(peersPath string) (*PeerRegistry, error) {
	pr := &PeerRegistry{
		peers:    make(map[string]*Peer),
		path:     peersPath,
		stopChan: make(chan struct{}),
	}

	// Try to load existing peers
	if err := pr.load(); err != nil {
		// No existing peers, that's ok
		pr.rebuildKDTree()
		return pr, nil
	}

	pr.rebuildKDTree()
	return pr, nil
}

// AddPeer adds a new peer to the registry.
// Note: Persistence is debounced (writes batched every 5s). Call Close() to ensure
// all changes are flushed to disk before shutdown.
func (r *PeerRegistry) AddPeer(peer *Peer) error {
	r.mu.Lock()

	if peer.ID == "" {
		r.mu.Unlock()
		return fmt.Errorf("peer ID is required")
	}

	if _, exists := r.peers[peer.ID]; exists {
		r.mu.Unlock()
		return fmt.Errorf("peer %s already exists", peer.ID)
	}

	// Set defaults
	if peer.AddedAt.IsZero() {
		peer.AddedAt = time.Now()
	}
	if peer.Score == 0 {
		peer.Score = 50 // Default neutral score
	}

	r.peers[peer.ID] = peer
	r.rebuildKDTree()
	r.mu.Unlock()

	return r.save()
}

// UpdatePeer updates an existing peer's information.
// Note: Persistence is debounced. Call Close() to flush before shutdown.
func (r *PeerRegistry) UpdatePeer(peer *Peer) error {
	r.mu.Lock()

	if _, exists := r.peers[peer.ID]; !exists {
		r.mu.Unlock()
		return fmt.Errorf("peer %s not found", peer.ID)
	}

	r.peers[peer.ID] = peer
	r.rebuildKDTree()
	r.mu.Unlock()

	return r.save()
}

// RemovePeer removes a peer from the registry.
// Note: Persistence is debounced. Call Close() to flush before shutdown.
func (r *PeerRegistry) RemovePeer(id string) error {
	r.mu.Lock()

	if _, exists := r.peers[id]; !exists {
		r.mu.Unlock()
		return fmt.Errorf("peer %s not found", id)
	}

	delete(r.peers, id)
	r.rebuildKDTree()
	r.mu.Unlock()

	return r.save()
}

// GetPeer returns a peer by ID.
func (r *PeerRegistry) GetPeer(id string) *Peer {
	r.mu.RLock()
	defer r.mu.RUnlock()

	peer, exists := r.peers[id]
	if !exists {
		return nil
	}

	// Return a copy
	peerCopy := *peer
	return &peerCopy
}

// ListPeers returns all registered peers.
func (r *PeerRegistry) ListPeers() []*Peer {
	r.mu.RLock()
	defer r.mu.RUnlock()

	peers := make([]*Peer, 0, len(r.peers))
	for _, peer := range r.peers {
		peerCopy := *peer
		peers = append(peers, &peerCopy)
	}
	return peers
}

// UpdateMetrics updates a peer's performance metrics.
// Note: Persistence is debounced. Call Close() to flush before shutdown.
func (r *PeerRegistry) UpdateMetrics(id string, pingMS, geoKM float64, hops int) error {
	r.mu.Lock()

	peer, exists := r.peers[id]
	if !exists {
		r.mu.Unlock()
		return fmt.Errorf("peer %s not found", id)
	}

	peer.PingMS = pingMS
	peer.GeoKM = geoKM
	peer.Hops = hops
	peer.LastSeen = time.Now()

	r.rebuildKDTree()
	r.mu.Unlock()

	return r.save()
}

// UpdateScore updates a peer's reliability score.
// Note: Persistence is debounced. Call Close() to flush before shutdown.
func (r *PeerRegistry) UpdateScore(id string, score float64) error {
	r.mu.Lock()

	peer, exists := r.peers[id]
	if !exists {
		r.mu.Unlock()
		return fmt.Errorf("peer %s not found", id)
	}

	// Clamp score to 0-100
	if score < 0 {
		score = 0
	} else if score > 100 {
		score = 100
	}

	peer.Score = score
	r.rebuildKDTree()
	r.mu.Unlock()

	return r.save()
}

// SetConnected updates a peer's connection state.
func (r *PeerRegistry) SetConnected(id string, connected bool) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if peer, exists := r.peers[id]; exists {
		peer.Connected = connected
		if connected {
			peer.LastSeen = time.Now()
		}
	}
}

// SelectOptimalPeer returns the best peer based on multi-factor optimization.
// Uses Poindexter KD-tree to find the peer closest to ideal metrics.
func (r *PeerRegistry) SelectOptimalPeer() *Peer {
	r.mu.RLock()
	defer r.mu.RUnlock()

	if r.kdTree == nil || len(r.peers) == 0 {
		return nil
	}

	// Target: ideal peer (0 ping, 0 hops, 0 geo, 100 score)
	// Score is inverted (100 - score) so lower is better in the tree
	target := []float64{0, 0, 0, 0}

	result, _, found := r.kdTree.Nearest(target)
	if !found {
		return nil
	}

	peer, exists := r.peers[result.Value]
	if !exists {
		return nil
	}

	peerCopy := *peer
	return &peerCopy
}

// SelectNearestPeers returns the n best peers based on multi-factor optimization.
func (r *PeerRegistry) SelectNearestPeers(n int) []*Peer {
	r.mu.RLock()
	defer r.mu.RUnlock()

	if r.kdTree == nil || len(r.peers) == 0 {
		return nil
	}

	// Target: ideal peer
	target := []float64{0, 0, 0, 0}

	results, _ := r.kdTree.KNearest(target, n)

	peers := make([]*Peer, 0, len(results))
	for _, result := range results {
		if peer, exists := r.peers[result.Value]; exists {
			peerCopy := *peer
			peers = append(peers, &peerCopy)
		}
	}

	return peers
}

// GetConnectedPeers returns all currently connected peers.
func (r *PeerRegistry) GetConnectedPeers() []*Peer {
	r.mu.RLock()
	defer r.mu.RUnlock()

	peers := make([]*Peer, 0)
	for _, peer := range r.peers {
		if peer.Connected {
			peerCopy := *peer
			peers = append(peers, &peerCopy)
		}
	}
	return peers
}

// Count returns the number of registered peers.
func (r *PeerRegistry) Count() int {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return len(r.peers)
}

// rebuildKDTree rebuilds the KD-tree from current peers.
// Must be called with lock held.
func (r *PeerRegistry) rebuildKDTree() {
	if len(r.peers) == 0 {
		r.kdTree = nil
		return
	}

	points := make([]poindexter.KDPoint[string], 0, len(r.peers))
	for _, peer := range r.peers {
		// Build 4D point with weighted, normalized values
		// Invert score so that higher score = lower value (better)
		point := poindexter.KDPoint[string]{
			ID: peer.ID,
			Coords: []float64{
				peer.PingMS * pingWeight,
				float64(peer.Hops) * hopsWeight,
				peer.GeoKM * geoWeight,
				(100 - peer.Score) * scoreWeight, // Invert score
			},
			Value: peer.ID,
		}
		points = append(points, point)
	}

	// Build KD-tree with Euclidean distance
	tree, err := poindexter.NewKDTree(points, poindexter.WithMetric(poindexter.EuclideanDistance{}))
	if err != nil {
		// Log error but continue - worst case we don't have optimal selection
		return
	}

	r.kdTree = tree
}

// scheduleSave schedules a debounced save operation.
// Multiple calls within saveDebounceInterval will be coalesced into a single save.
// Must NOT be called with r.mu held.
func (r *PeerRegistry) scheduleSave() {
	r.saveMu.Lock()
	defer r.saveMu.Unlock()

	r.dirty = true

	// If timer already running, let it handle the save
	if r.saveTimer != nil {
		return
	}

	// Start a new timer
	r.saveTimer = time.AfterFunc(saveDebounceInterval, func() {
		r.saveMu.Lock()
		r.saveTimer = nil
		shouldSave := r.dirty
		r.dirty = false
		r.saveMu.Unlock()

		if shouldSave {
			r.mu.RLock()
			err := r.saveNow()
			r.mu.RUnlock()
			if err != nil {
				// Log error but continue - best effort persistence
				log.Printf("Warning: failed to save peer registry: %v", err)
			}
		}
	})
}

// saveNow persists peers to disk immediately.
// Must be called with r.mu held (at least RLock).
func (r *PeerRegistry) saveNow() error {
	// Ensure directory exists
	dir := filepath.Dir(r.path)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return fmt.Errorf("failed to create peers directory: %w", err)
	}

	// Convert to slice for JSON
	peers := make([]*Peer, 0, len(r.peers))
	for _, peer := range r.peers {
		peers = append(peers, peer)
	}

	data, err := json.MarshalIndent(peers, "", "  ")
	if err != nil {
		return fmt.Errorf("failed to marshal peers: %w", err)
	}

	// Use atomic write pattern: write to temp file, then rename
	tmpPath := r.path + ".tmp"
	if err := os.WriteFile(tmpPath, data, 0644); err != nil {
		return fmt.Errorf("failed to write peers temp file: %w", err)
	}

	if err := os.Rename(tmpPath, r.path); err != nil {
		os.Remove(tmpPath) // Clean up temp file
		return fmt.Errorf("failed to rename peers file: %w", err)
	}

	return nil
}

// Close flushes any pending changes and releases resources.
func (r *PeerRegistry) Close() error {
	r.saveStopOnce.Do(func() {
		close(r.stopChan)
	})

	// Cancel pending timer and save immediately if dirty
	r.saveMu.Lock()
	if r.saveTimer != nil {
		r.saveTimer.Stop()
		r.saveTimer = nil
	}
	shouldSave := r.dirty
	r.dirty = false
	r.saveMu.Unlock()

	if shouldSave {
		r.mu.RLock()
		err := r.saveNow()
		r.mu.RUnlock()
		return err
	}

	return nil
}

// save is a helper that schedules a debounced save.
// Kept for backward compatibility but now debounces writes.
// Must NOT be called with r.mu held.
func (r *PeerRegistry) save() error {
	r.scheduleSave()
	return nil // Errors will be logged asynchronously
}

// load reads peers from disk.
func (r *PeerRegistry) load() error {
	data, err := os.ReadFile(r.path)
	if err != nil {
		return fmt.Errorf("failed to read peers: %w", err)
	}

	var peers []*Peer
	if err := json.Unmarshal(data, &peers); err != nil {
		return fmt.Errorf("failed to unmarshal peers: %w", err)
	}

	r.peers = make(map[string]*Peer)
	for _, peer := range peers {
		r.peers[peer.ID] = peer
	}

	return nil
}