AetherForge/server/internal/api/ai_handler_test.go

package api

import (
	"bytes"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

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

// Numeric constants from ai_handler.go (guard against silent drift).
const (
	aiReportBufferCap          = 1000
	aiHeartbeatIntervalSeconds = 60
	aiReasoningTruncateLen     = 120
	aiReportOutputTruncateLen  = 200
	aiEngineIdleTTL            = time.Hour
	aiActivityIdleTTL          = 24 * time.Hour
	aiCleanupTickerInterval    = 10 * time.Minute
)

func newTestAIHandler(t *testing.T) *AIHandler {
	t.Helper()
	database, err := db.New(t.TempDir())
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { database.Close() })
	return NewAIHandler(database)
}

func mockOllamaChatServer(t *testing.T, content string, status int) *httptest.Server {
	t.Helper()
	if status == 0 {
		status = http.StatusOK
	}
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/chat" {
			http.NotFound(w, r)
			return
		}
		if r.Method != http.MethodPost {
			http.Error(w, "method", http.StatusMethodNotAllowed)
			return
		}
		w.WriteHeader(status)
		if status != http.StatusOK {
			_, _ = w.Write([]byte("ollama down"))
			return
		}
		payload := map[string]interface{}{
			"message": map[string]string{"content": content},
			"done":    true,
		}
		_ = json.NewEncoder(w).Encode(payload)
	}))
}

func ollamaDecideContent(reasoning string, tools []ollama.ToolCall) string {
	resp := ollama.DecideResponse{Reasoning: reasoning, ToolCalls: tools}
	b, _ := json.Marshal(resp)
	return string(b)
}

func TestAIConstants(t *testing.T) {
	if aiReportBufferCap != 1000 || aiHeartbeatIntervalSeconds != 60 ||
		aiReasoningTruncateLen != 120 || aiReportOutputTruncateLen != 200 ||
		aiEngineIdleTTL != time.Hour || aiActivityIdleTTL != 24*time.Hour ||
		aiCleanupTickerInterval != 10*time.Minute {
		t.Fatal("ai_handler constants drifted from documented values")
	}
}

func TestAINewAIHandler(t *testing.T) {
	h := newTestAIHandler(t)
	if h.db == nil || h.engines == nil || h.activity == nil {
		t.Fatal("NewAIHandler did not initialize fields")
	}
	if cap(h.reports) != aiReportBufferCap {
		t.Fatalf("reports slice cap want %d got %d", aiReportBufferCap, cap(h.reports))
	}
	if len(h.ActivitySnapshot()) != 0 {
		t.Fatal("expected empty activity snapshot")
	}
}

func TestAITruncateStr(t *testing.T) {
	if truncateStr("short", aiReasoningTruncateLen) != "short" {
		t.Fatal("short string should be unchanged")
	}
	long := strings.Repeat("x", aiReasoningTruncateLen+10)
	out := truncateStr(long, aiReasoningTruncateLen)
	if len(out) != aiReasoningTruncateLen+3 || !strings.HasSuffix(out, "...") {
		t.Fatalf("truncate want len %d with ellipsis, got %q", aiReasoningTruncateLen+3, out)
	}
	if truncateStr("", 5) != "" {
		t.Fatal("empty string")
	}
	if truncateStr("exact", 5) != "exact" {
		t.Fatal("exact length boundary")
	}
}

func TestAISetEngineForAgentAndGetRemove(t *testing.T) {
	h := newTestAIHandler(t)
	srv := mockOllamaChatServer(t, ollamaDecideContent("ok", nil), http.StatusOK)
	defer srv.Close()

	h.SetEngineForAgent("agent-a", srv.URL, "test-model")
	if h.GetEngine("agent-a") == nil {
		t.Fatal("GetEngine returned nil after SetEngineForAgent")
	}
	h.RemoveEngine("agent-a")
	if h.GetEngine("agent-a") != nil {
		t.Fatal("GetEngine should be nil after RemoveEngine")
	}
}

