Mining/pkg/node/worker.go

package node

import (
	"encoding/base64"
	"path/filepath"
	"time"

	"forge.lthn.ai/Snider/Mining/pkg/logging"
	"github.com/adrg/xdg"
)

// worker.SetMinerManager(miningManager)
// miner, err := worker.minerManager.StartMiner("xmrig", config)
// worker.minerManager.StopMiner("xmrig")
// for _, miner := range worker.minerManager.ListMiners() { ... }
type MinerManager interface {
	StartMiner(minerType string, config interface{}) (MinerInstance, error)
	StopMiner(name string) error
	ListMiners() []MinerInstance
	GetMiner(name string) (MinerInstance, error)
}

// miner := worker.minerManager.GetMiner("xmrig")
// name := miner.GetName()           // "xmrig"
// stats, _ := miner.GetStats()      // current hashrate, shares
// logs := miner.GetConsoleHistory(50)
type MinerInstance interface {
	GetName() string
	GetType() string
	GetStats() (interface{}, error)
	GetConsoleHistory(lines int) []string
}

// profile, err := worker.profileManager.GetProfile("pool-main")
// if err != nil { return nil, err }
// worker.profileManager.SaveProfile(profile)
type ProfileManager interface {
	GetProfile(id string) (interface{}, error)
	SaveProfile(profile interface{}) error
}

// worker := node.NewWorker(nodeManager, transport)
// worker.SetMinerManager(miningManager)
// worker.SetProfileManager(profileManager)
// worker.RegisterWithTransport() // begin processing incoming messages
type Worker struct {
	nodeManager    *NodeManager
	transport      *Transport
	minerManager   MinerManager
	profileManager ProfileManager
	startTime      time.Time
}

// worker := node.NewWorker(nodeManager, transport)
// worker.SetMinerManager(miningManager)
// worker.RegisterWithTransport()
func NewWorker(nodeManager *NodeManager, transport *Transport) *Worker {
	return &Worker{
		nodeManager: nodeManager,
		transport:   transport,
		startTime:   time.Now(),
	}
}

// worker.SetMinerManager(miningManager)
func (worker *Worker) SetMinerManager(manager MinerManager) {
	worker.minerManager = manager
}

// worker.SetProfileManager(profileManager)
func (worker *Worker) SetProfileManager(manager ProfileManager) {
	worker.profileManager = manager
}

// worker.HandleMessage(connection, message)
// worker.RegisterWithTransport() // registers HandleMessage as the transport handler
func (worker *Worker) HandleMessage(connection *PeerConnection, message *Message) {
	var response *Message
	var err error

	switch message.Type {
	case MsgPing:
		response, err = worker.handlePing(message)
	case MsgGetStats:
		response, err = worker.handleGetStats(message)
	case MsgStartMiner:
		response, err = worker.handleStartMiner(message)
	case MsgStopMiner:
		response, err = worker.handleStopMiner(message)
	case MsgGetLogs:
		response, err = worker.handleGetLogs(message)
	case MsgDeploy:
		response, err = worker.handleDeploy(connection, message)
	default:
		// Unknown message type - ignore or send error
		return
	}

	if err != nil {
		// Send error response
		identity := worker.nodeManager.GetIdentity()
		if identity != nil {
			errorMessage, _ := NewErrorMessage(
				identity.ID,
				message.From,
				ErrCodeOperationFailed,
				err.Error(),
				message.ID,
			)
			connection.Send(errorMessage)
		}
		return
	}

	if response != nil {
		logging.Debug("sending response", logging.Fields{"type": response.Type, "to": message.From})
		if err := connection.Send(response); err != nil {
			logging.Error("failed to send response", logging.Fields{"error": err})
		} else {
			logging.Debug("response sent successfully")
		}
	}
}

// response, err := worker.handlePing(message)
func (worker *Worker) handlePing(message *Message) (*Message, error) {
	var ping PingPayload
	if err := message.ParsePayload(&ping); err != nil {
		return nil, &ProtocolError{Code: ErrCodeInvalidMessage, Message: "invalid ping payload: " + err.Error()}
	}

	pong := PongPayload{
		SentAt:     ping.SentAt,
		ReceivedAt: time.Now().UnixMilli(),
	}

	return message.Reply(MsgPong, pong)
}

