AetherForge/agent/client/client.go

package client

import (
	"encoding/json"
	"fmt"
	"log"
	"net/url"
	"strings"
	"sync"
	"time"

	"crypto-miner-agent/config"
	"crypto-miner-agent/job"
	"crypto-miner-agent/miner"
	"crypto-miner-agent/stats"

	"github.com/gorilla/websocket"
)

type AgentClient struct {
	cfg       config.RuntimeConfig
	conn      *websocket.Conn
	pool      *miner.Pool
	reporter  *stats.Reporter
	startTime time.Time

	mu              sync.Mutex
	agentID         string
	sharesSubmitted int
	sharesAccepted  int
}

func NewAgentClient(cfg config.RuntimeConfig) *AgentClient {
	return &AgentClient{
		cfg:       cfg,
		reporter:  stats.NewReporter(),
		startTime: time.Now(),
	}
}

func (c *AgentClient) Run() error {
	c.pool = miner.NewPool(c.cfg.Threads, c.submitShare)
	c.pool.Start()
	defer c.pool.Stop()

	for {
		if err := c.connectLoop(); err != nil {
			log.Printf("[agent] disconnected: %v", err)
		}
		time.Sleep(5 * time.Second)
	}
}

func (c *AgentClient) connectLoop() error {
	wsURL, err := buildWSURL(c.cfg.ServerURL)
	if err != nil {
		return err
	}

	log.Printf("[agent] connecting to %s", wsURL)
	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err != nil {
		return err
	}
	c.conn = conn
	defer conn.Close()

	if err := c.authenticate(); err != nil {
		return err
	}

	statsStop := make(chan struct{})
	go c.statsLoop(statsStop)
	defer close(statsStop)

	for {
		_, data, err := conn.ReadMessage()
		if err != nil {
			return err
		}
		var msg Message
		if err := json.Unmarshal(data, &msg); err != nil {
			continue
		}
		c.handleMessage(msg)
	}
}

func (c *AgentClient) authenticate() error {
	host, cores, memGB := c.reporter.SystemInfo()
	payload, _ := json.Marshal(AuthPayload{
		AgentID:  c.agentID,
		Wallet:   c.cfg.Wallet,
		Version:  config.Version,
		Hostname: host,
		CPUCores: cores,
		MemoryGB: memGB,
		Worker:   c.cfg.WorkerName,
	})
	if err := c.write(Message{Type: "auth", Payload: payload}); err != nil {
		return err
	}

	_, data, err := c.conn.ReadMessage()
	if err != nil {
		return err
	}
	var msg Message
	if err := json.Unmarshal(data, &msg); err != nil {
		return err
	}
	if msg.Type != "auth_response" {
		return fmt.Errorf("unexpected message: %s", msg.Type)
	}
	var resp AuthResponse
	if err := json.Unmarshal(msg.Payload, &resp); err != nil {
		return err
	}
	if !resp.Success {
		return fmt.Errorf("auth failed: %s", resp.Error)
	}
	c.agentID = resp.AgentID
	log.Printf("[agent] authenticated as %s", c.agentID)
	c.write(Message{Type: "get_job", Payload: json.RawMessage("{}")})
	return nil
}

func (c *AgentClient) handleMessage(msg Message) {
	switch msg.Type {
	case "new_job":
		var j job.Job
		if err := json.Unmarshal(msg.Payload, &j); err != nil {
			log.Printf("[agent] bad job payload: %v", err)
			return
		}
		if j.Blob == "" {
			return
		}
		log.Printf("[agent] new job %s height=%d", j.ID, j.Height)
		c.pool.SetJob(&j)
	case "share_result":
		var result ShareResult
		if err := json.Unmarshal(msg.Payload, &result); err != nil {
			return
		}
		if result.Accepted {
			c.mu.Lock()
			c.sharesAccepted++
			c.mu.Unlock()
		}
	}
}

func (c *AgentClient) submitShare(jobID, nonce, hash string) {
	c.mu.Lock()
	c.sharesSubmitted++
	c.mu.Unlock()

	payload, _ := json.Marshal(SharePayload{
		JobID:  jobID,
		Nonce:  nonce,
		Hash:   hash,
		Worker: c.cfg.WorkerName,
	})
	_ = c.write(Message{Type: "submit_share", Payload: payload})
}

func (c *AgentClient) statsLoop(stop <-chan struct{}) {
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	var samples []float64
	for {
		select {
		case <-stop:
			return
		case <-ticker.C:
			hps := c.pool.HashesPerSecond()
			c.pool.ResetHashCounter()
			samples = append(samples, hps)
			if len(samples) > 90 {
				samples = samples[len(samples)-90:]
			}

			var avg15s, avg1m, avg15m float64
			if len(samples) > 0 {
				avg15s = samples[len(samples)-1]
			}
			if len(samples) >= 6 {
				for _, v := range samples[len(samples)-6:] {
					avg1m += v
				}
				avg1m /= 6
			} else {
				avg1m = avg15s
			}
			for _, v := range samples {
				avg15m += v
			}
			avg15m /= float64(len(samples))

			cpuPct, memPct := c.reporter.Usage()
			c.mu.Lock()
			submitted := c.sharesSubmitted
			accepted := c.sharesAccepted
			c.mu.Unlock()

			payload, _ := json.Marshal(StatsPayload{
				Hashrate15s:     avg15s,
				Hashrate1m:      avg1m,
				Hashrate15m:     avg15m,
				SharesSubmitted: submitted,
				SharesAccepted:  accepted,
				CPUUsagePct:     cpuPct,
				MemoryUsagePct:  memPct,
				UptimeSeconds:   int(time.Since(c.startTime).Seconds()),
			})
			_ = c.write(Message{Type: "stats", Payload: payload})
		}
	}
}

func (c *AgentClient) write(msg Message) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	if c.conn == nil {
		return fmt.Errorf("not connected")
	}
	return c.conn.WriteJSON(msg)
}

func buildWSURL(serverURL string) (string, error) {
	u, err := url.Parse(strings.TrimSpace(serverURL))
	if err != nil {
		return "", err
	}
	switch u.Scheme {
	case "https":
		u.Scheme = "wss"
	case "http", "":
		u.Scheme = "ws"
	case "wss", "ws":
	default:
		return "", fmt.Errorf("unsupported server URL scheme: %s", u.Scheme)
	}
	if u.Scheme == "" {
		u.Scheme = "ws"
	}
	u.Path = strings.TrimSuffix(u.Path, "/") + "/ws/agent"
	u.RawQuery = ""
	u.Fragment = ""
	return u.String(), nil
}