osmedeus/internal/executor/agent_executor_test.go

package executor

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"testing"

	"github.com/j3ssie/osmedeus/v5/internal/config"
	"github.com/j3ssie/osmedeus/v5/internal/core"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// mockLLMResponse creates a ChatCompletionResponse with text content and no tool calls
func mockLLMResponse(content string) ChatCompletionResponse {
	return ChatCompletionResponse{
		ID:    "test-id",
		Model: "test-model",
		Choices: []ChatChoice{
			{
				Index: 0,
				Message: ChatMessage{
					Role:    "assistant",
					Content: content,
				},
				FinishReason: "stop",
			},
		},
		Usage: ChatUsage{
			PromptTokens:     10,
			CompletionTokens: 20,
			TotalTokens:      30,
		},
	}
}

// mockLLMToolCallResponse creates a response with tool calls
func mockLLMToolCallResponse(toolCalls []core.LLMToolCall) ChatCompletionResponse {
	return ChatCompletionResponse{
		ID:    "test-id",
		Model: "test-model",
		Choices: []ChatChoice{
			{
				Index: 0,
				Message: ChatMessage{
					Role:      "assistant",
					Content:   "",
					ToolCalls: toolCalls,
				},
				FinishReason: "tool_calls",
			},
		},
		Usage: ChatUsage{
			PromptTokens:     10,
			CompletionTokens: 20,
			TotalTokens:      30,
		},
	}
}

// newMockLLMServer creates a mock OpenAI-compatible API server.
// handler is called for each request and should write the response.
func newMockLLMServer(handler http.HandlerFunc) *httptest.Server {
	return httptest.NewServer(handler)
}

// newMockConfig creates a config with a mock LLM provider pointing to the given URL
func newMockConfig(t *testing.T, serverURL string) *config.Config {
	t.Helper()
	baseDir := t.TempDir()
	return &config.Config{
		BaseFolder:     baseDir,
		WorkspacesPath: filepath.Join(baseDir, "workspaces"),
		LLM: config.LLMConfig{
			LLMProviders: []config.LLMProvider{
				{
					Provider:  "mock",
					BaseURL:   serverURL,
					AuthToken: "test-token",
					Model:     "test-model",
				},
			},
			MaxTokens:  1000,
			MaxRetries: 1,
			Timeout:    "30s",
		},
	}
}

// ============================================================================
// Agent Executor Unit Tests
// ============================================================================

func TestAgentExecutor_Name(t *testing.T) {
	executor := NewAgentExecutor(nil, nil)
	assert.Equal(t, "agent", executor.Name())
}

func TestAgentExecutor_StepTypes(t *testing.T) {
	executor := NewAgentExecutor(nil, nil)
	types := executor.StepTypes()
	assert.Len(t, types, 1)
	assert.Equal(t, core.StepTypeAgent, types[0])
}

func TestAgentExecutor_ValidationErrors(t *testing.T) {
	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	t.Run("missing config", func(t *testing.T) {
		executor := NewAgentExecutor(nil, nil)
		step := &core.Step{Name: "test", Type: core.StepTypeAgent}

		result, err := executor.Execute(ctx, step, execCtx)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "config not set")
		assert.Equal(t, core.StepStatusFailed, result.Status)
	})

	t.Run("missing query", func(t *testing.T) {
		server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {})
		defer server.Close()
		cfg := newMockConfig(t, server.URL)

		executor := NewAgentExecutor(nil, nil)
		executor.SetConfig(cfg)

		step := &core.Step{
			Name:          "test",
			Type:          core.StepTypeAgent,
			MaxIterations: 5,
			AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
		}

		result, err := executor.Execute(ctx, step, execCtx)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "requires 'query'")
		assert.Equal(t, core.StepStatusFailed, result.Status)
	})

	t.Run("missing max_iterations", func(t *testing.T) {
		server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {})
		defer server.Close()
		cfg := newMockConfig(t, server.URL)

		executor := NewAgentExecutor(nil, nil)
		executor.SetConfig(cfg)

		step := &core.Step{
			Name:       "test",
			Type:       core.StepTypeAgent,
			Query:      "test query",
			AgentTools: []core.AgentToolDef{{Preset: "bash"}},
		}

		result, err := executor.Execute(ctx, step, execCtx)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "requires 'max_iterations'")
		assert.Equal(t, core.StepStatusFailed, result.Status)
	})

	t.Run("missing agent_tools", func(t *testing.T) {
		server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {})
		defer server.Close()
		cfg := newMockConfig(t, server.URL)

		executor := NewAgentExecutor(nil, nil)
		executor.SetConfig(cfg)

		step := &core.Step{
			Name:          "test",
			Type:          core.StepTypeAgent,
			Query:         "test query",
			MaxIterations: 5,
		}

		result, err := executor.Execute(ctx, step, execCtx)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "requires 'agent_tools'")
		assert.Equal(t, core.StepStatusFailed, result.Status)
	})

	t.Run("unknown preset tool", func(t *testing.T) {
		server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {})
		defer server.Close()
		cfg := newMockConfig(t, server.URL)

		executor := NewAgentExecutor(nil, nil)
		executor.SetConfig(cfg)

		step := &core.Step{
			Name:          "test",
			Type:          core.StepTypeAgent,
			Query:         "test query",
			MaxIterations: 5,
			AgentTools:    []core.AgentToolDef{{Preset: "nonexistent_tool"}},
		}

		result, err := executor.Execute(ctx, step, execCtx)
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "unknown preset tool")
		assert.Equal(t, core.StepStatusFailed, result.Status)
	})
}

