AetherForge/server/internal/api/websocket.go

package api

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

	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/models"
	"crypto-miner-server/internal/pool"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
)

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
	poolManager     *pool.Manager
	defaultPool     pool.Config
	aiHandler       *AIHandler
	agentConfigs    map[string]AgentForgeConfig
	agentLogs       map[string]string
	serverPolicy    ServerPolicy
	pingIntervalSec int
	mu              sync.RWMutex
}

func NewWSHub(database *db.Database) *WSHub {
	return &WSHub{
		db:              database,
		agents:          make(map[string]*AgentConnection),
		dashboards:      make(map[string]*websocket.Conn),
		agentConfigs:    make(map[string]AgentForgeConfig),
		agentLogs:       make(map[string]string),
		pingIntervalSec: 30,
	}
}

func (h *WSHub) SetServerPolicy(p ServerPolicy) {
	h.mu.Lock()
	h.serverPolicy = p
	h.mu.Unlock()
}

func (h *WSHub) SetPingInterval(seconds int) {
	if seconds < 10 {
		seconds = 30
	}
	h.mu.Lock()
	h.pingIntervalSec = seconds
	h.mu.Unlock()
}

func (h *WSHub) pingInterval() time.Duration {
	h.mu.RLock()
	sec := h.pingIntervalSec
	h.mu.RUnlock()
	if sec < 10 {
		sec = 30
	}
	return time.Duration(sec) * time.Second
}

func (h *WSHub) runPingLoop(conn *websocket.Conn) {
	interval := h.pingInterval()
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	_ = conn.SetReadDeadline(time.Now().Add(interval * 2))
	conn.SetPongHandler(func(string) error {
		return conn.SetReadDeadline(time.Now().Add(interval * 2))
	})

	for range ticker.C {
		if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil {
			return
		}
	}
}

func (h *WSHub) serverPolicySnapshot() ServerPolicy {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.serverPolicy
}

func (h *WSHub) connectedAgentCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.agents)
}

func (h *WSHub) isAgentConnected(agentID string) bool {
	h.mu.RLock()
	defer h.mu.RUnlock()
	_, ok := h.agents[agentID]
	return ok
}

func (h *WSHub) SetPoolManager(manager *pool.Manager, defaultCfg pool.Config) {
	h.mu.Lock()
	h.poolManager = manager
	h.defaultPool = defaultCfg
	h.mu.Unlock()
}

func (h *WSHub) SetAIHandler(ai *AIHandler) {
	h.mu.Lock()
	h.aiHandler = ai
	h.mu.Unlock()
}

func (h *WSHub) agentPoolConfig(agentID string) pool.Config {
	h.mu.RLock()
	defer h.mu.RUnlock()
	cfg := h.agentConfigs[agentID]
	poolCfg := pool.Config{
		Host:   cfg.poolHostOrDefault(h.defaultPool.Host),
		Port:   cfg.poolPortOrDefault(h.defaultPool.Port),
		Wallet: cfg.Wallet,
	}
	if cfg.PoolHost == "" {
		poolCfg.UseTLS = h.defaultPool.UseTLS
	} else {
		poolCfg.UseTLS = cfg.PoolTLS
	}
	if cfg.PoolPass != "" {
		poolCfg.Password = cfg.PoolPass
	} else if h.defaultPool.Password != "" {
		poolCfg.Password = h.defaultPool.Password
	} else {
		poolCfg.Password = "x"
	}
	if poolCfg.Wallet == "" {
		poolCfg.Wallet = h.defaultPool.Wallet
	}
	return poolCfg
}