// response, err := worker.handleGetStats(message)
func (worker *Worker) handleGetStats(message *Message) (*Message, error) {
	identity := worker.nodeManager.GetIdentity()
	if identity == nil {
		return nil, &ProtocolError{Code: ErrCodeUnauthorized, Message: "node identity not initialized"}
	}

	stats := StatsPayload{
		NodeID:   identity.ID,
		NodeName: identity.Name,
		Miners:   []MinerStatsItem{},
		Uptime:   int64(time.Since(worker.startTime).Seconds()),
	}

	if worker.minerManager != nil {
		miners := worker.minerManager.ListMiners()
		for _, miner := range miners {
			minerStats, err := miner.GetStats()
			if err != nil {
				continue
			}

			item := convertMinerStats(miner, minerStats)
			stats.Miners = append(stats.Miners, item)
		}
	}

	return message.Reply(MsgStats, stats)
}

// item := convertMinerStats(miner, rawStats)
func convertMinerStats(miner MinerInstance, rawStats interface{}) MinerStatsItem {
	item := MinerStatsItem{
		Name: miner.GetName(),
		Type: miner.GetType(),
	}

	if statsMap, ok := rawStats.(map[string]interface{}); ok {
		if hashrate, ok := statsMap["hashrate"].(float64); ok {
			item.Hashrate = hashrate
		}
		if shares, ok := statsMap["shares"].(int); ok {
			item.Shares = shares
		}
		if rejected, ok := statsMap["rejected"].(int); ok {
			item.Rejected = rejected
		}
		if uptime, ok := statsMap["uptime"].(int); ok {
			item.Uptime = uptime
		}
		if pool, ok := statsMap["pool"].(string); ok {
			item.Pool = pool
		}
		if algorithm, ok := statsMap["algorithm"].(string); ok {
			item.Algorithm = algorithm
		}
	}

	return item
}

// response, err := worker.handleStartMiner(message)
func (worker *Worker) handleStartMiner(message *Message) (*Message, error) {
	if worker.minerManager == nil {
		return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "miner manager not configured"}
	}

	var payload StartMinerPayload
	if err := message.ParsePayload(&payload); err != nil {
		return nil, &ProtocolError{Code: ErrCodeInvalidMessage, Message: "invalid start miner payload: " + err.Error()}
	}

	// Validate miner type is provided
	if payload.MinerType == "" {
		return nil, &ProtocolError{Code: ErrCodeInvalidMessage, Message: "miner type is required"}
	}

	var config interface{}
	if payload.Config != nil {
		config = payload.Config
	} else if worker.profileManager != nil {
		profile, err := worker.profileManager.GetProfile(payload.ProfileID)
		if err != nil {
			return nil, &ProtocolError{Code: ErrCodeNotFound, Message: "profile not found: " + payload.ProfileID}
		}
		config = profile
	} else {
		return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "no config provided and no profile manager configured"}
	}

	miner, err := worker.minerManager.StartMiner(payload.MinerType, config)
	if err != nil {
		acknowledgement := MinerAckPayload{
			Success: false,
			Error:   err.Error(),
		}
		return message.Reply(MsgMinerAck, acknowledgement)
	}

	acknowledgement := MinerAckPayload{
		Success:   true,
		MinerName: miner.GetName(),
	}
	return message.Reply(MsgMinerAck, acknowledgement)
}

// response, err := worker.handleStopMiner(message)
func (worker *Worker) handleStopMiner(message *Message) (*Message, error) {
	if worker.minerManager == nil {
		return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "miner manager not configured"}
	}

	var payload StopMinerPayload
	if err := message.ParsePayload(&payload); err != nil {
		return nil, &ProtocolError{Code: ErrCodeInvalidMessage, Message: "invalid stop miner payload: " + err.Error()}
	}

	err := worker.minerManager.StopMiner(payload.MinerName)
	acknowledgement := MinerAckPayload{
		Success:   err == nil,
		MinerName: payload.MinerName,
	}
	if err != nil {
		acknowledgement.Error = err.Error()
	}

	return message.Reply(MsgMinerAck, acknowledgement)
}