func TestAgentExecutor_SimpleCompletion(t *testing.T) {
	// Mock server that responds without tool calls (agent should complete in 1 iteration)
	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		resp := mockLLMResponse("The analysis is complete.")
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(resp)
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "simple-agent",
		Type:          core.StepTypeAgent,
		Query:         "Analyze the target",
		MaxIterations: 5,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)

	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Equal(t, "The analysis is complete.", result.Output)
	assert.Equal(t, "The analysis is complete.", result.Exports["agent_content"])
	assert.Equal(t, 1, result.Exports["agent_iterations"])
	assert.Equal(t, 30, result.Exports["agent_total_tokens"])
	assert.NotEmpty(t, result.Exports["agent_history"])
}

func TestAgentExecutor_ToolCallLoop(t *testing.T) {
	// Mock server that:
	// 1st call: returns a tool call for exec_cmd
	// 2nd call: returns final text response
	var callCount int32

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		if count == 1 {
			// First call: return tool call
			resp := mockLLMToolCallResponse([]core.LLMToolCall{
				{
					ID:   "call_1",
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "bash",
						Arguments: `{"command": "echo hello"}`,
					},
				},
			})
			_ = json.NewEncoder(w).Encode(resp)
		} else {
			// Second call: return final response
			resp := mockLLMResponse("Command output was: hello")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "tool-call-agent",
		Type:          core.StepTypeAgent,
		Query:         "Run echo hello",
		MaxIterations: 5,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)

	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Equal(t, "Command output was: hello", result.Output)
	assert.Equal(t, 2, result.Exports["agent_iterations"])
	// 2 calls * 30 tokens each
	assert.Equal(t, 60, result.Exports["agent_total_tokens"])
}

func TestAgentExecutor_MaxIterationsLimit(t *testing.T) {
	// Mock server that always returns tool calls (never finishes)
	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		resp := mockLLMToolCallResponse([]core.LLMToolCall{
			{
				ID:   "call_1",
				Type: "function",
				Function: core.LLMToolCallFunction{
					Name:      "bash",
					Arguments: `{"command": "echo loop"}`,
				},
			},
		})
		_ = json.NewEncoder(w).Encode(resp)
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "loop-agent",
		Type:          core.StepTypeAgent,
		Query:         "Keep running",
		MaxIterations: 3,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)

	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Equal(t, 3, result.Exports["agent_iterations"])
}

func TestAgentExecutor_StopCondition(t *testing.T) {
	// Mock server: first call returns "checking...", second returns "DONE"
	var callCount int32

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		if count == 1 {
			// Return tool call first
			resp := mockLLMToolCallResponse([]core.LLMToolCall{
				{
					ID:   "call_1",
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "bash",
						Arguments: `{"command": "echo check"}`,
					},
				},
			})
			_ = json.NewEncoder(w).Encode(resp)
		} else {
			// Return text with DONE keyword
			resp := mockLLMResponse("Analysis DONE successfully")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "stop-agent",
		Type:          core.StepTypeAgent,
		Query:         "Analyze and say DONE when finished",
		MaxIterations: 10,
		StopCondition: `contains(agent_content, "DONE")`,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)

	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Contains(t, result.Output, "DONE")
	// Should have stopped after 2 iterations (not 10)
	assert.Equal(t, 2, result.Exports["agent_iterations"])
}

func TestAgentExecutor_SystemPrompt(t *testing.T) {
	// Verify that system prompt is included in the request
	var receivedMessages []ChatMessage

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var req ChatCompletionRequest
		_ = json.Unmarshal(body, &req)
		receivedMessages = req.Messages

		w.Header().Set("Content-Type", "application/json")
		resp := mockLLMResponse("Response")
		_ = json.NewEncoder(w).Encode(resp)
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "system-prompt-agent",
		Type:          core.StepTypeAgent,
		SystemPrompt:  "You are a security analyst.",
		Query:         "Analyze target",
		MaxIterations: 3,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	_, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)

	// Verify messages contain system prompt and user query
	require.Len(t, receivedMessages, 2)
	assert.Equal(t, "system", receivedMessages[0].Role)
	assert.Equal(t, "You are a security analyst.", receivedMessages[0].Content)
	assert.Equal(t, "user", receivedMessages[1].Role)
	assert.Equal(t, "Analyze target", receivedMessages[1].Content)
}

func TestAgentExecutor_ToolsInRequest(t *testing.T) {
	// Verify that tools are included in the LLM request
	var receivedTools []core.LLMTool

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var req ChatCompletionRequest
		_ = json.Unmarshal(body, &req)
		receivedTools = req.Tools

		w.Header().Set("Content-Type", "application/json")
		resp := mockLLMResponse("Done")
		_ = json.NewEncoder(w).Encode(resp)
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "tools-check",
		Type:          core.StepTypeAgent,
		Query:         "Test",
		MaxIterations: 3,
		AgentTools: []core.AgentToolDef{
			{Preset: "bash"},
			{Preset: "read_file"},
			{
				Name:        "custom",
				Description: "Custom tool",
				Parameters:  map[string]interface{}{"type": "object"},
			},
		},
	}

	_, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)

	// Verify 3 tools were sent to the LLM
	assert.Len(t, receivedTools, 3)
	assert.Equal(t, "bash", receivedTools[0].Function.Name)
	assert.Equal(t, "read_file", receivedTools[1].Function.Name)
	assert.Equal(t, "custom", receivedTools[2].Function.Name)
}

