agents_and_robots/shell/llm/claudecode_test.go

package llm

import (
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/enmanuel/agents/internal/config"
	coretypes "github.com/enmanuel/agents/pkg/llm"
)

var discardLog = slog.New(slog.NewTextHandler(io.Discard, nil))

// ── buildClaudeArgs ──────────────────────────────────────────────────────

func TestBuildClaudeArgs_Minimal(t *testing.T) {
	cfg := config.ClaudeCodeCfg{}
	req := coretypes.CompletionRequest{}

	args := buildClaudeArgs(cfg, req)

	// Must always start with --print --output-format json
	want := []string{"--print", "--output-format", "json"}
	if len(args) != len(want) {
		t.Fatalf("got %v, want %v", args, want)
	}
	for i := range want {
		if args[i] != want[i] {
			t.Errorf("args[%d] = %q, want %q", i, args[i], want[i])
		}
	}
}

func TestBuildClaudeArgs_AllOptions(t *testing.T) {
	cfg := config.ClaudeCodeCfg{
		Model:          "sonnet",
		FallbackModel:  "haiku",
		PermissionMode: "bypassPermissions",
		AllowedTools:   []string{"Bash(git:*)", "Read"},
		SessionID:      "abc-123",
		AddDirs:        []string{"/tmp/extra"},
	}
	req := coretypes.CompletionRequest{
		SystemPrompt: "You are a helpful bot",
	}

	args := buildClaudeArgs(cfg, req)

	assertContains(t, args, "--system-prompt", "You are a helpful bot")
	assertContains(t, args, "--model", "sonnet")
	assertContains(t, args, "--fallback-model", "haiku")
	assertContains(t, args, "--permission-mode", "bypassPermissions")
	assertContains(t, args, "--session-id", "abc-123")
	assertContains(t, args, "--add-dir", "/tmp/extra")
	assertContains(t, args, "--allowedTools", "Bash(git:*)")
}

func TestBuildClaudeArgs_DisableTools(t *testing.T) {
	cfg := config.ClaudeCodeCfg{
		DisableTools: true,
		AllowedTools: []string{"Bash"}, // should be ignored
	}
	req := coretypes.CompletionRequest{}

	args := buildClaudeArgs(cfg, req)

	assertContains(t, args, "--tools", "")
	// --allowedTools must NOT appear when disable_tools is set
	for _, a := range args {
		if a == "--allowedTools" {
			t.Error("--allowedTools should not appear when DisableTools=true")
		}
	}
}

func TestBuildClaudeArgs_DisallowedTools(t *testing.T) {
	cfg := config.ClaudeCodeCfg{
		DisallowedTools: []string{"Edit", "Write"},
	}
	req := coretypes.CompletionRequest{}

	args := buildClaudeArgs(cfg, req)
	assertContains(t, args, "--disallowedTools", "Edit")
}

// ── flattenMessages ──────────────────────────────────────────────────────

func TestFlattenMessages_Empty(t *testing.T) {
	got := flattenMessages(nil)
	if got != "" {
		t.Errorf("expected empty, got %q", got)
	}
}

func TestFlattenMessages_MultiRole(t *testing.T) {
	msgs := []coretypes.Message{
		{Role: coretypes.RoleUser, Content: "hello"},
		{Role: coretypes.RoleAssistant, Content: "hi there"},
		{Role: coretypes.RoleTool, Content: `{"time":"12:00"}`},
		{Role: coretypes.RoleUser, Content: "thanks"},
	}

	got := flattenMessages(msgs)

	expects := []string{
		"User: hello",
		"Assistant: hi there",
		`Tool result: {"time":"12:00"}`,
		"User: thanks",
	}
	for _, e := range expects {
		if !contains(got, e) {
			t.Errorf("missing %q in:\n%s", e, got)
		}
	}
}

func TestFlattenMessages_SkipsSystem(t *testing.T) {
	msgs := []coretypes.Message{
		{Role: coretypes.RoleSystem, Content: "system prompt"},
		{Role: coretypes.RoleUser, Content: "hello"},
	}

	got := flattenMessages(msgs)
	if contains(got, "system prompt") {
		t.Error("system messages should not appear in flattened output")
	}
	if !contains(got, "User: hello") {
		t.Error("user message missing")
	}
}

// ── parseClaudeOutput ────────────────────────────────────────────────────