func TestAISetEngineForAgentDefaults(t *testing.T) {
	h := newTestAIHandler(t)
	var gotModel string
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req struct {
			Model string `json:"model"`
		}
		_ = json.NewDecoder(r.Body).Decode(&req)
		gotModel = req.Model
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(map[string]interface{}{
			"message": map[string]string{"content": ollamaDecideContent("ok", nil)},
			"done":    true,
		})
	}))
	defer srv.Close()

	h.SetEngineForAgent("agent-defaults", srv.URL, "")
	eng := h.GetEngine("agent-defaults")
	if eng == nil {
		t.Fatal("engine missing")
	}
	_, err := eng.Decide(&ollama.AgentState{AgentID: "agent-defaults"})
	if err != nil {
		t.Fatalf("Decide: %v", err)
	}
	if gotModel != "llama3.2" {
		t.Fatalf("default model want llama3.2 got %q", gotModel)
	}
}

func TestAIHandleDecideErrors(t *testing.T) {
	h := newTestAIHandler(t)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/decide", strings.NewReader("{"))
	w := httptest.NewRecorder()
	h.HandleDecide(w, req)
	if w.Code != http.StatusBadRequest {
		t.Fatalf("invalid JSON status %d", w.Code)
	}

	req = httptest.NewRequest(http.MethodPost, "/api/v1/agent/decide", strings.NewReader(`{"worker_name":"x"}`))
	w = httptest.NewRecorder()
	h.HandleDecide(w, req)
	if w.Code != http.StatusBadRequest || !strings.Contains(w.Body.String(), "agent_id") {
		t.Fatalf("missing agent_id status %d body %s", w.Code, w.Body.String())
	}
}

func TestAIHandleDecideSuccess(t *testing.T) {
	h := newTestAIHandler(t)
	content := ollamaDecideContent("all good", []ollama.ToolCall{{
		Tool: "check_miner", Args: map[string]string{"process_name": "xmrig"}, Reason: "verify",
	}})
	srv := mockOllamaChatServer(t, content, http.StatusOK)
	defer srv.Close()

	body, _ := json.Marshal(map[string]interface{}{
		"agent_id":        "agent-decide-ok",
		"ollama_endpoint": srv.URL,
		"model":           "test-model",
		"hostname":        "host1",
	})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/decide", bytes.NewReader(body))
	w := httptest.NewRecorder()
	h.HandleDecide(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("status %d body %s", w.Code, w.Body.String())
	}
	var resp ollama.DecideResponse
	if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
		t.Fatal(err)
	}
	if len(resp.ToolCalls) != 1 || resp.ToolCalls[0].Tool != "check_miner" {
		t.Fatalf("unexpected tool calls: %+v", resp.ToolCalls)
	}

	snap := h.ActivitySnapshot()
	var found bool
	for _, e := range snap {
		if e.AgentID == "agent-decide-ok" && e.LastAction == "check_miner" && e.ToolCallCount == 1 {
			found = true
			if len(e.LastReasoning) > aiReasoningTruncateLen+3 {
				t.Fatalf("reasoning should be truncated to %d", aiReasoningTruncateLen)
			}
		}
	}
	if !found {
		t.Fatalf("activity not recorded: %+v", snap)
	}
}

func TestAIHandleDecideOllamaFailureFallback(t *testing.T) {
	h := newTestAIHandler(t)
	srv := mockOllamaChatServer(t, "", http.StatusInternalServerError)
	defer srv.Close()

	body, _ := json.Marshal(map[string]interface{}{
		"agent_id":        "agent-decide-fail",
		"ollama_endpoint": srv.URL,
	})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/decide", bytes.NewReader(body))
	w := httptest.NewRecorder()
	h.HandleDecide(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("fallback should return 200, got %d", w.Code)
	}
	var out map[string]interface{}
	if err := json.NewDecoder(w.Body).Decode(&out); err != nil {
		t.Fatal(err)
	}
	if out["error"] == nil {
		t.Fatal("expected error field")
	}
	tools, ok := out["tool_calls"].([]interface{})
	if !ok || len(tools) != 1 {
		t.Fatalf("expected sleep fallback tool_calls, got %v", out["tool_calls"])
	}
	first, _ := tools[0].(map[string]interface{})
	if first["tool"] != "sleep" {
		t.Fatalf("expected sleep tool, got %v", first["tool"])
	}
	args, _ := first["args"].(map[string]interface{})
	if args["seconds"] != "120" {
		t.Fatalf("fallback sleep seconds want 120 got %v", args["seconds"])
	}
}