func TestAgentExecutor_MemoryPersist(t *testing.T) {
	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		resp := mockLLMResponse("Persisted response")
		_ = json.NewEncoder(w).Encode(resp)
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	persistPath := filepath.Join(t.TempDir(), "agent", "conversation.json")

	step := &core.Step{
		Name:          "persist-agent",
		Type:          core.StepTypeAgent,
		Query:         "Test persist",
		MaxIterations: 3,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
		Memory: &core.AgentMemoryConfig{
			PersistPath: persistPath,
		},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	assert.Equal(t, core.StepStatusSuccess, result.Status)

	// Verify conversation was persisted
	data, err := os.ReadFile(persistPath)
	require.NoError(t, err)

	var messages []ChatMessage
	err = json.Unmarshal(data, &messages)
	require.NoError(t, err)
	assert.GreaterOrEqual(t, len(messages), 2) // at least user + assistant
}

func TestAgentExecutor_MemoryResume(t *testing.T) {
	// Create a conversation file to resume from
	tmpDir := t.TempDir()
	resumePath := filepath.Join(tmpDir, "resume.json")

	priorMessages := []ChatMessage{
		{Role: "system", Content: "You are a helper."},
		{Role: "user", Content: "Previous question"},
		{Role: "assistant", Content: "Previous answer"},
	}
	data, _ := json.MarshalIndent(priorMessages, "", "  ")
	require.NoError(t, os.WriteFile(resumePath, data, 0o644))

	var receivedMessages []ChatMessage

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var req ChatCompletionRequest
		_ = json.Unmarshal(body, &req)
		receivedMessages = req.Messages

		w.Header().Set("Content-Type", "application/json")
		resp := mockLLMResponse("Resumed response")
		_ = json.NewEncoder(w).Encode(resp)
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "resume-agent",
		Type:          core.StepTypeAgent,
		Query:         "Follow-up question",
		MaxIterations: 3,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
		Memory: &core.AgentMemoryConfig{
			ResumePath: resumePath,
		},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	assert.Equal(t, core.StepStatusSuccess, result.Status)

	// Verify resumed messages + new query
	// Should be: system + previous user + previous assistant + new user query
	assert.Len(t, receivedMessages, 4)
	assert.Equal(t, "system", receivedMessages[0].Role)
	assert.Equal(t, "user", receivedMessages[1].Role)
	assert.Equal(t, "Previous question", receivedMessages[1].Content)
	assert.Equal(t, "assistant", receivedMessages[2].Role)
	assert.Equal(t, "user", receivedMessages[3].Role)
	assert.Equal(t, "Follow-up question", receivedMessages[3].Content)
}

func TestAgentExecutor_CustomToolHandler(t *testing.T) {
	// First call returns a tool call for "custom_tool", second returns final answer
	var callCount int32

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		if count == 1 {
			resp := mockLLMToolCallResponse([]core.LLMToolCall{
				{
					ID:   "call_1",
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "greet",
						Arguments: `{"name": "World"}`,
					},
				},
			})
			_ = json.NewEncoder(w).Encode(resp)
		} else {
			resp := mockLLMResponse("Greeting sent")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "custom-handler-agent",
		Type:          core.StepTypeAgent,
		Query:         "Greet World",
		MaxIterations: 5,
		AgentTools: []core.AgentToolDef{
			{
				Name:        "greet",
				Description: "Greet someone",
				Parameters: map[string]interface{}{
					"type": "object",
					"properties": map[string]interface{}{
						"name": map[string]interface{}{"type": "string"},
					},
				},
				Handler: `"Hello, " + args.name + "!"`,
			},
		},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Equal(t, "Greeting sent", result.Output)
}

func TestAgentExecutor_ContextCancellation(t *testing.T) {
	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		// This should not be reached because context is already cancelled
		w.Header().Set("Content-Type", "application/json")
		resp := mockLLMResponse("Should not reach")
		_ = json.NewEncoder(w).Encode(resp)
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx, cancel := context.WithCancel(context.Background())
	cancel() // Cancel immediately

	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "cancel-agent",
		Type:          core.StepTypeAgent,
		Query:         "This should be cancelled",
		MaxIterations: 5,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	assert.Error(t, err)
	assert.Equal(t, core.StepStatusFailed, result.Status)
}

func TestAgentExecutor_MultipleToolCalls(t *testing.T) {
	// Mock server: first call returns 2 tool calls, second returns final answer
	var callCount int32

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		if count == 1 {
			resp := mockLLMToolCallResponse([]core.LLMToolCall{
				{
					ID:   "call_1",
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "bash",
						Arguments: `{"command": "echo first"}`,
					},
				},
				{
					ID:   "call_2",
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "bash",
						Arguments: `{"command": "echo second"}`,
					},
				},
			})
			_ = json.NewEncoder(w).Encode(resp)
		} else {
			resp := mockLLMResponse("Both commands executed")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "multi-tool-agent",
		Type:          core.StepTypeAgent,
		Query:         "Run two commands",
		MaxIterations: 5,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Equal(t, "Both commands executed", result.Output)
	assert.Equal(t, 2, result.Exports["agent_iterations"])
}

