AetherForge/agent/client/gpu_miner.go

package client

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

	"crypto-miner-agent/config"
)


// GPUVendor identifies the discrete GPU brand on the host.
type GPUVendor int

const (
	GPUVendorNone  GPUVendor = iota
	GPUVendorNVIDIA          // use T-Rex miner (KawPoW)
	GPUVendorAMD             // use TeamRedMiner (KawPoW)
	GPUVendorOther           // generic / Intel — not supported for KawPoW
)

// GPUInfo holds detected GPU metadata.
type GPUInfo struct {
	Vendor GPUVendor
	Model  string
}

// GPUMinerStats is polled from the miner's local HTTP API.
type GPUMinerStats struct {
	Hashrate15s  float64
	Hashrate1m   float64
	Hashrate15m  float64
	GPUTempC     *int
	GPUUsagePct  *int
	ActiveAlgo   string
}

// rvnEndpoint is one pool entry for the GPU miner (primary or backup).
type rvnEndpoint struct {
	host string
	port int
	tls  bool
	pass string
}

// GPUMiner manages one GPU miner sub-process (T-Rex or TeamRedMiner).
type GPUMiner struct {
	cfg        config.RuntimeConfig
	info       GPUInfo
	installDir string

	mu     sync.RWMutex
	stats  GPUMinerStats
	active bool
	paused bool
	proc   *os.Process // currently running subprocess (nil if stopped)

	stopCh   chan struct{}
	pauseCh  chan struct{} // closed when paused, re-created on resume
	resumeCh chan struct{} // closed when resuming from pause
	pauseMu  sync.Mutex
	wg       sync.WaitGroup
}

// newGPUMiner creates a GPUMiner if GPU mining is configured and a supported GPU is detected.
// Returns nil if GPU mining should not run.
func newGPUMiner(cfg config.RuntimeConfig) *GPUMiner {
	if !cfg.GPUEnabled || cfg.RVNWallet == "" {
		return nil
	}
	info := detectGPU()
	if info.Vendor == GPUVendorNone || info.Vendor == GPUVendorOther {
		log.Printf("[gpu] GPU mining enabled but no supported GPU detected (vendor=%v model=%q)", info.Vendor, info.Model)
		return nil
	}
	installDir, err := cfg.InstallDirectory()
	if err != nil {
		log.Printf("[gpu] cannot determine install dir: %v", err)
		return nil
	}
	log.Printf("[gpu] detected %s — will run KawPoW miner for RVN", info.Model)
	g := &GPUMiner{
		cfg:        cfg,
		info:       info,
		installDir: installDir,
		stopCh:     make(chan struct{}),
		pauseCh:    make(chan struct{}),
		resumeCh:   make(chan struct{}),
	}
	// pauseCh starts open; waitIfPaused hits the default branch and returns
	// true immediately, so no pre-close of resumeCh is needed (and
	// pre-closing it would break the first Pause() — the inner select would
	// fire on the already-closed channel instead of blocking).
	return g
}

// Start downloads (if needed) and launches the GPU miner, then polls stats.
func (g *GPUMiner) Start() {
	g.wg.Add(1)
	go func() {
		defer g.wg.Done()
		g.run()
	}()
}

// Stop shuts down the GPU miner and waits for it to exit.
func (g *GPUMiner) Stop() {
	// Resume first so the run loop is not blocked on pauseCh when stop fires.
	g.Resume()
	select {
	case <-g.stopCh:
	default:
		close(g.stopCh)
	}
	g.wg.Wait()
}

// Pause suspends KawPoW polling and kills the running miner subprocess until
// Resume is called. Safe to call multiple times.
func (g *GPUMiner) Pause() {
	g.pauseMu.Lock()
	defer g.pauseMu.Unlock()
	g.mu.Lock()
	already := g.paused
	if !already {
		g.paused = true
		// Kill the running process so it stops consuming GPU.
		if g.proc != nil {
			_ = g.proc.Kill()
		}
	}
	g.mu.Unlock()
	if !already {
		// Signal the run loop to enter the paused wait.
		select {
		case <-g.pauseCh:
		default:
			close(g.pauseCh)
		}
		log.Printf("[gpu] miner paused by remote command")
	}
}

// Resume restarts the KawPoW miner after a Pause. Safe to call when not paused.
func (g *GPUMiner) Resume() {
	g.pauseMu.Lock()
	defer g.pauseMu.Unlock()
	g.mu.Lock()
	wasPaused := g.paused
	g.paused = false
	g.mu.Unlock()
	if wasPaused {
		// Unblock the run loop waiting on resumeCh, then reset both channels.
		select {
		case <-g.resumeCh:
		default:
			close(g.resumeCh)
		}
		g.pauseCh = make(chan struct{})
		g.resumeCh = make(chan struct{})
		log.Printf("[gpu] miner resumed by remote command")
	}
}