func TestAIHandleDecideCreatesEngineOnFirstRequest(t *testing.T) {
	h := newTestAIHandler(t)
	content := ollamaDecideContent("idle", nil)
	srv := mockOllamaChatServer(t, content, http.StatusOK)
	defer srv.Close()

	body, _ := json.Marshal(map[string]interface{}{
		"agent_id":        "agent-new",
		"ollama_endpoint": srv.URL,
		"model":           "m1",
	})
	req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/decide", bytes.NewReader(body))
	w := httptest.NewRecorder()
	h.HandleDecide(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("status %d %s", w.Code, w.Body.String())
	}
	if h.GetEngine("agent-new") == nil {
		t.Fatal("engine should exist after first decide")
	}
}

func TestAIHandleReportSingleAndArray(t *testing.T) {
	h := newTestAIHandler(t)

	single := ollama.Report{
		AgentID: "agent-1",
		Tool:    "sleep",
		Success: true,
		Output:  "ok",
	}
	body, _ := json.Marshal(single)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/report", bytes.NewReader(body))
	w := httptest.NewRecorder()
	h.HandleReport(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("single report status %d body %s", w.Code, w.Body.String())
	}

	arr := []ollama.Report{{
		AgentID: "agent-2",
		Tool:    "check_log",
		Success: false,
		Output:  "fail",
	}}
	body, _ = json.Marshal(arr)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/agent/report", bytes.NewReader(body))
	w = httptest.NewRecorder()
	h.HandleReport(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("array report status %d body %s", w.Code, w.Body.String())
	}

	activity := h.ActivitySnapshot()
	if len(activity) < 2 {
		t.Fatalf("expected activity entries, got %d", len(activity))
	}
}

func TestAIHandleReportErrors(t *testing.T) {
	h := newTestAIHandler(t)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/report", errReader{})
	w := httptest.NewRecorder()
	h.HandleReport(w, req)
	if w.Code != http.StatusBadRequest {
		t.Fatalf("read error status %d", w.Code)
	}

	req = httptest.NewRequest(http.MethodPost, "/api/v1/agent/report", strings.NewReader("not-json"))
	w = httptest.NewRecorder()
	h.HandleReport(w, req)
	if w.Code != http.StatusBadRequest {
		t.Fatalf("invalid json status %d", w.Code)
	}
}

func TestAIHandleReportBufferCap(t *testing.T) {
	h := newTestAIHandler(t)
	for i := 0; i < aiReportBufferCap+5; i++ {
		r := ollama.Report{AgentID: "cap-agent", Tool: "sleep", Success: true}
		b, _ := json.Marshal(r)
		req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/report", bytes.NewReader(b))
		w := httptest.NewRecorder()
		h.HandleReport(w, req)
		if w.Code != http.StatusOK {
			t.Fatalf("report %d status %d", i, w.Code)
		}
	}
	h.mu.RLock()
	n := len(h.reports)
	h.mu.RUnlock()
	if n != aiReportBufferCap {
		t.Fatalf("reports len want %d got %d", aiReportBufferCap, n)
	}
}

func TestAIHandleReportEventBroadcaster(t *testing.T) {
	h := newTestAIHandler(t)
	var events []AIActivityEntry
	h.SetEventBroadcaster(func(e AIActivityEntry) {
		events = append(events, e)
	})

	r := ollama.Report{AgentID: "evt-agent", Tool: "upload_log", Success: true, Output: "done"}
	b, _ := json.Marshal(r)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/report", bytes.NewReader(b))
	w := httptest.NewRecorder()
	h.HandleReport(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("status %d", w.Code)
	}
	if len(events) != 1 || events[0].AgentID != "evt-agent" || events[0].LastTool != "upload_log" {
		t.Fatalf("broadcaster events: %+v", events)
	}
	var resp map[string]interface{}
	_ = json.NewDecoder(bytes.NewReader(w.Body.Bytes())).Decode(&resp)
	if resp["received"].(float64) != 1 || resp["success"] != true {
		t.Fatalf("response: %v", resp)
	}
}