// ============================================================================
// Helper Function Tests
// ============================================================================

func TestBuildPresetCallExpr(t *testing.T) {
	tests := []struct {
		name     string
		funcName string
		args     map[string]interface{}
		expected string
	}{
		{
			name:     "bash",
			funcName: "bash",
			args:     map[string]interface{}{"command": "echo hello"},
			expected: `bash("echo hello")`,
		},
		{
			name:     "read_file",
			funcName: "read_file",
			args:     map[string]interface{}{"path": "/tmp/test.txt"},
			expected: `read_file("/tmp/test.txt")`,
		},
		{
			name:     "save_content",
			funcName: "save_content",
			args:     map[string]interface{}{"content": "hello", "path": "/tmp/out.txt"},
			expected: `save_content("hello", "/tmp/out.txt")`,
		},
		{
			name:     "grep_string",
			funcName: "grep_string",
			args:     map[string]interface{}{"source": "/tmp/file.txt", "str": "pattern"},
			expected: `grep_string("/tmp/file.txt", "pattern")`,
		},
		{
			name:     "jq",
			funcName: "jq",
			args:     map[string]interface{}{"json_data": `{"key":"value"}`, "expression": ".key"},
			expected: `jq("{\"key\":\"value\"}", ".key")`,
		},
		{
			name:     "exec_python",
			funcName: "exec_python",
			args:     map[string]interface{}{"code": "print('hello')"},
			expected: `exec_python("print('hello')")`,
		},
		{
			name:     "exec_python_file",
			funcName: "exec_python_file",
			args:     map[string]interface{}{"path": "/tmp/script.py"},
			expected: `exec_python_file("/tmp/script.py")`,
		},
		{
			name:     "run_module",
			funcName: "run_module",
			args:     map[string]interface{}{"module": "subdomain", "target": "example.com", "params": "threads=10"},
			expected: `run_module("subdomain", "example.com", "threads=10")`,
		},
		{
			name:     "run_flow",
			funcName: "run_flow",
			args:     map[string]interface{}{"flow": "general", "target": "example.com", "params": ""},
			expected: `run_flow("general", "example.com", "")`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := buildPresetCallExpr(tt.funcName, tt.args)
			assert.Equal(t, tt.expected, result)
		})
	}
}

func TestJsQuote(t *testing.T) {
	tests := []struct {
		input    string
		expected string
	}{
		{`hello`, `"hello"`},
		{`say "hi"`, `"say \"hi\""`},
		{"line1\nline2", `"line1\nline2"`},
		{`path\to\file`, `"path\\to\\file"`},
		{"tab\there", `"tab\there"`},
	}

	for _, tt := range tests {
		t.Run(tt.input, func(t *testing.T) {
			assert.Equal(t, tt.expected, jsQuote(tt.input))
		})
	}
}

func TestFormatToolResult(t *testing.T) {
	assert.Equal(t, "hello", formatToolResult("hello"))
	assert.Equal(t, "true", formatToolResult(true))
	assert.Equal(t, "false", formatToolResult(false))
	assert.Equal(t, "42", formatToolResult(42))
	assert.Equal(t, "3.14", formatToolResult(3.14))
	assert.Equal(t, "", formatToolResult(nil))

	// Complex type gets JSON serialized
	result := formatToolResult(map[string]string{"key": "value"})
	assert.Contains(t, result, "key")
	assert.Contains(t, result, "value")
}

func TestGetStringArg(t *testing.T) {
	args := map[string]interface{}{
		"str_arg": "hello",
		"int_arg": 42,
		"nil_arg": nil,
	}

	assert.Equal(t, "hello", getStringArg(args, "str_arg"))
	assert.Equal(t, "42", getStringArg(args, "int_arg"))
	assert.Equal(t, "", getStringArg(args, "missing"))
}

// ============================================================================
// Integration Test: Full Agent Loop via Executor
// ============================================================================

func TestExecutor_AgentStep_FullIntegration(t *testing.T) {
	// Mock LLM server that simulates a tool call loop
	var callCount int32
	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		if count == 1 {
			// First: ask to run a command
			resp := mockLLMToolCallResponse([]core.LLMToolCall{
				{
					ID:   "call_1",
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "bash",
						Arguments: `{"command": "echo integration_test"}`,
					},
				},
			})
			_ = json.NewEncoder(w).Encode(resp)
		} else {
			resp := mockLLMResponse("Integration test passed")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := testConfig(t)
	cfg.LLM = config.LLMConfig{
		LLMProviders: []config.LLMProvider{
			{
				Provider: "mock",
				BaseURL:  server.URL,
				Model:    "test-model",
			},
		},
		MaxTokens:  1000,
		MaxRetries: 1,
		Timeout:    "30s",
	}

	module := &core.Workflow{
		Name: "test-agent-integration",
		Kind: core.KindModule,
		Steps: []core.Step{
			{
				Name:          "agent-step",
				Type:          core.StepTypeAgent,
				Query:         "Run echo integration_test",
				MaxIterations: 5,
				AgentTools: []core.AgentToolDef{
					{Preset: "bash"},
				},
				Exports: map[string]string{
					"result": "{{agent_content}}",
				},
			},
			{
				Name:    "verify-step",
				Type:    core.StepTypeBash,
				Command: "echo 'agent result: {{result}}'",
			},
		},
	}

	executor := NewExecutor()
	executor.SetDryRun(false)
	executor.SetSpinner(false)

	ctx := context.Background()
	result, err := executor.ExecuteModule(ctx, module, map[string]string{
		"target": "test",
	}, cfg)

	require.NoError(t, err)
	assert.Equal(t, core.RunStatusCompleted, result.Status)
	assert.Len(t, result.Steps, 2)
	assert.Equal(t, core.StepStatusSuccess, result.Steps[0].Status)
	assert.Equal(t, "Integration test passed", result.Steps[0].Output)
	assert.Equal(t, core.StepStatusSuccess, result.Steps[1].Status)
}