// response, err := worker.handleGetLogs(message)
func (worker *Worker) handleGetLogs(message *Message) (*Message, error) {
	if worker.minerManager == nil {
		return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "miner manager not configured"}
	}

	var payload GetLogsPayload
	if err := message.ParsePayload(&payload); err != nil {
		return nil, &ProtocolError{Code: ErrCodeInvalidMessage, Message: "invalid get logs payload: " + err.Error()}
	}

	// Validate and limit the Lines parameter to prevent resource exhaustion
	const maxLogLines = 10000
	if payload.Lines <= 0 || payload.Lines > maxLogLines {
		payload.Lines = maxLogLines
	}

	miner, err := worker.minerManager.GetMiner(payload.MinerName)
	if err != nil {
		return nil, &ProtocolError{Code: ErrCodeNotFound, Message: "miner not found: " + payload.MinerName}
	}

	lines := miner.GetConsoleHistory(payload.Lines)

	logs := LogsPayload{
		MinerName: payload.MinerName,
		Lines:     lines,
		HasMore:   len(lines) >= payload.Lines,
	}

	return message.Reply(MsgLogs, logs)
}

// response, err := worker.handleDeploy(connection, message)
func (worker *Worker) handleDeploy(connection *PeerConnection, message *Message) (*Message, error) {
	var payload DeployPayload
	if err := message.ParsePayload(&payload); err != nil {
		return nil, &ProtocolError{Code: ErrCodeInvalidMessage, Message: "invalid deploy payload: " + err.Error()}
	}

	// Reconstruct Bundle object from payload
	bundle := &Bundle{
		Type:     BundleType(payload.BundleType),
		Name:     payload.Name,
		Data:     payload.Data,
		Checksum: payload.Checksum,
	}

	// Use shared secret as password (base64 encoded)
	password := ""
	if connection != nil && len(connection.SharedSecret) > 0 {
		password = base64.StdEncoding.EncodeToString(connection.SharedSecret)
	}

	switch bundle.Type {
	case BundleProfile:
		if worker.profileManager == nil {
			return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "profile manager not configured"}
		}

		// Decrypt and extract profile data
		profileData, err := ExtractProfileBundle(bundle, password)
		if err != nil {
			return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to extract profile bundle: " + err.Error()}
		}

		// Unmarshal into interface{} to pass to ProfileManager
		var profile interface{}
		if err := UnmarshalJSON(profileData, &profile); err != nil {
			return nil, &ProtocolError{Code: ErrCodeInvalidMessage, Message: "invalid profile data JSON: " + err.Error()}
		}

		if err := worker.profileManager.SaveProfile(profile); err != nil {
			acknowledgement := DeployAckPayload{
				Success: false,
				Name:    payload.Name,
				Error:   err.Error(),
			}
			return message.Reply(MsgDeployAck, acknowledgement)
		}

		acknowledgement := DeployAckPayload{
			Success: true,
			Name:    payload.Name,
		}
		return message.Reply(MsgDeployAck, acknowledgement)

	case BundleMiner, BundleFull:
		// Determine installation directory
		// We use xdg.DataHome/lethean-desktop/miners/<bundle_name>
		minersDir := filepath.Join(xdg.DataHome, "lethean-desktop", "miners")
		installDir := filepath.Join(minersDir, payload.Name)

		logging.Info("deploying miner bundle", logging.Fields{
			"name": payload.Name,
			"path": installDir,
			"type": payload.BundleType,
		})

		// Extract miner bundle
		minerPath, profileData, err := ExtractMinerBundle(bundle, password, installDir)
		if err != nil {
			return nil, &ProtocolError{Code: ErrCodeOperationFailed, Message: "failed to extract miner bundle: " + err.Error()}
		}

		// If the bundle contained a profile config, save it
		if len(profileData) > 0 && worker.profileManager != nil {
			var profile interface{}
			if err := UnmarshalJSON(profileData, &profile); err != nil {
				logging.Warn("failed to parse profile from miner bundle", logging.Fields{"error": err})
			} else {
				if err := worker.profileManager.SaveProfile(profile); err != nil {
					logging.Warn("failed to save profile from miner bundle", logging.Fields{"error": err})
				}
			}
		}

		// Success response
		acknowledgement := DeployAckPayload{
			Success: true,
			Name:    payload.Name,
		}

		// Log the installation
		logging.Info("miner bundle installed successfully", logging.Fields{
			"name":       payload.Name,
			"miner_path": minerPath,
		})

		return message.Reply(MsgDeployAck, acknowledgement)

	default:
		return nil, &ProtocolError{Code: ErrCodeInvalidMessage, Message: "unknown bundle type: " + payload.BundleType}
	}
}

// worker.RegisterWithTransport() // call once after setup to start handling messages
func (worker *Worker) RegisterWithTransport() {
	worker.transport.OnMessage(worker.HandleMessage)
}