// waitIfPaused blocks the run loop while paused, returning false if stop fires.
func (g *GPUMiner) waitIfPaused() bool {
	g.pauseMu.Lock()
	pauseCh := g.pauseCh
	resumeCh := g.resumeCh
	g.pauseMu.Unlock()

	select {
	case <-pauseCh:
		// Paused — wait for resume or stop.
		select {
		case <-g.stopCh:
			return false
		case <-resumeCh:
			return true
		}
	default:
		return true
	}
}

// Stats returns the latest GPU mining statistics.
func (g *GPUMiner) Stats() (GPUMinerStats, bool) {
	g.mu.RLock()
	defer g.mu.RUnlock()
	return g.stats, g.active
}

// GPUModel returns the detected GPU model string.
func (g *GPUMiner) GPUModel() string {
	return g.info.Model
}

// buildPoolList returns the primary pool followed by any configured backups.
func (g *GPUMiner) buildPoolList() []rvnEndpoint {
	eps := []rvnEndpoint{{
		host: g.cfg.RVNPoolHost,
		port: g.cfg.RVNPoolPort,
		tls:  g.cfg.RVNPoolTLS,
		pass: g.cfg.RVNPoolPass,
	}}
	for _, bp := range g.cfg.RVNBackupPools {
		if bp.Host != "" && bp.Port > 0 {
			eps = append(eps, rvnEndpoint{
				host: bp.Host,
				port: bp.Port,
				tls:  bp.TLS,
				pass: bp.Pass,
			})
		}
	}
	return eps
}

func (g *GPUMiner) run() {
	binPath, err := g.ensureMinerBinary()
	if err != nil {
		log.Printf("[gpu] could not obtain miner binary: %v", err)
		return
	}

	pools := g.buildPoolList()
	poolIdx := 0
	const retryDelay = 30 * time.Second

	for {
		select {
		case <-g.stopCh:
			return
		default:
		}

		if !g.waitIfPaused() {
			return
		}

		ep := pools[poolIdx%len(pools)]
		proc, err := g.startProcessOnPool(binPath, ep)
		if err != nil {
			log.Printf("[gpu] failed to start miner: %v — retry in %s (pool %d/%d)", err, retryDelay, poolIdx%len(pools)+1, len(pools))
			select {
			case <-g.stopCh:
				return
			case <-time.After(retryDelay):
			}
			poolIdx++
			continue
		}

		g.mu.Lock()
		g.active = true
		g.proc = proc
		g.mu.Unlock()

		log.Printf("[gpu] %s started (pid=%d) → %s:%d", g.spec().fileName, proc.Pid, ep.host, ep.port)

		// pollStop signals pollStats to exit; closed when this iteration ends.
		pollStop := make(chan struct{})
		pollDone := make(chan struct{})
		go func() {
			defer close(pollDone)
			g.pollStats(pollStop)
		}()

		// Wait for process exit in a goroutine so we can also listen for stop.
		waitDone := make(chan error, 1)
		go func() {
			_, werr := proc.Wait()
			waitDone <- werr
		}()

		var stopRequested bool
		select {
		case <-g.stopCh:
			// Agent shutting down — kill the miner process immediately.
			stopRequested = true
			_ = proc.Kill()
			<-waitDone
		case waitErr := <-waitDone:
			if waitErr != nil {
				log.Printf("[gpu] miner exited: %v — rotating to next pool", waitErr)
			}
			// Miner crashed or exited cleanly — rotate to next pool on retry.
			poolIdx++
		}

		close(pollStop)
		<-pollDone

		g.mu.Lock()
		g.active = false
		g.proc = nil
		g.mu.Unlock()

		if stopRequested {
			return
		}

		// Wait before retrying, but exit cleanly if Stop() is called.
		select {
		case <-g.stopCh:
			return
		case <-time.After(retryDelay):
		}
	}
}

// pollStats polls the miner's HTTP API until stop is closed.
func (g *GPUMiner) pollStats(stop <-chan struct{}) {
	apiPort := g.apiPort()
	ticker := time.NewTicker(10 * time.Second)
	defer ticker.Stop()

	samples := make([]float64, 0, 90) // 15 min at 10s intervals

	for {
		select {
		case <-stop:
			return
		case <-ticker.C:
			hr, tempC, usage, err := fetchMinerStats(g.info.Vendor, apiPort)
			if err != nil {
				continue
			}
			samples = append(samples, hr)
			if len(samples) > 90 {
				samples = samples[len(samples)-90:]
			}

			g.mu.Lock()
			g.stats = GPUMinerStats{
				Hashrate15s: hr,
				Hashrate1m:  avg(samples, 6),
				Hashrate15m: avg(samples, len(samples)),
				GPUTempC:    tempC,
				GPUUsagePct: usage,
				ActiveAlgo:  "kawpow",
			}
			g.mu.Unlock()
		}
	}
}