// ============================================================================
// Group 2: Planning Stage Tests
// ============================================================================

func TestAgentExecutor_PlanningStage(t *testing.T) {
	// Mock server: 1st call = planning response (no tools), 2nd call = final response
	var callCount int32
	var receivedMessages [][]ChatMessage

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		body, _ := io.ReadAll(r.Body)
		var req ChatCompletionRequest
		_ = json.Unmarshal(body, &req)
		receivedMessages = append(receivedMessages, req.Messages)

		if count == 1 {
			// Planning phase: return plan text
			resp := mockLLMResponse("Step 1: Scan ports\nStep 2: Check services\nStep 3: Report findings")
			_ = json.NewEncoder(w).Encode(resp)
		} else {
			// Execution phase: return final response
			resp := mockLLMResponse("Plan executed successfully")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "plan-agent",
		Type:          core.StepTypeAgent,
		PlanPrompt:    "Create a plan for scanning test.com",
		Query:         "Execute the plan",
		MaxIterations: 5,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)

	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Equal(t, "Plan executed successfully", result.Output)

	// Verify agent_plan export is populated
	planContent, ok := result.Exports["agent_plan"]
	assert.True(t, ok)
	assert.Contains(t, planContent, "Step 1")

	// Verify planning request had the plan prompt
	require.Len(t, receivedMessages, 2)
	planMsgs := receivedMessages[0]
	assert.Equal(t, "user", planMsgs[len(planMsgs)-1].Role)
	assert.Equal(t, "Create a plan for scanning test.com", planMsgs[len(planMsgs)-1].Content)

	// Verify execution request includes the plan as assistant message
	execMsgs := receivedMessages[1]
	foundPlan := false
	for _, msg := range execMsgs {
		if msg.Role == "assistant" {
			if content, ok := msg.Content.(string); ok && strings.Contains(content, "Step 1") {
				foundPlan = true
			}
		}
	}
	assert.True(t, foundPlan, "execution messages should include the plan as assistant message")

	// Verify tokens are tracked from both planning and execution
	assert.Equal(t, 60, result.Exports["agent_total_tokens"]) // 30 from planning + 30 from execution
}

func TestAgentExecutor_PlanningStage_WithMaxTokens(t *testing.T) {
	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")

		body, _ := io.ReadAll(r.Body)
		var req ChatCompletionRequest
		_ = json.Unmarshal(body, &req)

		// Return plan or final based on whether tools are present
		if len(req.Tools) == 0 {
			// Planning request — verify max_tokens
			assert.Equal(t, 500, req.MaxTokens)
			resp := mockLLMResponse("Short plan")
			_ = json.NewEncoder(w).Encode(resp)
		} else {
			resp := mockLLMResponse("Done")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	planMaxTokens := 500
	step := &core.Step{
		Name:          "plan-tokens-agent",
		Type:          core.StepTypeAgent,
		PlanPrompt:    "Create a plan",
		PlanMaxTokens: &planMaxTokens,
		Query:         "Execute",
		MaxIterations: 3,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	assert.Equal(t, core.StepStatusSuccess, result.Status)
}

// ============================================================================
// Group 3: Multi-Goal Execution Tests
// ============================================================================

func TestAgentExecutor_MultiGoal(t *testing.T) {
	// Mock server: responds to each query with a unique response
	var callCount int32

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		resp := mockLLMResponse(fmt.Sprintf("Goal %d complete", count))
		_ = json.NewEncoder(w).Encode(resp)
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "multi-goal-agent",
		Type:          core.StepTypeAgent,
		Queries:       []string{"List files", "Summarize findings", "Report"},
		MaxIterations: 3,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)

	assert.Equal(t, core.StepStatusSuccess, result.Status)
	// Final content should be from the last goal
	assert.Equal(t, "Goal 3 complete", result.Output)

	// Verify agent_goal_results export
	goalResultsJSON, ok := result.Exports["agent_goal_results"]
	assert.True(t, ok, "agent_goal_results should be exported for multi-goal")
	assert.Contains(t, goalResultsJSON, "List files")
	assert.Contains(t, goalResultsJSON, "Summarize findings")
	assert.Contains(t, goalResultsJSON, "Report")

	// Tokens should be accumulated from all goals
	assert.Equal(t, 90, result.Exports["agent_total_tokens"]) // 30 * 3 goals
}

func TestAgentExecutor_QueryAndQueriesMutualExclusion(t *testing.T) {
	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {})
	defer server.Close()
	cfg := newMockConfig(t, server.URL)

	executor := NewAgentExecutor(nil, nil)
	executor.SetConfig(cfg)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "both-query",
		Type:          core.StepTypeAgent,
		Query:         "single query",
		Queries:       []string{"query 1", "query 2"},
		MaxIterations: 3,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "cannot have both")
	assert.Equal(t, core.StepStatusFailed, result.Status)
}

