package agent

import (
	"context"
	"strings"
	"testing"

	"github.com/ollama/ollama/api"
)

type mockTool struct {
	name string
}

func (m mockTool) Name() string        { return m.name }
func (m mockTool) Description() string { return "" }
func (m mockTool) Schema() api.ToolFunction {
	return api.ToolFunction{Name: m.name}
}

func (m mockTool) Execute(context.Context, ToolContext, map[string]any) (ToolResult, error) {
	return ToolResult{}, nil
}

func TestToolApprovalScopeUsesScopedTool(t *testing.T) {
	shellTool := mockScopedTool{
		mockTool: mockTool{name: "bash"},
		scope: func(args map[string]any) string {
			if cmd, ok := args["command"].(string); ok {
				cmd = strings.TrimSpace(cmd)
				if cmd != "" {
					return "bash\x00" + cmd
				}
			}
			return "bash"
		},
	}
	plainTool := mockTool{name: "edit"}

	tests := []struct {
		tool Tool
		name string
		args map[string]any
		want string
	}{
		{shellTool, "bash", map[string]any{"command": " pwd "}, "bash\x00pwd"},
		{shellTool, "bash", map[string]any{"command": "Get-ChildItem"}, "bash\x00Get-ChildItem"},
		{plainTool, "edit", map[string]any{"path": "README.md"}, "edit"},
	}
	for _, tt := range tests {
		if got := toolApprovalScope(tt.tool, tt.name, tt.args); got != tt.want {
			t.Fatalf("toolApprovalScope(%q) = %q, want %q", tt.name, got, tt.want)
		}
	}
}

type mockScopedTool struct {
	mockTool
	scope func(args map[string]any) string
}

func (m mockScopedTool) ApprovalScope(args map[string]any) string {
	return m.scope(args)
}

func TestSessionApplyApprovalScopes(t *testing.T) {
	session := &Session{}
	result := Approval{AllowScopes: []string{"edit", "bash\x00pwd", " "}}

	session.applyApproval(&result)

	if !result.Allow {
		t.Fatal("scoped approval should allow the current request")
	}
	if !session.allows("edit") || !session.allows("bash\x00pwd") {
		t.Fatal("scoped approval was not saved")
	}
	if session.allows("bash") || session.allows("bash\x00ls") {
		t.Fatal("shell approval was too broad")
	}
	if session.ApprovalState.AllGranted() {
		t.Fatal("allow all = true, want false for scoped approval")
	}
}

func TestSessionApplyApprovalAllowAll(t *testing.T) {
	session := &Session{}
	result := Approval{AllowAll: true}

	session.applyApproval(&result)

	if !result.Allow || !session.allows("anything") {
		t.Fatalf("allow all = %v result = %#v, want allow all", session.ApprovalState.AllGranted(), result)
	}
}
