AetherForge/server/internal/api/ai_handler.go

package api

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

	"crypto-miner-server/internal/db"
	"crypto-miner-server/internal/ollama"
)

// ──────────────────────────────────────────────
//  AI Autonomy Handler
//  ──────────────────────────────────────────────
//  Provides REST endpoints for the AI Autonomy feature.
//  Agents call these endpoints to get decisions from Ollama
//  and report tool execution results.
// ──────────────────────────────────────────────

// AIHandler manages AI autonomy endpoints.
type AIHandler struct {
	db       *db.Database
	engines  map[string]*ollama.Engine // agentID -> engine (each agent can have its own Ollama config)
	reports  []ollama.Report           // recent tool execution reports
	activity map[string]AIActivityEntry
	onEvent  func(AIActivityEntry)
	mu       sync.RWMutex
}

// AIActivityEntry summarizes recent AI cycles per agent.
type AIActivityEntry struct {
	AgentID       string    `json:"agent_id"`
	LastDecideAt  time.Time `json:"last_decide_at,omitempty"`
	LastAction    string    `json:"last_action,omitempty"`
	LastTool      string    `json:"last_tool,omitempty"`
	ToolCallCount int       `json:"tool_call_count"`
	LastReasoning string    `json:"last_reasoning,omitempty"`
	LastReportAt  time.Time `json:"last_report_at,omitempty"`
	LastSuccess   bool      `json:"last_success"`
}

// NewAIHandler creates a new AI handler.
func NewAIHandler(database *db.Database) *AIHandler {
	return &AIHandler{
		db:       database,
		engines:  make(map[string]*ollama.Engine),
		reports:  make([]ollama.Report, 0, 1000),
		activity: make(map[string]AIActivityEntry),
	}
}

func (h *AIHandler) SetEventBroadcaster(fn func(AIActivityEntry)) {
	h.mu.Lock()
	h.onEvent = fn
	h.mu.Unlock()
}

// SetEngineForAgent sets or updates the Ollama engine for a specific agent.
// This is called when an agent authenticates with AI settings.
func (h *AIHandler) SetEngineForAgent(agentID, ollamaEndpoint, model string) {
	h.mu.Lock()
	defer h.mu.Unlock()

	if ollamaEndpoint == "" {
		ollamaEndpoint = "http://localhost:11434"
	}
	if model == "" {
		model = "llama3.2"
	}

	h.engines[agentID] = ollama.NewEngine(ollamaEndpoint, model)
	log.Printf("[AI] Engine set for agent %s (endpoint=%s, model=%s)", agentID, ollamaEndpoint, model)
}

// RemoveEngine removes an agent's engine (on disconnect).
func (h *AIHandler) RemoveEngine(agentID string) {
	h.mu.Lock()
	defer h.mu.Unlock()
	delete(h.engines, agentID)
}

// GetEngine returns the engine for an agent.
func (h *AIHandler) GetEngine(agentID string) *ollama.Engine {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return h.engines[agentID]
}

// HandleDecide handles POST /api/v1/agent/decide
func (h *AIHandler) HandleDecide(w http.ResponseWriter, r *http.Request) {
	h.handleDecide(w, r)
}

// HandleReport handles POST /api/v1/agent/report
func (h *AIHandler) HandleReport(w http.ResponseWriter, r *http.Request) {
	h.handleReport(w, r)
}

// HandleHeartbeat handles POST /api/v1/agent/heartbeat
func (h *AIHandler) HandleHeartbeat(w http.ResponseWriter, r *http.Request) {
	h.handleHeartbeat(w, r)
}

// ─── Decide ───────────────────────────────────

type decideRequest struct {
	AgentID        string `json:"agent_id"`
	OllamaEndpoint string `json:"ollama_endpoint,omitempty"`
	Model          string `json:"model,omitempty"`
	ollama.AgentState
}

