Mining/pkg/node/controller.go

package node

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"

	"forge.lthn.ai/Snider/Mining/pkg/logging"
)

// ctrl := node.NewController(nodeManager, peerRegistry, transport)
// rtt, _ := ctrl.PingPeer("abc123def456")
// stats, _ := ctrl.GetRemoteStats("abc123def456")
type Controller struct {
	node      *NodeManager
	peers     *PeerRegistry
	transport *Transport
	mutex     sync.RWMutex

	// Pending requests awaiting responses
	pending map[string]chan *Message // message ID -> response channel
}

// ctrl := node.NewController(nodeManager, peerRegistry, transport)
// rtt, err := ctrl.PingPeer("abc123def456")
func NewController(node *NodeManager, peers *PeerRegistry, transport *Transport) *Controller {
	controller := &Controller{
		node:      node,
		peers:     peers,
		transport: transport,
		pending:   make(map[string]chan *Message),
	}

	// Register message handler for responses
	transport.OnMessage(controller.handleResponse)

	return controller
}

// transport.OnMessage(ctrl.handleResponse) // registered in NewController
func (ctrl *Controller) handleResponse(conn *PeerConnection, msg *Message) {
	if msg.ReplyTo == "" {
		return // Not a response, let worker handle it
	}

	ctrl.mutex.Lock()
	channel, exists := ctrl.pending[msg.ReplyTo]
	if exists {
		delete(ctrl.pending, msg.ReplyTo)
	}
	ctrl.mutex.Unlock()

	if exists && channel != nil {
		select {
		case channel <- msg:
		default:
			// Channel full or closed
		}
	}
}

// resp, err := ctrl.sendRequest("abc123def456", msg, 10*time.Second)
// if err != nil { return nil, err }
func (ctrl *Controller) sendRequest(peerID string, msg *Message, timeout time.Duration) (*Message, error) {
	actualPeerID := peerID

	// Auto-connect if not already connected
	if ctrl.transport.GetConnection(peerID) == nil {
		peer := ctrl.peers.GetPeer(peerID)
		if peer == nil {
			return nil, fmt.Errorf("peer not found: %s", peerID)
		}
		conn, err := ctrl.transport.Connect(peer)
		if err != nil {
			return nil, fmt.Errorf("failed to connect to peer: %w", err)
		}
		// Use the real peer ID after handshake (it may have changed)
		actualPeerID = conn.Peer.ID
		// Update the message destination
		msg.To = actualPeerID
	}

	// Create response channel
	responseChannel := make(chan *Message, 1)

	ctrl.mutex.Lock()
	ctrl.pending[msg.ID] = responseChannel
	ctrl.mutex.Unlock()

	// Clean up on exit - ensure channel is closed and removed from map
	defer func() {
		ctrl.mutex.Lock()
		delete(ctrl.pending, msg.ID)
		ctrl.mutex.Unlock()
		close(responseChannel) // Close channel to allow garbage collection
	}()

	// Send the message
	if err := ctrl.transport.Send(actualPeerID, msg); err != nil {
		return nil, fmt.Errorf("failed to send message: %w", err)
	}

	// Wait for response
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	select {
	case resp := <-responseChannel:
		return resp, nil
	case <-ctx.Done():
		return nil, fmt.Errorf("request timeout")
	}
}

// stats, err := ctrl.GetRemoteStats("abc123def456")
// if err == nil { log("miners:", len(stats.Miners)) }
func (ctrl *Controller) GetRemoteStats(peerID string) (*StatsPayload, error) {
	identity := ctrl.node.GetIdentity()
	if identity == nil {
		return nil, fmt.Errorf("node identity not initialized")
	}

	msg, err := NewMessage(MsgGetStats, identity.ID, peerID, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create message: %w", err)
	}

	resp, err := ctrl.sendRequest(peerID, msg, 10*time.Second)
	if err != nil {
		return nil, err
	}

	var stats StatsPayload
	if err := ParseResponse(resp, MsgStats, &stats); err != nil {
		return nil, err
	}

	return &stats, nil
}

// err := ctrl.StartRemoteMiner("abc123def456", "xmrig", "profile-1", nil)
// err := ctrl.StartRemoteMiner("abc123def456", "xmrig", "", json.RawMessage(`{"pool":"stratum+tcp://pool.lthn.io:3333"}`))
func (ctrl *Controller) StartRemoteMiner(peerID, minerType, profileID string, configOverride json.RawMessage) error {
	identity := ctrl.node.GetIdentity()
	if identity == nil {
		return fmt.Errorf("node identity not initialized")
	}

	if minerType == "" {
		return fmt.Errorf("miner type is required")
	}

	payload := StartMinerPayload{
		MinerType: minerType,
		ProfileID: profileID,
		Config:    configOverride,
	}

	msg, err := NewMessage(MsgStartMiner, identity.ID, peerID, payload)
	if err != nil {
		return fmt.Errorf("failed to create message: %w", err)
	}

	resp, err := ctrl.sendRequest(peerID, msg, 30*time.Second)
	if err != nil {
		return err
	}

	var ack MinerAckPayload
	if err := ParseResponse(resp, MsgMinerAck, &ack); err != nil {
		return err
	}

	if !ack.Success {
		return fmt.Errorf("miner start failed: %s", ack.Error)
	}

	return nil
}

