package agent

import (
	"strings"
	"testing"
	"ollie/toolsrv"
)

func TestBuildPreamble(t *testing.T) {
	toolInfos := []toolsrv.ToolInfo{
		{
			Name:        "file_read",
			Description: "Read a file with line numbers",
			Prompt:      "**file_read**\n\nRead a file with line numbers...",
			InputSchema: []byte(`{"type":"object"}`),
			Server:      "",
		},
		{
			Name:        "file_write",
			Description: "Create or overwrite a file",
			Prompt:      "**file_write**\n\nCreate or overwrite a file...",
			InputSchema: []byte(`{"type":"object"}`),
			Server:      "remote",
		},
		{
			Name:        "shell",
			Description: "Execute shell command",
			Prompt:      "",
			InputSchema: []byte(`{}`),
			Server:      "",
		},
	}

	systemPrompt := "# Core System Prompt\nTest"
	envBlock := "# Environment\n- Test env"
	agentPrompt := "# Agent Prompt\nTest instructions"

	preamble := buildPreamble(systemPrompt, envBlock, agentPrompt, toolInfos)

	// Must contain all sections in order
	expectedSections := []string{
		"# Core System Prompt\nTest\n",
		"# Environment\n- Test env\n",
		"# Agent Prompt\nTest instructions\n",
		"# Tool Sections\n",
		"# Available Tools\n",
		"- **file_read** — Read a file with line numbers\n",
		"# Tool Documentation\n",
		"**file_read**\n\nRead a file with line numbers...\n\n**Schema**:\n```json\n{\"type\":\"object\"}\n```\n\n",
		"## Tool discovery\n",
	}

	for _, section := range expectedSections {
		if !strings.Contains(preamble, section) {
			t.Errorf("buildPreamble missing section: %q", section)
		}
	}

	// Should NOT contain remote server tools
	if strings.Contains(preamble, "file_write") {
		t.Error("buildPreamble should exclude tools with Server != \"\"")
	}

	// Should NOT contain tools without prompts in docs section
	if strings.Contains(preamble, "**shell**\n\n") {
		t.Error("buildPreamble should exclude tools without prompts from docs section")
	}

	// Verify single instances of key headers
	if strings.Count(preamble, "# Available Tools") != 1 {
		t.Errorf("buildPreamble should have exactly one '# Available Tools' section, got %d", 
			strings.Count(preamble, "# Available Tools"))
	}
	if strings.Count(preamble, "# Tool Documentation") != 1 {
		t.Errorf("buildPreamble should have exactly one '# Tool Documentation' section, got %d", 
			strings.Count(preamble, "# Tool Documentation"))
	}

	// Verify tool listing includes all non-server tools with descriptions
	if !strings.Contains(preamble, "**file_read** — Read a file with line numbers") {
		t.Error("buildPreamble should include file_read in tool listing")
	}
	if !strings.Contains(preamble, "**shell** — Execute shell command") {
		t.Error("buildPreamble should include shell in tool listing (has description)")
	}
}

func TestBuildPreambleEmptyTools(t *testing.T) {
	systemPrompt := "# Core System Prompt\nTest"
	envBlock := "# Environment\n- Test env"
	agentPrompt := "# Agent Prompt\nTest instructions"

	preamble := buildPreamble(systemPrompt, envBlock, agentPrompt, nil)

	// Should not contain tool sections when no tools
	if strings.Contains(preamble, "# Tool Sections") {
		t.Error("buildPreamble should not include tool sections when no tools")
	}
	if strings.Contains(preamble, "# Available Tools") {
		t.Error("buildPreamble should not include Available Tools when no tools")
	}
	if strings.Contains(preamble, "# Tool Documentation") {
		t.Error("buildPreamble should not include Tool Documentation when no tools")
	}

	// Should still contain core sections
	expectedSections := []string{
		"# Core System Prompt\nTest\n",
		"# Environment\n- Test env\n",
		"# Agent Prompt\nTest instructions\n",
		"## Tool discovery\n",
	}

	for _, section := range expectedSections {
		if !strings.Contains(preamble, section) {
			t.Errorf("buildPreamble missing core section: %q", section)
		}
	}
}

func TestBuildPreambleOnlyListingNoDocs(t *testing.T) {
	toolInfos := []toolsrv.ToolInfo{
		{
			Name:        "test_tool",
			Description: "Test description",
			Prompt:      "", // No prompt means no docs
			InputSchema: []byte(`{}`),
			Server:      "",
		},
	}

	systemPrompt := "# Core System Prompt\nTest"
	envBlock := "# Environment\n- Test env"
	agentPrompt := "# Agent Prompt\nTest instructions"

	preamble := buildPreamble(systemPrompt, envBlock, agentPrompt, toolInfos)

	// Should have tool sections with listing but no docs
	if !strings.Contains(preamble, "# Tool Sections") {
		t.Error("buildPreamble should include # Tool Sections")
	}
	if !strings.Contains(preamble, "# Available Tools") {
		t.Error("buildPreamble should include # Available Tools")
	}
	if !strings.Contains(preamble, "**test_tool** — Test description") {
		t.Error("buildPreamble should include tool in listing")
	}
	if strings.Contains(preamble, "# Tool Documentation") {
		t.Error("buildPreamble should NOT include # Tool Documentation when no tool has prompt")
	}
}

func TestBuildPreambleOnlyDocsNoListing(t *testing.T) {
	toolInfos := []toolsrv.ToolInfo{
		{
			Name:        "test_tool",
			Description: "", // No description means not in listing
			Prompt:      "**test_tool**\n\nTest docs...",
			InputSchema: []byte(`{}`),
			Server:      "",
		},
	}

	systemPrompt := "# Core System Prompt\nTest"
	envBlock := "# Environment\n- Test env"
	agentPrompt := "# Agent Prompt\nTest instructions"

	preamble := buildPreamble(systemPrompt, envBlock, agentPrompt, toolInfos)

	// Should have tool sections with docs but no listing
	if !strings.Contains(preamble, "# Tool Sections") {
		t.Error("buildPreamble should include # Tool Sections")
	}
	if strings.Contains(preamble, "# Available Tools") {
		t.Error("buildPreamble should NOT include # Available Tools when no tool has description")
	}
	if !strings.Contains(preamble, "# Tool Documentation") {
		t.Error("buildPreamble should include # Tool Documentation")
	}
	if !strings.Contains(preamble, "**test_tool**") {
		t.Error("buildPreamble should include tool in docs")
	}
}