func TestAIHandleHeartbeat(t *testing.T) {
	h := newTestAIHandler(t)

	req := httptest.NewRequest(http.MethodPost, "/api/v1/agent/heartbeat", strings.NewReader("{"))
	w := httptest.NewRecorder()
	h.HandleHeartbeat(w, req)
	if w.Code != http.StatusBadRequest {
		t.Fatalf("invalid JSON status %d", w.Code)
	}

	req = httptest.NewRequest(http.MethodPost, "/api/v1/agent/heartbeat", strings.NewReader(`{"status":"alive"}`))
	w = httptest.NewRecorder()
	h.HandleHeartbeat(w, req)
	if w.Code != http.StatusBadRequest {
		t.Fatalf("missing agent_id status %d", w.Code)
	}

	body, _ := json.Marshal(heartbeatRequest{AgentID: "hb-agent", Status: "alive", Message: "ok"})
	req = httptest.NewRequest(http.MethodPost, "/api/v1/agent/heartbeat", bytes.NewReader(body))
	w = httptest.NewRecorder()
	h.HandleHeartbeat(w, req)
	if w.Code != http.StatusOK {
		t.Fatalf("status %d %s", w.Code, w.Body.String())
	}
	var out map[string]interface{}
	if err := json.NewDecoder(w.Body).Decode(&out); err != nil {
		t.Fatal(err)
	}
	if out["success"] != true {
		t.Fatalf("success: %v", out["success"])
	}
	interval, ok := out["interval"].(float64)
	if !ok || int(interval) != aiHeartbeatIntervalSeconds {
		t.Fatalf("interval want %d got %v", aiHeartbeatIntervalSeconds, out["interval"])
	}
}

func TestAIRecordActivityMergeAndBroadcast(t *testing.T) {
	h := newTestAIHandler(t)
	now := time.Now()
	h.mu.Lock()
	h.activity["merge-agent"] = AIActivityEntry{
		AgentID:       "merge-agent",
		LastReportAt:  now,
		LastSuccess:   true,
		ToolCallCount: 3,
	}
	h.mu.Unlock()

	var broadcasted AIActivityEntry
	h.SetEventBroadcaster(func(e AIActivityEntry) {
		broadcasted = e
	})

	h.recordActivity(AIActivityEntry{
		AgentID:       "merge-agent",
		LastDecideAt:  now.Add(time.Minute),
		LastAction:    "sleep",
		LastTool:      "sleep",
		LastReasoning: "rest",
		ToolCallCount: 0, // zero must not wipe previous count
	})

	h.mu.RLock()
	prev := h.activity["merge-agent"]
	h.mu.RUnlock()
	if prev.ToolCallCount != 3 {
		t.Fatalf("zero ToolCallCount should not clear prior count, got %d", prev.ToolCallCount)
	}
	if prev.LastAction != "sleep" || prev.LastReasoning != "rest" {
		t.Fatalf("merged activity: %+v", prev)
	}
	if broadcasted.AgentID != "merge-agent" {
		t.Fatalf("broadcast: %+v", broadcasted)
	}
}

func TestAIActivitySnapshot(t *testing.T) {
	h := newTestAIHandler(t)
	h.mu.Lock()
	h.activity["a"] = AIActivityEntry{AgentID: "a"}
	h.activity["b"] = AIActivityEntry{AgentID: "b"}
	h.mu.Unlock()
	snap := h.ActivitySnapshot()
	if len(snap) != 2 {
		t.Fatalf("want 2 entries got %d", len(snap))
	}
}

type errReader struct{}

func (errReader) Read([]byte) (int, error) { return 0, errors.New("read fail") }
func (errReader) Close() error             { return nil }

// errReader must satisfy io.ReadCloser for request body
var _ io.ReadCloser = errReader{}

func TestEstimateXMRPerDay(t *testing.T) {
	out := EstimateXMRPerDay(3_000_000_000)
	xmr, ok := out["xmr_per_day"].(float64)
	if !ok || xmr <= 0 {
		t.Fatalf("expected positive xmr estimate at network hashrate, got %v", out)
	}
}