// err := ctrl.StopRemoteMiner("abc123def456", "xmrig")
// if err != nil { log("stop failed:", err) }
func (ctrl *Controller) StopRemoteMiner(peerID, minerName string) error {
	identity := ctrl.node.GetIdentity()
	if identity == nil {
		return fmt.Errorf("node identity not initialized")
	}

	payload := StopMinerPayload{
		MinerName: minerName,
	}

	msg, err := NewMessage(MsgStopMiner, identity.ID, peerID, payload)
	if err != nil {
		return fmt.Errorf("failed to create message: %w", err)
	}

	resp, err := ctrl.sendRequest(peerID, msg, 30*time.Second)
	if err != nil {
		return err
	}

	var ack MinerAckPayload
	if err := ParseResponse(resp, MsgMinerAck, &ack); err != nil {
		return err
	}

	if !ack.Success {
		return fmt.Errorf("miner stop failed: %s", ack.Error)
	}

	return nil
}

// lines, err := ctrl.GetRemoteLogs("abc123def456", "xmrig", 100)
// if err == nil { log("last log:", lines[len(lines)-1]) }
func (ctrl *Controller) GetRemoteLogs(peerID, minerName string, lines int) ([]string, error) {
	identity := ctrl.node.GetIdentity()
	if identity == nil {
		return nil, fmt.Errorf("node identity not initialized")
	}

	payload := GetLogsPayload{
		MinerName: minerName,
		Lines:     lines,
	}

	msg, err := NewMessage(MsgGetLogs, identity.ID, peerID, payload)
	if err != nil {
		return nil, fmt.Errorf("failed to create message: %w", err)
	}

	resp, err := ctrl.sendRequest(peerID, msg, 10*time.Second)
	if err != nil {
		return nil, err
	}

	var logs LogsPayload
	if err := ParseResponse(resp, MsgLogs, &logs); err != nil {
		return nil, err
	}

	return logs.Lines, nil
}

// allStats := ctrl.GetAllStats()
// for peerID, stats := range allStats { log(peerID, stats.Miners) }
func (ctrl *Controller) GetAllStats() map[string]*StatsPayload {
	peers := ctrl.peers.GetConnectedPeers()
	results := make(map[string]*StatsPayload)
	var resultsMutex sync.Mutex
	var waitGroup sync.WaitGroup

	for _, peer := range peers {
		waitGroup.Add(1)
		go func(connectedPeer *Peer) {
			defer waitGroup.Done()
			stats, err := ctrl.GetRemoteStats(connectedPeer.ID)
			if err != nil {
				logging.Debug("failed to get stats from peer", logging.Fields{
					"peer_id": connectedPeer.ID,
					"peer":    connectedPeer.Name,
					"error":   err.Error(),
				})
				return // Skip failed peers
			}
			resultsMutex.Lock()
			results[connectedPeer.ID] = stats
			resultsMutex.Unlock()
		}(peer)
	}

	waitGroup.Wait()
	return results
}

// rtt, err := ctrl.PingPeer("abc123def456")
// if err == nil { log("latency:", rtt, "ms") }
func (ctrl *Controller) PingPeer(peerID string) (float64, error) {
	identity := ctrl.node.GetIdentity()
	if identity == nil {
		return 0, fmt.Errorf("node identity not initialized")
	}
	sentAt := time.Now()

	payload := PingPayload{
		SentAt: sentAt.UnixMilli(),
	}

	msg, err := NewMessage(MsgPing, identity.ID, peerID, payload)
	if err != nil {
		return 0, fmt.Errorf("failed to create message: %w", err)
	}

	resp, err := ctrl.sendRequest(peerID, msg, 5*time.Second)
	if err != nil {
		return 0, err
	}

	if err := ValidateResponse(resp, MsgPong); err != nil {
		return 0, err
	}

	roundTripTime := time.Since(sentAt).Seconds() * 1000 // milliseconds

	// Update peer metrics
	peer := ctrl.peers.GetPeer(peerID)
	if peer != nil {
		ctrl.peers.UpdateMetrics(peerID, roundTripTime, peer.GeoKM, peer.Hops)
	}

	return roundTripTime, nil
}

// if err := ctrl.ConnectToPeer("abc123def456"); err != nil { log(err) }
func (ctrl *Controller) ConnectToPeer(peerID string) error {
	peer := ctrl.peers.GetPeer(peerID)
	if peer == nil {
		return fmt.Errorf("peer not found: %s", peerID)
	}

	_, err := ctrl.transport.Connect(peer)
	return err
}

// if err := ctrl.DisconnectFromPeer("abc123def456"); err != nil { log(err) }
func (ctrl *Controller) DisconnectFromPeer(peerID string) error {
	conn := ctrl.transport.GetConnection(peerID)
	if conn == nil {
		return fmt.Errorf("peer not connected: %s", peerID)
	}

	return conn.Close()
}