func (h *AIHandler) handleDecide(w http.ResponseWriter, r *http.Request) {
	var req decideRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid request: "+err.Error(), http.StatusBadRequest)
		return
	}

	if req.AgentID == "" {
		http.Error(w, "agent_id is required", http.StatusBadRequest)
		return
	}

	// Get or create engine for this agent
	engine := h.GetEngine(req.AgentID)
	if engine == nil {
		// Create engine on first request
		h.SetEngineForAgent(req.AgentID, req.OllamaEndpoint, req.Model)
		engine = h.GetEngine(req.AgentID)
	}

	if engine == nil {
		http.Error(w, "failed to create AI engine", http.StatusInternalServerError)
		return
	}

	// Call Ollama for decision
	resp, err := engine.Decide(&req.AgentState)
	if err != nil {
		log.Printf("[AI] Decide error for agent %s: %v", req.AgentID, err)
		writeJSON(w, map[string]interface{}{
			"error": err.Error(),
			"tool_calls": []ollama.ToolCall{
				{
					Tool:   "sleep",
					Args:   map[string]string{"seconds": "120"}, // Should be parsed by agent to include random jitter
					Reason: "Ollama decision failed, backing off for 120 seconds to prevent thundering herd",
				},
			},
		})
		return
	}

	log.Printf("[AI] Agent %s decision: %s (%d tool calls)", req.AgentID, resp.Reasoning, len(resp.ToolCalls))

	lastAction := "decide"
	lastTool := ""
	if len(resp.ToolCalls) > 0 {
		lastTool = resp.ToolCalls[0].Tool
		lastAction = resp.ToolCalls[0].Tool
	}
	h.recordActivity(AIActivityEntry{
		AgentID:       req.AgentID,
		LastDecideAt:  time.Now(),
		LastAction:    lastAction,
		LastTool:      lastTool,
		ToolCallCount: len(resp.ToolCalls),
		LastReasoning: truncateStr(resp.Reasoning, 120),
	})

	writeJSON(w, resp)
}

// ─── Report ───────────────────────────────────

func (h *AIHandler) handleReport(w http.ResponseWriter, r *http.Request) {
	body, err := io.ReadAll(r.Body)
	if err != nil {
		http.Error(w, "invalid report body", http.StatusBadRequest)
		return
	}

	var reports []ollama.Report
	if err := json.Unmarshal(body, &reports); err != nil {
		var single ollama.Report
		if err2 := json.Unmarshal(body, &single); err2 != nil {
			http.Error(w, "invalid report: "+err.Error(), http.StatusBadRequest)
			return
		}
		reports = []ollama.Report{single}
	}

	h.mu.Lock()
	for _, report := range reports {
		report.Timestamp = time.Now()
		h.reports = append(h.reports, report)
		log.Printf("[AI] Report from agent %s: tool=%s success=%v output=%s",
			report.AgentID, report.Tool, report.Success, truncateStr(report.Output, 200))

		prev := h.activity[report.AgentID]
		prev.AgentID = report.AgentID
		prev.LastReportAt = report.Timestamp
		prev.LastTool = report.Tool
		prev.LastAction = report.Tool
		prev.LastSuccess = report.Success
		h.activity[report.AgentID] = prev
		if h.onEvent != nil {
			h.onEvent(prev)
		}
	}
	// Keep only last 1000 reports
	if len(h.reports) > 1000 {
		h.reports = h.reports[len(h.reports)-1000:]
	}
	h.mu.Unlock()

	writeJSON(w, map[string]interface{}{
		"success":  true,
		"received": len(reports),
	})
}

// ─── Heartbeat ────────────────────────────────

type heartbeatRequest struct {
	AgentID string `json:"agent_id"`
	Status  string `json:"status"` // "alive", "restarting", "error"
	Message string `json:"message,omitempty"`
}

func (h *AIHandler) handleHeartbeat(w http.ResponseWriter, r *http.Request) {
	var req heartbeatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "invalid heartbeat", http.StatusBadRequest)
		return
	}

	if req.AgentID == "" {
		http.Error(w, "agent_id is required", http.StatusBadRequest)
		return
	}

	log.Printf("[AI] Heartbeat from agent %s: status=%s", req.AgentID, req.Status)

	writeJSON(w, map[string]interface{}{
		"success":  true,
		"interval": 60, // seconds until next heartbeat
	})
}

// ─── Helpers ──────────────────────────────────

func truncateStr(s string, maxLen int) string {
	if len(s) <= maxLen {
		return s
	}
	return s[:maxLen] + "..."
}

func (h *AIHandler) recordActivity(entry AIActivityEntry) {
	h.mu.Lock()
	prev := h.activity[entry.AgentID]
	if !entry.LastDecideAt.IsZero() {
		prev.LastDecideAt = entry.LastDecideAt
	}
	if entry.LastAction != "" {
		prev.LastAction = entry.LastAction
	}
	if entry.LastTool != "" {
		prev.LastTool = entry.LastTool
	}
	if entry.ToolCallCount > 0 {
		prev.ToolCallCount = entry.ToolCallCount
	}
	if entry.LastReasoning != "" {
		prev.LastReasoning = entry.LastReasoning
	}
	prev.AgentID = entry.AgentID
	h.activity[entry.AgentID] = prev
	fn := h.onEvent
	h.mu.Unlock()
	if fn != nil {
		fn(prev)
	}
}

func (h *AIHandler) ActivitySnapshot() []AIActivityEntry {
	h.mu.RLock()
	defer h.mu.RUnlock()
	out := make([]AIActivityEntry, 0, len(h.activity))
	for _, v := range h.activity {
		out = append(out, v)
	}
	return out
}