AetherForge/server/internal/api/websocket.go

package api

import (
	"encoding/json"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/models"
	"crypto-miner-server/internal/pool"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  4096,
	WriteBufferSize: 4096,
	CheckOrigin: func(r *http.Request) bool {
		return true // Allow all origins for local use
	},
}

type Message struct {
	Type    string          `json:"type"`
	Payload json.RawMessage `json:"payload"`
}

type AgentConnection struct {
	AgentID string
	Conn    *websocket.Conn
	mu      sync.Mutex
}

func (c *AgentConnection) SendJSON(v interface{}) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.Conn.WriteJSON(v)
}

type WSHub struct {
	db            *db.Database
	agents        map[string]*AgentConnection
	dashboards    map[string]*websocket.Conn
	poolProxy     *pool.Proxy
	defaultAgent  AgentDefaults
	mu            sync.RWMutex
}

type AgentDefaults struct {
	Threads     int
	CPUPriority string
}

func NewWSHub(database *db.Database) *WSHub {
	return &WSHub{
		db:         database,
		agents:     make(map[string]*AgentConnection),
		dashboards: make(map[string]*websocket.Conn),
		defaultAgent: AgentDefaults{Threads: 4, CPUPriority: "below_normal"},
	}
}

func (h *WSHub) SetDefaultAgentConfig(cfg interface{}) {
	type defaults struct {
		Threads     int    `json:"threads"`
		CPUPriority string `json:"cpu_priority"`
	}
	if cfg == nil {
		return
	}
	data, err := json.Marshal(cfg)
	if err != nil {
		return
	}
	var d defaults
	if err := json.Unmarshal(data, &d); err != nil {
		return
	}
	if d.Threads <= 0 {
		d.Threads = 4
	}
	if d.CPUPriority == "" {
		d.CPUPriority = "below_normal"
	}
	h.mu.Lock()
	h.defaultAgent = AgentDefaults{Threads: d.Threads, CPUPriority: d.CPUPriority}
	h.mu.Unlock()
}

// SetPoolProxy sets the pool proxy for share submission forwarding
func (h *WSHub) SetPoolProxy(proxy *pool.Proxy) {
	h.poolProxy = proxy
}