func TestAgentExecutor_MultiGoal_SharedConversation(t *testing.T) {
	// Verify subsequent goals see messages from previous goals
	var receivedMessages [][]ChatMessage

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")

		body, _ := io.ReadAll(r.Body)
		var req ChatCompletionRequest
		_ = json.Unmarshal(body, &req)
		receivedMessages = append(receivedMessages, req.Messages)

		resp := mockLLMResponse("response")
		_ = json.NewEncoder(w).Encode(resp)
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "shared-conv",
		Type:          core.StepTypeAgent,
		Queries:       []string{"First query", "Second query"},
		MaxIterations: 3,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	assert.Equal(t, core.StepStatusSuccess, result.Status)

	// Second goal should have more messages (accumulated from first goal)
	require.Len(t, receivedMessages, 2)
	assert.True(t, len(receivedMessages[1]) > len(receivedMessages[0]),
		"second goal should see accumulated messages from first goal")
}

// ============================================================================
// Group 5: Model Fallback Tests
// ============================================================================

func TestAgentExecutor_ModelFallback(t *testing.T) {
	var receivedModels []string

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")

		body, _ := io.ReadAll(r.Body)
		var req ChatCompletionRequest
		_ = json.Unmarshal(body, &req)
		receivedModels = append(receivedModels, req.Model)

		// Fail for first model, succeed for second
		if req.Model == "bad-model" {
			w.WriteHeader(http.StatusInternalServerError)
			_ = json.NewEncoder(w).Encode(map[string]interface{}{
				"error": map[string]interface{}{
					"message": "model not found",
					"type":    "invalid_request_error",
				},
			})
			return
		}

		resp := mockLLMResponse("Fallback succeeded")
		_ = json.NewEncoder(w).Encode(resp)
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "fallback-agent",
		Type:          core.StepTypeAgent,
		Query:         "Test fallback",
		Models:        []string{"bad-model", "test-model"},
		MaxIterations: 3,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Equal(t, "Fallback succeeded", result.Output)

	// Verify both models were attempted
	assert.Contains(t, receivedModels, "bad-model")
	assert.Contains(t, receivedModels, "test-model")
}

// ============================================================================
// Group 6: Structured Output Tests
// ============================================================================