func avg(samples []float64, last int) float64 {
	if len(samples) == 0 || last <= 0 {
		return 0
	}
	if last > len(samples) {
		last = len(samples)
	}
	slice := samples[len(samples)-last:]
	var sum float64
	for _, v := range slice {
		sum += v
	}
	return sum / float64(len(slice))
}

func (g *GPUMiner) apiPort() int {
	switch g.info.Vendor {
	case GPUVendorNVIDIA:
		return 4067
	case GPUVendorAMD:
		return 4068
	default:
		return 4067
	}
}

// buildPoolURL constructs the stratum URL for a given pool endpoint.
func buildPoolURL(ep rvnEndpoint) string {
	scheme := "stratum+tcp"
	if ep.tls {
		scheme = "stratum+ssl"
	}
	return fmt.Sprintf("%s://%s:%d", scheme, ep.host, ep.port)
}

// ---- Miner binary management ----

type minerSpec struct {
	fileName   string
	downloadURL string
}

func (g *GPUMiner) spec() minerSpec {
	switch g.info.Vendor {
	case GPUVendorNVIDIA:
		return minerSpec{
			fileName:    "t-rex.exe",
			downloadURL: "https://github.com/trexminer/T-Rex/releases/download/0.26.8/t-rex-0.26.8-win.zip",
		}
	default: // AMD
		return minerSpec{
			fileName:    "teamredminer.exe",
			downloadURL: "https://github.com/todxx/teamredminer/releases/download/v0.10.21/teamredminer-v0.10.21-win.zip",
		}
	}
}

func (g *GPUMiner) ensureMinerBinary() (string, error) {
	spec := g.spec()

	// 1. Check the agent's install directory first.
	binPath := filepath.Join(g.installDir, spec.fileName)
	if _, err := os.Stat(binPath); err == nil {
		return binPath, nil
	}

	// 2. Check the directory that contains the running agent binary (side-by-side).
	if exePath, err := os.Executable(); err == nil {
		sideBySide := filepath.Join(filepath.Dir(exePath), spec.fileName)
		if _, err := os.Stat(sideBySide); err == nil {
			log.Printf("[gpu] found %s next to agent binary, using local copy", spec.fileName)
			return sideBySide, nil
		}
	}

	// 3. Fall back to downloading from GitHub.
	log.Printf("[gpu] GPU miner binary not found locally, downloading from GitHub (this may fail on restricted networks)")
	if err := downloadAndExtract(spec.downloadURL, g.installDir, spec.fileName); err != nil {
		return "", fmt.Errorf("download failed: %w", err)
	}
	if _, err := os.Stat(binPath); err != nil {
		return "", fmt.Errorf("binary not found after download: %s", binPath)
	}
	return binPath, nil
}

func downloadAndExtract(url, destDir, targetFile string) error {
	client := &http.Client{Timeout: 5 * time.Minute}
	resp, err := client.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP %d from %s", resp.StatusCode, url)
	}
	data, err := io.ReadAll(io.LimitReader(resp.Body, 512<<20))
	if err != nil {
		return err
	}
	return extractZipFile(data, destDir, targetFile)
}

// ---- Miner HTTP API polling ----

// T-Rex summary response (subset we care about).
type trexSummary struct {
	Hashrate int `json:"hashrate"`
	GPUs     []struct {
		Temperature int `json:"temperature"`
		GpuLoad     int `json:"gpu_load"`
	} `json:"gpus"`
}

// TeamRedMiner status response (subset).
type trmStatus struct {
	Algorithms []struct {
		Name      string  `json:"algorithm"`
		TotalMHs  float64 `json:"mhsh_total"`
	} `json:"algorithms"`
	GPUs []struct {
		TempC   int `json:"temp_c"`
		Fan     int `json:"fan_pct"`
	} `json:"gpus"`
}

func fetchMinerStats(vendor GPUVendor, port int) (hashrate float64, tempC, usagePct *int, err error) {
	url := fmt.Sprintf("http://127.0.0.1:%d/summary", port)
	resp, e := http.Get(url) //nolint:noctx
	if e != nil {
		return 0, nil, nil, e
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(resp.Body)

	switch vendor {
	case GPUVendorNVIDIA:
		var s trexSummary
		if e := json.Unmarshal(body, &s); e != nil {
			return 0, nil, nil, e
		}
		hashrate = float64(s.Hashrate)
		if len(s.GPUs) > 0 {
			t := s.GPUs[0].Temperature
			u := s.GPUs[0].GpuLoad
			tempC = &t
			usagePct = &u
		}
	case GPUVendorAMD:
		var s trmStatus
		if e := json.Unmarshal(body, &s); e != nil {
			return 0, nil, nil, e
		}
		for _, a := range s.Algorithms {
			if a.Name == "kawpow" || a.Name == "KawPoW" {
				hashrate = a.TotalMHs * 1e6 // convert MH/s → H/s
			}
		}
		if len(s.GPUs) > 0 {
			t := s.GPUs[0].TempC
			u := s.GPUs[0].Fan
			tempC = &t
			usagePct = &u
		}
	}
	return
}