func TestParseClaudeOutput_Success(t *testing.T) {
	output := claudeJSONOutput{
		Type:      "result",
		Subtype:   "success",
		IsError:   false,
		NumTurns:  1,
		Result:    "Hello! I'm Claude.",
		TotalCost: 0.025,
		Usage:     claudeUsage{InputTokens: 10, OutputTokens: 50},
	}
	stdout, _ := json.Marshal(output)

	resp, err := parseClaudeOutput(stdout, nil, nil, 2*time.Second, discardLog)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if resp.Content != "Hello! I'm Claude." {
		t.Errorf("content = %q, want %q", resp.Content, "Hello! I'm Claude.")
	}
	if resp.Usage.InputTokens != 10 {
		t.Errorf("input tokens = %d, want 10", resp.Usage.InputTokens)
	}
	if resp.Usage.OutputTokens != 50 {
		t.Errorf("output tokens = %d, want 50", resp.Usage.OutputTokens)
	}
	if resp.Usage.TotalTokens != 60 {
		t.Errorf("total tokens = %d, want 60", resp.Usage.TotalTokens)
	}
	if resp.FinishReason != "stop" {
		t.Errorf("finish reason = %q, want %q", resp.FinishReason, "stop")
	}
}

func TestParseClaudeOutput_ErrorResponse(t *testing.T) {
	output := claudeJSONOutput{
		IsError: true,
		Result:  "Invalid API key",
	}
	stdout, _ := json.Marshal(output)

	_, err := parseClaudeOutput(stdout, nil, nil, time.Second, discardLog)
	if err == nil {
		t.Fatal("expected error for IsError=true")
	}
	if !contains(err.Error(), "Invalid API key") {
		t.Errorf("error = %q, should contain 'Invalid API key'", err.Error())
	}
}

func TestParseClaudeOutput_ProcessFailedNoStdout(t *testing.T) {
	_, err := parseClaudeOutput(nil, []byte("unknown option\n"), errors.New("exit 1"), time.Second, discardLog)
	if err == nil {
		t.Fatal("expected error when process fails with no stdout")
	}
	if !contains(err.Error(), "unknown option") {
		t.Errorf("error = %q, should contain stderr message", err.Error())
	}
}

func TestParseClaudeOutput_ProcessFailedNoStderr(t *testing.T) {
	_, err := parseClaudeOutput(nil, nil, errors.New("exit 1"), time.Second, discardLog)
	if err == nil {
		t.Fatal("expected error")
	}
	if !contains(err.Error(), "exit 1") {
		t.Errorf("error = %q, should contain exec error", err.Error())
	}
}

func TestParseClaudeOutput_FallbackPlainText(t *testing.T) {
	// Non-JSON stdout should be treated as plain text
	resp, err := parseClaudeOutput([]byte("just plain text\n"), nil, nil, time.Second, discardLog)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if resp.Content != "just plain text" {
		t.Errorf("content = %q, want %q", resp.Content, "just plain text")
	}
}

func TestParseClaudeOutput_ContentBlocks(t *testing.T) {
	output := claudeJSONOutput{
		Result: "", // empty result, content in blocks
		ContentBlock: []claudeContent{
			{Type: "text", Text: "First part."},
			{Type: "text", Text: "Second part."},
		},
		Usage: claudeUsage{InputTokens: 5, OutputTokens: 20},
	}
	stdout, _ := json.Marshal(output)

	resp, err := parseClaudeOutput(stdout, nil, nil, time.Second, discardLog)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if resp.Content != "First part.\nSecond part." {
		t.Errorf("content = %q, want joined blocks", resp.Content)
	}
}

func TestParseClaudeOutput_ExecErrWithStdout(t *testing.T) {
	// Process failed but produced valid JSON output — should parse and set finish_reason=error
	output := claudeJSONOutput{
		Result: "partial answer",
		Usage:  claudeUsage{InputTokens: 3, OutputTokens: 10},
	}
	stdout, _ := json.Marshal(output)

	resp, err := parseClaudeOutput(stdout, nil, errors.New("timeout"), time.Second, discardLog)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if resp.FinishReason != "error" {
		t.Errorf("finish reason = %q, want %q", resp.FinishReason, "error")
	}
	if resp.Content != "partial answer" {
		t.Errorf("content = %q", resp.Content)
	}
}

// ── filterEnv ────────────────────────────────────────────────────────────