func TestAgentExecutor_StructuredOutput(t *testing.T) {
	var callCount int32

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		body, _ := io.ReadAll(r.Body)
		var req ChatCompletionRequest
		_ = json.Unmarshal(body, &req)

		if count == 1 {
			// First call: agent completes with unstructured text
			resp := mockLLMResponse("Found some results")
			_ = json.NewEncoder(w).Encode(resp)
		} else {
			// Second call: structured output request — verify response_format is set
			assert.NotNil(t, req.ResponseFormat, "structured output call should have response_format")
			resp := mockLLMResponse(`{"subdomains":[{"name":"sub.test.com","status":"active"}]}`)
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "structured-agent",
		Type:          core.StepTypeAgent,
		Query:         "Find subdomains",
		MaxIterations: 5,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
		OutputSchema:  `{"type":"object","properties":{"subdomains":{"type":"array"}}}`,
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	assert.Equal(t, core.StepStatusSuccess, result.Status)
	// Final content should be the structured JSON
	assert.Contains(t, result.Output, "subdomains")
	assert.Contains(t, result.Output, "sub.test.com")
}

func TestAgentExecutor_StructuredOutput_AlreadyJSON(t *testing.T) {
	// If the agent already returns valid JSON, no extra call should be made
	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		resp := mockLLMResponse(`{"results":["a","b"]}`)
		_ = json.NewEncoder(w).Encode(resp)
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "already-json-agent",
		Type:          core.StepTypeAgent,
		Query:         "Return JSON",
		MaxIterations: 3,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
		OutputSchema:  `{"type":"object"}`,
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	assert.Equal(t, core.StepStatusSuccess, result.Status)
	// Should only have 1 LLM call (no extra structured output request)
	assert.Equal(t, 30, result.Exports["agent_total_tokens"])
}

// ============================================================================
// Group 7: Tool Tracing Hooks Tests
// ============================================================================

func TestAgentExecutor_TracingHooks(t *testing.T) {
	var callCount int32

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		if count == 1 {
			resp := mockLLMToolCallResponse([]core.LLMToolCall{
				{
					ID:   "call_1",
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "bash",
						Arguments: `{"command": "echo traced"}`,
					},
				},
			})
			_ = json.NewEncoder(w).Encode(resp)
		} else {
			resp := mockLLMResponse("Traced complete")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	// Use hooks that call log_info (which is a registered function)
	step := &core.Step{
		Name:          "traced-agent",
		Type:          core.StepTypeAgent,
		Query:         "Run traced command",
		MaxIterations: 5,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
		OnToolStart:   `log_info("HOOK_START: " + tool_name)`,
		OnToolEnd:     `log_info("HOOK_END: " + tool_name + " dur=" + duration)`,
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Equal(t, "Traced complete", result.Output)
}

func TestAgentExecutor_TracingHooks_EmptyHooks(t *testing.T) {
	// Verify hooks are no-ops when empty strings
	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		resp := mockLLMResponse("No hooks")
		_ = json.NewEncoder(w).Encode(resp)
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "no-hooks-agent",
		Type:          core.StepTypeAgent,
		Query:         "Test without hooks",
		MaxIterations: 3,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
		OnToolStart:   "",
		OnToolEnd:     "",
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	assert.Equal(t, core.StepStatusSuccess, result.Status)
}

// ============================================================================
// Sub-Agent Orchestration Tests
// ============================================================================

func TestAgentExecutor_SubAgentSpawn(t *testing.T) {
	// Mock server:
	// 1st call (parent): returns spawn_agent tool call
	// 2nd call (child): returns child final response
	// 3rd call (parent): returns parent final response using child result
	var callCount int32

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		switch count {
		case 1:
			// Parent: spawn recon_agent
			resp := mockLLMToolCallResponse([]core.LLMToolCall{
				{
					ID:   "call_spawn",
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "spawn_agent",
						Arguments: `{"agent":"recon_agent","query":"scan ports on target"}`,
					},
				},
			})
			_ = json.NewEncoder(w).Encode(resp)
		case 2:
			// Child: return final result
			resp := mockLLMResponse("Found open ports: 80, 443, 8080")
			_ = json.NewEncoder(w).Encode(resp)
		case 3:
			// Parent: final response using child result
			resp := mockLLMResponse("Recon complete. Open ports: 80, 443, 8080")
			_ = json.NewEncoder(w).Encode(resp)
		default:
			resp := mockLLMResponse("unexpected call")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "orchestrator",
		Type:          core.StepTypeAgent,
		Query:         "Analyze target by coordinating specialists",
		SystemPrompt:  "You are an orchestrator. Delegate tasks to sub-agents.",
		MaxIterations: 5,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
		SubAgents: []core.SubAgentDef{
			{
				Name:          "recon_agent",
				Description:   "Specialized agent for reconnaissance",
				SystemPrompt:  "You are a recon specialist",
				MaxIterations: 5,
				AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
			},
		},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)

	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Contains(t, result.Output, "Recon complete")
	// 3 LLM calls × 30 tokens each = 90
	assert.Equal(t, 90, result.Exports["agent_total_tokens"])
}

func TestAgentExecutor_SubAgentDepthLimit(t *testing.T) {
	// Mock server: parent calls spawn_agent which tries to exceed depth
	var callCount int32

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		if count == 1 {
			// Parent: spawn sub-agent
			resp := mockLLMToolCallResponse([]core.LLMToolCall{
				{
					ID:   "call_spawn",
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "spawn_agent",
						Arguments: `{"agent":"child","query":"do work"}`,
					},
				},
			})
			_ = json.NewEncoder(w).Encode(resp)
		} else {
			resp := mockLLMResponse("Done with depth error")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "depth-limit-parent",
		Type:          core.StepTypeAgent,
		Query:         "Delegate work",
		MaxIterations: 5,
		MaxAgentDepth: 1, // Only 1 level allowed, but parent is at 0 so child at 1 is OK
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
		SubAgents: []core.SubAgentDef{
			{
				Name:          "child",
				Description:   "Child agent",
				MaxIterations: 3,
				AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
				SubAgents: []core.SubAgentDef{
					{
						Name:          "grandchild",
						Description:   "Grandchild agent",
						MaxIterations: 3,
						AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
					},
				},
			},
		},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	// Parent should succeed even though the child depth limit was hit
	// The error is returned as a tool result to the parent LLM
	assert.Equal(t, core.StepStatusSuccess, result.Status)
}

func TestAgentExecutor_SubAgentFailure(t *testing.T) {
	// Mock server: child agent fails (e.g., config not available or bad response)
	var callCount int32

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		switch count {
		case 1:
			// Parent: spawn sub-agent
			resp := mockLLMToolCallResponse([]core.LLMToolCall{
				{
					ID:   "call_spawn",
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "spawn_agent",
						Arguments: `{"agent":"failing_agent","query":"do work"}`,
					},
				},
			})
			_ = json.NewEncoder(w).Encode(resp)
		case 2:
			// Child: return server error
			w.WriteHeader(http.StatusInternalServerError)
			_ = json.NewEncoder(w).Encode(map[string]interface{}{
				"error": map[string]interface{}{
					"message": "internal server error",
					"type":    "server_error",
				},
			})
		default:
			// Parent: handle error gracefully
			resp := mockLLMResponse("Sub-agent failed but I handled it")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "failure-parent",
		Type:          core.StepTypeAgent,
		Query:         "Delegate work",
		MaxIterations: 5,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
		SubAgents: []core.SubAgentDef{
			{
				Name:          "failing_agent",
				Description:   "Agent that will fail",
				MaxIterations: 3,
				AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
			},
		},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	// Parent should succeed — child failure is returned as tool result string
	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Contains(t, result.Output, "handled")
}

func TestAgentExecutor_SubAgentRecursive(t *testing.T) {
	// Test 3-level nesting: parent → child → grandchild
	// Mock server handles 5 calls:
	// 1. Parent spawns child
	// 2. Child spawns grandchild
	// 3. Grandchild returns result
	// 4. Child returns using grandchild result
	// 5. Parent returns using child result
	var callCount int32

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		switch count {
		case 1:
			// Parent: spawn child
			resp := mockLLMToolCallResponse([]core.LLMToolCall{
				{
					ID:   "call_1",
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "spawn_agent",
						Arguments: `{"agent":"child","query":"intermediate task"}`,
					},
				},
			})
			_ = json.NewEncoder(w).Encode(resp)
		case 2:
			// Child: spawn grandchild
			resp := mockLLMToolCallResponse([]core.LLMToolCall{
				{
					ID:   "call_2",
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "spawn_agent",
						Arguments: `{"agent":"grandchild","query":"leaf task"}`,
					},
				},
			})
			_ = json.NewEncoder(w).Encode(resp)
		case 3:
			// Grandchild: return result
			resp := mockLLMResponse("Grandchild result: data collected")
			_ = json.NewEncoder(w).Encode(resp)
		case 4:
			// Child: return using grandchild data
			resp := mockLLMResponse("Child processed grandchild data")
			_ = json.NewEncoder(w).Encode(resp)
		case 5:
			// Parent: final response
			resp := mockLLMResponse("Parent summarized: all levels complete")
			_ = json.NewEncoder(w).Encode(resp)
		default:
			resp := mockLLMResponse("unexpected")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "recursive-parent",
		Type:          core.StepTypeAgent,
		Query:         "Coordinate 3-level task",
		MaxIterations: 5,
		MaxAgentDepth: 3,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
		SubAgents: []core.SubAgentDef{
			{
				Name:          "child",
				Description:   "Intermediate agent",
				MaxIterations: 5,
				AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
				SubAgents: []core.SubAgentDef{
					{
						Name:          "grandchild",
						Description:   "Leaf agent",
						MaxIterations: 5,
						AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
					},
				},
			},
		},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)

	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Contains(t, result.Output, "all levels complete")
	// 5 LLM calls × 30 tokens = 150
	assert.Equal(t, 150, result.Exports["agent_total_tokens"])
}