func (h *WSHub) getAgentConn(agentID string) *AgentConnection {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.agents[agentID]
}

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
	}

	go h.runPingLoop(conn)

	agentID := ""
	defer func() {
		if agentID != "" {
			h.mu.Lock()
			delete(h.agents, agentID)
			delete(h.agentConfigs, agentID)
			h.mu.Unlock()
			if h.aiHandler != nil {
				h.aiHandler.RemoveEngine(agentID)
			}
			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"`
				WorkerName       string `json:"worker_name"`
				CPUCores         int    `json:"cpu_cores"`
				MemoryGB         int    `json:"memory_gb"`
				PoolHost         string `json:"pool_host"`
				PoolPort         int    `json:"pool_port"`
				PoolTLS          bool   `json:"pool_tls"`
				PoolPass         string `json:"pool_pass"`
				AIEnabled        bool   `json:"ai_enabled"`
				AIOllamaEndpoint string `json:"ai_ollama_endpoint"`
				AIModel          string `json:"ai_model"`
			}
			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.WorkerName
			if displayName == "" {
				displayName = auth.Worker
			}
			if displayName == "" {
				displayName = auth.Hostname
			}
			if displayName == "" {
				displayName = agentID[:8]
			}

			policy := h.serverPolicySnapshot()
			if policy.MaxAgents > 0 && !h.isAgentConnected(agentID) && h.connectedAgentCount() >= policy.MaxAgents {
				conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
					"success": false, "error": "fleet agent limit reached",
				})})
				continue
			}

			forgeCfg := AgentForgeConfig{
				Wallet:           auth.Wallet,
				PoolHost:         auth.PoolHost,
				PoolPort:         auth.PoolPort,
				PoolTLS:          auth.PoolTLS,
				PoolPass:         auth.PoolPass,
				AIEnabled:        auth.AIEnabled,
				AIOllamaEndpoint: auth.AIOllamaEndpoint,
				AIModel:          auth.AIModel,
			}

			h.mu.Lock()
			h.agentConfigs[agentID] = forgeCfg
			h.mu.Unlock()

			if h.poolManager != nil {
				poolCfg := h.agentPoolConfig(agentID)
				if poolCfg.Password == "" {
					poolCfg.Password = "x"
				}
				if _, err := h.poolManager.EnsurePool(&poolCfg); err != nil {
					log.Printf("[WS] Failed to ensure forged pool for agent %s: %v", agentID, err)
				}
			}

			if h.aiHandler != nil && forgeCfg.AIEnabled {
				h.aiHandler.SetEngineForAgent(agentID, forgeCfg.AIOllamaEndpoint, forgeCfg.AIModel)
			}

			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)
				conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(map[string]interface{}{
					"success": false, "error": "database error",
				})})
				continue
			}

			if policy.LogAgentConnections {
				log.Printf("[WS] Agent connected: id=%s name=%s ip=%s", agentID, displayName, clientIP)
			}

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

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

			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()
			share.Accepted = false

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

			sendShareResult := func(accepted bool, errMsg string) {
				share.Accepted = accepted
				share.Error = errMsg
				if err := h.db.UpdateShareResult(shareID, accepted, errMsg); err != nil {
					log.Printf("Failed to update share result: %v", err)
				}
				if h.serverPolicySnapshot().LogShareSubmissions {
					log.Printf("[WS] Share agent=%s job=%s accepted=%v err=%q", agentID, share.JobID, accepted, errMsg)
				}

				agentConn := h.getAgentConn(agentID)
				if agentConn != nil {
					result := map[string]interface{}{
						"job_id":   share.JobID,
						"accepted": accepted,
					}
					if errMsg != "" {
						result["error"] = errMsg
					}
					_ = agentConn.SendJSON(Message{Type: "share_result", Payload: mustMarshal(result)})
				}

				h.broadcastDashboard(Message{
					Type: "new_share",
					Payload: mustMarshal(map[string]interface{}{
						"id":        shareID,
						"agent_id":  agentID,
						"job_id":    share.JobID,
						"accepted":  accepted,
						"hash":      share.Hash,
						"nonce":     share.Nonce,
						"error":     errMsg,
						"timestamp": share.Timestamp,
					}),
				})
			}

			if h.poolManager == nil {
				sendShareResult(false, "pool manager not configured")
				continue
			}

			poolCfg := h.agentPoolConfig(agentID)
			proxy := h.poolManager.GetPool(&poolCfg)
			if proxy == nil {
				if p, err := h.poolManager.EnsurePool(&poolCfg); err == nil {
					proxy = p
				} else {
					sendShareResult(false, "pool not connected: "+err.Error())
					continue
				}
			}

			if !proxy.IsConnected() {
				sendShareResult(false, "pool not connected")
				continue
			}

			wallet := poolCfg.Wallet
			if wallet == "" {
				wallet = h.defaultPool.Wallet
			}

			proxy.SubmitShare(agentID, wallet, share.JobID, share.Nonce, share.Hash, sendShareResult)

		case "get_job":
			var proxy *pool.Proxy
			if h.poolManager != nil {
				poolCfg := h.agentPoolConfig(agentID)
				proxy = h.poolManager.GetPool(&poolCfg)
				if proxy == nil {
					if p, err := h.poolManager.EnsurePool(&poolCfg); err == nil {
						proxy = p
					}
				}
			}
			if proxy != nil {
				job := proxy.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"})})
			}

		case "log_tail":
			var payload struct {
				Content string `json:"content"`
				Lines   int    `json:"lines"`
			}
			if err := json.Unmarshal(msg.Payload, &payload); err != nil {
				continue
			}
			h.mu.Lock()
			h.agentLogs[agentID] = payload.Content
			h.mu.Unlock()
			h.broadcastDashboard(Message{
				Type:    "agent_log",
				Payload: mustMarshal(map[string]interface{}{"agent_id": agentID, "content": payload.Content}),
			})

		case "command_result":
			var payload map[string]interface{}
			if err := json.Unmarshal(msg.Payload, &payload); err != nil {
				continue
			}
			payload["agent_id"] = agentID
			h.broadcastDashboard(Message{Type: "command_result", Payload: mustMarshal(payload)})
		}
	}
}

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,
	})})

	go h.runPingLoop(conn)

	// 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()
			id := id
			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)
		}
	}
}

// SendToAgent sends a message to one connected agent.
func (h *WSHub) SendToAgent(agentID string, msg Message) error {
	agent := h.getAgentConn(agentID)
	if agent == nil {
		return fmt.Errorf("agent %s not connected", agentID)
	}
	return agent.SendJSON(msg)
}

// SendAgentCommand sends a remote command to an agent.
func (h *WSHub) SendAgentCommand(agentID, action string, args map[string]interface{}) error {
	payload := map[string]interface{}{"action": action}
	for k, v := range args {
		payload[k] = v
	}
	return h.SendToAgent(agentID, Message{Type: "command", Payload: mustMarshal(payload)})
}

// BroadcastAgentCommand sends a remote command to all connected agents.
func (h *WSHub) BroadcastAgentCommand(action string, args map[string]interface{}) {
	payload := map[string]interface{}{"action": action}
	for k, v := range args {
		payload[k] = v
	}
	h.BroadcastToAgents(Message{Type: "command", Payload: mustMarshal(payload)})
}

func (h *WSHub) GetAgentLog(agentID string) string {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.agentLogs[agentID]
}

func (h *WSHub) BroadcastFleetAlert(ev interface{}) {
	h.broadcastDashboard(Message{Type: "fleet_alert", Payload: mustMarshal(ev)})
}

func (h *WSHub) BroadcastPoolStatus(status interface{}) {
	h.broadcastDashboard(Message{Type: "pool_status", Payload: mustMarshal(status)})
}

func (h *WSHub) BroadcastAIActivity(entry interface{}) {
	h.broadcastDashboard(Message{Type: "ai_activity", Payload: mustMarshal(entry)})
}