func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade error: %v", err)
		return
	}

	agentID := ""
	defer func() {
		if agentID != "" {
			h.mu.Lock()
			delete(h.agents, agentID)
			h.mu.Unlock()
			h.db.SetAgentOffline(agentID)
			h.broadcastDashboard(Message{
				Type: "agent_offline",
				Payload: mustMarshal(map[string]string{"agent_id": agentID}),
			})
		}
		conn.Close()
	}()

	for {
		_, msgBytes, err := conn.ReadMessage()
		if err != nil {
			log.Printf("Agent read error: %v", err)
			break
		}

		var msg Message
		if err := json.Unmarshal(msgBytes, &msg); err != nil {
			log.Printf("Invalid message from agent: %v", err)
			continue
		}

		switch msg.Type {
		case "auth":
			var auth struct {
				AgentID  string `json:"agent_id"`
				Wallet   string `json:"wallet"`
				Version  string `json:"version"`
				Hostname string `json:"hostname"`
				Worker   string `json:"worker"`
				CPUCores int    `json:"cpu_cores"`
				MemoryGB int    `json:"memory_gb"`
			}
			if err := json.Unmarshal(msg.Payload, &auth); err != nil {
				conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
					"success": false, "error": "invalid auth payload",
				})})
				continue
			}

			agentID = auth.AgentID
			if agentID == "" {
				agentID = uuid.New().String()
			}

			displayName := auth.Worker
			if displayName == "" {
				displayName = auth.Hostname
			}
			if displayName == "" {
				displayName = agentID[:8]
			}

			clientIP := r.Header.Get("X-Forwarded-For")
			if clientIP == "" {
				clientIP = r.RemoteAddr
			}
			if idx := strings.LastIndex(clientIP, ":"); idx > 0 && strings.Count(clientIP, ":") == 1 {
				clientIP = clientIP[:idx]
			}

			agent := &models.Agent{
				ID:       agentID,
				Name:     displayName,
				Wallet:   auth.Wallet,
				IP:       clientIP,
				Version:  auth.Version,
				Status:   "online",
				CPUCores: auth.CPUCores,
				MemoryGB: auth.MemoryGB,
				LastSeen: time.Now(),
			}

			if err := h.db.UpsertAgent(agent); err != nil {
				log.Printf("Failed to upsert agent: %v", err)
			}

			h.mu.Lock()
			h.agents[agentID] = &AgentConnection{AgentID: agentID, Conn: conn}
			h.mu.Unlock()

			h.mu.RLock()
			defaults := h.defaultAgent
			h.mu.RUnlock()

			conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
				"success":  true,
				"agent_id": agentID,
				"config": map[string]interface{}{
					"threads":  defaults.Threads,
					"priority": defaults.CPUPriority,
				},
			})})

			h.broadcastDashboard(Message{
				Type: "agent_online",
				Payload: mustMarshal(agent),
			})

		case "stats":
			var stats struct {
				Hashrate15s     float64 `json:"hashrate_15s"`
				Hashrate1m      float64 `json:"hashrate_1m"`
				Hashrate15m     float64 `json:"hashrate_15m"`
				SharesSubmitted int     `json:"shares_submitted"`
				SharesAccepted  int     `json:"shares_accepted"`
				CPUUsagePct     float64 `json:"cpu_usage_pct"`
				MemoryUsagePct  float64 `json:"memory_usage_pct"`
				UptimeSeconds   int     `json:"uptime_seconds"`
			}
			if err := json.Unmarshal(msg.Payload, &stats); err != nil {
				continue
			}

			sharesBad := stats.SharesSubmitted - stats.SharesAccepted
			if sharesBad < 0 {
				sharesBad = 0
			}

			h.db.UpdateAgentStats(agentID, stats.Hashrate15s, stats.Hashrate1m, stats.Hashrate15m,
				stats.SharesSubmitted, stats.SharesAccepted, sharesBad,
				stats.CPUUsagePct, stats.MemoryUsagePct, stats.UptimeSeconds)

			h.db.InsertHashrateSample(agentID, stats.Hashrate15m)

			h.broadcastDashboard(Message{
				Type: "stats_update",
				Payload: mustMarshal(map[string]interface{}{
					"agent_id": agentID,
					"hashrate_15s": stats.Hashrate15s,
					"hashrate_1m":  stats.Hashrate1m,
					"hashrate_15m": stats.Hashrate15m,
					"cpu_usage_pct": stats.CPUUsagePct,
				}),
			})

		case "submit_share":
			var share models.Share
			if err := json.Unmarshal(msg.Payload, &share); err != nil {
				continue
			}
			share.AgentID = agentID
			share.Timestamp = time.Now()

			// Forward share to pool proxy if connected
			if h.poolProxy != nil && h.poolProxy.IsConnected() {
				h.poolProxy.SubmitShare(agentID, share.JobID, share.Nonce, share.Hash)
				share.Accepted = true // Pool will validate; we assume accepted initially
			} else {
				// Pool not connected - mark as accepted locally for testing
				share.Accepted = true
				log.Printf("[WS] Pool not connected, marking share as accepted locally")
			}

			if err := h.db.InsertShare(&share); err != nil {
				log.Printf("Failed to insert share: %v", err)
			}

			conn.WriteJSON(Message{Type: "share_result", Payload: mustMarshal(map[string]interface{}{
				"job_id":   share.JobID,
				"accepted": share.Accepted,
			})})

			h.broadcastDashboard(Message{
				Type: "new_share",
				Payload: mustMarshal(map[string]interface{}{
					"agent_id":  agentID,
					"accepted":  share.Accepted,
					"hash":      share.Hash,
				}),
			})

		case "get_job":
			// Agent requesting current job from pool
			if h.poolProxy != nil {
				job := h.poolProxy.GetCurrentJob()
				if job != nil {
					conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(job)})
				} else {
					conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(map[string]string{"error": "no job available"})})
				}
			} else {
				conn.WriteJSON(Message{Type: "new_job", Payload: mustMarshal(map[string]string{"error": "pool not connected"})})
			}
		}
	}
}

func (h *WSHub) HandleDashboardWS(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("Dashboard WebSocket upgrade error: %v", err)
		return
	}

	dashID := uuid.New().String()
	h.mu.Lock()
	h.dashboards[dashID] = conn
	h.mu.Unlock()

	defer func() {
		h.mu.Lock()
		delete(h.dashboards, dashID)
		h.mu.Unlock()
		conn.Close()
	}()

	// Send initial data
	agents, _ := h.db.ListAgents()
	stats, _ := h.db.GetFleetStats()

	conn.WriteJSON(Message{Type: "init", Payload: mustMarshal(map[string]interface{}{
		"agents": agents,
		"stats":  stats,
	})})

	// Keep connection alive, read close messages
	for {
		_, _, err := conn.ReadMessage()
		if err != nil {
			break
		}
	}
}

func (h *WSHub) broadcastDashboard(msg Message) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	data, err := json.Marshal(msg)
	if err != nil {
		return
	}

	for id, conn := range h.dashboards {
		if err := conn.WriteMessage(websocket.TextMessage, data); err != nil {
			log.Printf("Failed to send to dashboard %s: %v", id, err)
			conn.Close()
			go func() {
				h.mu.Lock()
				delete(h.dashboards, id)
				h.mu.Unlock()
			}()
		}
	}
}

func mustMarshal(v interface{}) json.RawMessage {
	data, _ := json.Marshal(v)
	return data
}

// BroadcastToAgents sends a message to all connected agents
func (h *WSHub) BroadcastToAgents(msg Message) {
	h.mu.RLock()
	defer h.mu.RUnlock()

	for id, agent := range h.agents {
		if err := agent.SendJSON(msg); err != nil {
			log.Printf("Failed to send to agent %s: %v", id, err)
		}
	}
}

// mustMarshalRaw marshals a value to json.RawMessage, panicking on error
func mustMarshalRaw(v interface{}) json.RawMessage {
	data, err := json.Marshal(v)
	if err != nil {
		panic(err)
	}
	return data
}