func TestFilterEnv_RemovesSingleKey(t *testing.T) {
	env := []string{
		"HOME=/home/user",
		"ANTHROPIC_API_KEY=sk-secret",
		"PATH=/usr/bin",
	}

	got := filterEnv(env, "ANTHROPIC_API_KEY")

	if len(got) != 2 {
		t.Fatalf("expected 2 entries, got %d: %v", len(got), got)
	}
	for _, e := range got {
		if contains(e, "ANTHROPIC_API_KEY") {
			t.Errorf("ANTHROPIC_API_KEY should have been removed: %v", got)
		}
	}
}

func TestFilterEnv_RemovesMultipleKeys(t *testing.T) {
	env := []string{
		"HOME=/home/user",
		"ANTHROPIC_API_KEY=sk-secret",
		"OPENAI_API_KEY=sk-openai",
		"PATH=/usr/bin",
	}

	got := filterEnv(env, "ANTHROPIC_API_KEY", "OPENAI_API_KEY")

	if len(got) != 2 {
		t.Fatalf("expected 2 entries, got %d: %v", len(got), got)
	}
}

func TestFilterEnv_NoMatchKeepsAll(t *testing.T) {
	env := []string{"HOME=/home/user", "PATH=/usr/bin"}

	got := filterEnv(env, "NONEXISTENT")

	if len(got) != 2 {
		t.Fatalf("expected 2, got %d", len(got))
	}
}

func TestFilterEnv_PrefixSafety(t *testing.T) {
	// ANTHROPIC_API_KEY_V2 should NOT be removed when filtering ANTHROPIC_API_KEY
	env := []string{
		"ANTHROPIC_API_KEY=secret",
		"ANTHROPIC_API_KEY_V2=other",
	}

	got := filterEnv(env, "ANTHROPIC_API_KEY")

	if len(got) != 1 {
		t.Fatalf("expected 1, got %d: %v", len(got), got)
	}
	if got[0] != "ANTHROPIC_API_KEY_V2=other" {
		t.Errorf("wrong entry kept: %q", got[0])
	}
}

// ── resolveWorkDir ──────────────────────────────────────────────────────

func TestResolveWorkDir_EmptyCreatesTempDir(t *testing.T) {
	dir := resolveWorkDir("", discardLog)
	if dir == "" {
		t.Fatal("expected a temp directory, got empty string")
	}
	defer os.RemoveAll(dir)

	if !strings.Contains(dir, "claude-agent-") {
		t.Errorf("temp dir %q should contain 'claude-agent-' prefix", dir)
	}

	info, err := os.Stat(dir)
	if err != nil {
		t.Fatalf("temp dir should exist: %v", err)
	}
	if !info.IsDir() {
		t.Error("temp dir should be a directory")
	}
}

func TestResolveWorkDir_ConfiguredValueUsed(t *testing.T) {
	want := filepath.Join(t.TempDir(), "custom-workdir")

	got := resolveWorkDir(want, discardLog)

	if got != want {
		t.Errorf("got %q, want %q", got, want)
	}

	info, err := os.Stat(got)
	if err != nil {
		t.Fatalf("configured dir should be created: %v", err)
	}
	if !info.IsDir() {
		t.Error("configured dir should be a directory")
	}
}

func TestResolveWorkDir_ConfiguredAlreadyExists(t *testing.T) {
	want := t.TempDir() // already exists

	got := resolveWorkDir(want, discardLog)

	if got != want {
		t.Errorf("got %q, want %q", got, want)
	}
}

// ── helpers ──────────────────────────────────────────────────────────────

func contains(s, substr string) bool {
	return len(s) >= len(substr) && (s == substr || len(substr) == 0 ||
		(len(s) > 0 && stringContains(s, substr)))
}

func stringContains(s, sub string) bool {
	for i := 0; i <= len(s)-len(sub); i++ {
		if s[i:i+len(sub)] == sub {
			return true
		}
	}
	return false
}

func assertContains(t *testing.T, args []string, flag, value string) {
	t.Helper()
	for i, a := range args {
		if a == flag && i+1 < len(args) && args[i+1] == value {
			return
		}
		// For --tools "" where value is empty string
		if a == flag && value == "" && i+1 < len(args) && args[i+1] == "" {
			return
		}
	}
	t.Errorf("args %v missing %s %q", args, flag, value)
}