func TestAgentState_MergeTokens(t *testing.T) {
	state := &agentState{
		totalTokens:      100,
		promptTokens:     60,
		completionTokens: 40,
	}

	state.MergeTokens(30, 20, 10)
	assert.Equal(t, 130, state.totalTokens)
	assert.Equal(t, 80, state.promptTokens)
	assert.Equal(t, 50, state.completionTokens)
}

// ============================================================================
// Group 4: Conversation Compression Tests
// ============================================================================

func TestAgentExecutor_MessageWindowWithSummary(t *testing.T) {
	// Mock server: returns tool calls for several iterations, then final answer.
	// Also handles summarization request.
	var callCount int32

	server := newMockLLMServer(func(w http.ResponseWriter, r *http.Request) {
		count := atomic.AddInt32(&callCount, 1)
		w.Header().Set("Content-Type", "application/json")

		body, _ := io.ReadAll(r.Body)
		var req ChatCompletionRequest
		_ = json.Unmarshal(body, &req)

		// Check if this is a summarization request (no tools)
		if len(req.Tools) == 0 && len(req.Messages) > 0 {
			lastMsg := req.Messages[len(req.Messages)-1]
			if content, ok := lastMsg.Content.(string); ok && strings.Contains(content, "Summarize") {
				resp := mockLLMResponse("Summary of earlier context")
				_ = json.NewEncoder(w).Encode(resp)
				return
			}
		}

		if count <= 3 {
			// Return tool calls to build up conversation
			resp := mockLLMToolCallResponse([]core.LLMToolCall{
				{
					ID:   fmt.Sprintf("call_%d", count),
					Type: "function",
					Function: core.LLMToolCallFunction{
						Name:      "bash",
						Arguments: fmt.Sprintf(`{"command": "echo iter_%d"}`, count),
					},
				},
			})
			_ = json.NewEncoder(w).Encode(resp)
		} else {
			resp := mockLLMResponse("Compression test complete")
			_ = json.NewEncoder(w).Encode(resp)
		}
	})
	defer server.Close()

	cfg := newMockConfig(t, server.URL)
	dispatcher := NewStepDispatcher()
	executor := NewAgentExecutor(dispatcher.GetTemplateEngine(), dispatcher.GetFunctionRegistry())
	executor.SetConfig(cfg)
	executor.SetSilent(true)

	ctx := context.Background()
	execCtx := core.NewExecutionContext("test", core.KindModule, "run-1", "test.com")

	step := &core.Step{
		Name:          "compress-agent",
		Type:          core.StepTypeAgent,
		Query:         "Run multiple commands",
		MaxIterations: 10,
		AgentTools:    []core.AgentToolDef{{Preset: "bash"}},
		Memory: &core.AgentMemoryConfig{
			MaxMessages:         5,
			SummarizeOnTruncate: true,
		},
	}

	result, err := executor.Execute(ctx, step, execCtx)
	require.NoError(t, err)
	assert.Equal(t, core.StepStatusSuccess, result.Status)
	assert.Equal(t, "Compression test complete", result.Output)
}