package llm

import (
	"context"
	"strings"
	"testing"

	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/config"
	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime/document"
	brtypes "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types"
)

func TestAWSCredentialModeFor(t *testing.T) {
	if got := awsCredentialModeFor("AKIA", "secret", "rohit"); got != awsCredentialStatic {
		t.Fatalf("static keys should win over a profile, got %v", got)
	}
	if got := awsCredentialModeFor("AKIA", "", "rohit"); got != awsCredentialProfile {
		t.Fatalf("access key without a secret should fall through to the profile, got %v", got)
	}
	if got := awsCredentialModeFor("", "", "rohit"); got != awsCredentialProfile {
		t.Fatalf("profile only, got %v", got)
	}
	if got := awsCredentialModeFor("", "", ""); got != awsCredentialDefault {
		t.Fatalf("empty config should use the default chain, got %v", got)
	}
}

func TestBedrockProvider_FactoryRequiresExplicitModel(t *testing.T) {
	for _, model := range []string{"", "claude-sonnet-4-6"} {
		_, err := NewProvider(config.OrchestratorConfig{
			Provider:   "bedrock",
			Model:      model,
			AWSProfile: "rohit",
		})
		if err == nil {
			t.Fatalf("expected error for model %q — bedrock needs an explicit model", model)
		}
		if !strings.Contains(err.Error(), "model") {
			t.Errorf("error should mention model: %v", err)
		}
	}
}

func TestBedrockProvider_FactoryDefaults(t *testing.T) {
	p, err := NewProvider(config.OrchestratorConfig{
		Provider:   "bedrock",
		Model:      "zai.glm-5",
		AWSProfile: "rohit",
	})
	if err != nil {
		t.Fatalf("NewProvider(bedrock): %v", err)
	}
	bp, ok := p.(*BedrockProvider)
	if !ok {
		t.Fatalf("got %T, want *BedrockProvider", p)
	}
	if bp.model != "zai.glm-5" {
		t.Errorf("model = %q, want zai.glm-5", bp.model)
	}
	if bp.region != "us-east-1" {
		t.Errorf("region = %q, want us-east-1", bp.region)
	}
	if bp.profile != "rohit" {
		t.Errorf("profile = %q, want rohit", bp.profile)
	}
	if awsCredentialModeFor(bp.accessKey, bp.secret, bp.profile) != awsCredentialProfile {
		t.Error("expected profile credential mode")
	}
	if !bp.SupportsToolUse() {
		t.Error("bedrock should advertise tool use")
	}
}

func TestBedrockProvider_FactoryStaticKeysAndExplicitModel(t *testing.T) {
	p, err := NewProvider(config.OrchestratorConfig{
		Provider:           "bedrock",
		Model:              "zai.glm-4.7",
		Region:             "eu-west-1",
		AWSProfile:         "ignored",
		AWSAccessKeyID:     "AKIAEXAMPLE",
		AWSSecretAccessKey: "secret",
		AWSSessionToken:    "token",
		ContextWindow:      128000,
	})
	if err != nil {
		t.Fatalf("NewProvider: %v", err)
	}
	bp := p.(*BedrockProvider)
	if bp.model != "zai.glm-4.7" {
		t.Errorf("model = %q", bp.model)
	}
	if bp.region != "eu-west-1" {
		t.Errorf("region = %q", bp.region)
	}
	if bp.contextWindow != 128000 {
		t.Errorf("context window = %d", bp.contextWindow)
	}
	if awsCredentialModeFor(bp.accessKey, bp.secret, bp.profile) != awsCredentialStatic {
		t.Error("static keys should take precedence over the profile")
	}
	if bp.sessionToken != "token" {
		t.Errorf("session token = %q", bp.sessionToken)
	}
}

func TestBedrockProvider_AgentInheritsProfileAndModel(t *testing.T) {
	p, err := NewAgentProvider(config.AgentModelConfig{}, config.OrchestratorConfig{
		Provider:   "bedrock",
		Model:      "zai.glm-5",
		AWSProfile: "rohit",
		Region:     "us-east-1",
	})
	if err != nil {
		t.Fatalf("NewAgentProvider: %v", err)
	}
	bp := p.(*BedrockProvider)
	if bp.model != "zai.glm-5" {
		t.Errorf("inherited model = %q, want zai.glm-5", bp.model)
	}
	if bp.profile != "rohit" {
		t.Errorf("profile = %q", bp.profile)
	}
}

type fakeConverse struct {
	got *bedrockruntime.ConverseInput
	out *bedrockruntime.ConverseOutput
	err error
}

func (f *fakeConverse) Converse(_ context.Context, params *bedrockruntime.ConverseInput, _ ...func(*bedrockruntime.Options)) (*bedrockruntime.ConverseOutput, error) {
	f.got = params
	return f.out, f.err
}

func (f *fakeConverse) ConverseStream(context.Context, *bedrockruntime.ConverseStreamInput, ...func(*bedrockruntime.Options)) (*bedrockruntime.ConverseStreamOutput, error) {
	return nil, nil
}

func TestBedrockProvider_CompleteMapsRequestAndResponse(t *testing.T) {
	fake := &fakeConverse{
		out: &bedrockruntime.ConverseOutput{
			StopReason: brtypes.StopReasonToolUse,
			Usage: &brtypes.TokenUsage{
				InputTokens:           aws.Int32(11),
				OutputTokens:          aws.Int32(7),
				TotalTokens:           aws.Int32(18),
				CacheReadInputTokens:  aws.Int32(3),
				CacheWriteInputTokens: aws.Int32(2),
			},
			Output: &brtypes.ConverseOutputMemberMessage{
				Value: brtypes.Message{
					Role: brtypes.ConversationRoleAssistant,
					Content: []brtypes.ContentBlock{
						&brtypes.ContentBlockMemberText{Value: "running a probe"},
						&brtypes.ContentBlockMemberToolUse{Value: brtypes.ToolUseBlock{
							ToolUseId: aws.String("call_1"),
							Name:      aws.String("http_request"),
							Input:     document.NewLazyDocument(map[string]any{"url": "https://example.test"}),
						}},
					},
				},
			},
		},
	}
	p := NewBedrockProvider(BedrockProviderConfig{
		Model:      "zai.glm-5",
		Region:     "us-east-1",
		AWSProfile: "rohit",
	})
	p.converse = fake

	resp, err := p.Complete(context.Background(), CompletionRequest{
		SystemPrompt: "you are a tester",
		Messages: []Message{
			{Role: "system", Content: "stay in scope"},
			{Role: "user", Content: "scan the host"},
			{Role: "assistant", Content: "checking"},
			{Role: "tool", ToolCallID: "call_0", Content: "open port 443"},
			{Role: "tool", ToolCallID: "call_0b", Content: "title example"},
		},
		Tools: []Tool{{
			Name:        "http_request",
			Description: "fetch a URL",
			Parameters:  []byte(`{"type":"object","properties":{"url":{"type":"string"}}}`),
		}},
		MaxTokens:   128,
		Temperature: 0.1,
	})
	if err != nil {
		t.Fatalf("Complete: %v", err)
	}
	in := fake.got
	if in == nil {
		t.Fatal("Converse was not called")
	}
	if aws.ToString(in.ModelId) != "zai.glm-5" {
		t.Errorf("model = %q", aws.ToString(in.ModelId))
	}
	if aws.ToInt32(in.InferenceConfig.MaxTokens) != 128 {
		t.Errorf("max tokens = %d", aws.ToInt32(in.InferenceConfig.MaxTokens))
	}
	if len(in.System) != 2 {
		t.Fatalf("system blocks = %d, want 2", len(in.System))
	}
	if len(in.Messages) != 3 {
		t.Fatalf("messages = %d, want user, assistant, and coalesced tool results", len(in.Messages))
	}
	if in.Messages[0].Role != brtypes.ConversationRoleUser {
		t.Errorf("first role = %s", in.Messages[0].Role)
	}
	if in.Messages[1].Role != brtypes.ConversationRoleAssistant {
		t.Errorf("second role = %s", in.Messages[1].Role)
	}
	if in.Messages[2].Role != brtypes.ConversationRoleUser {
		t.Errorf("tool results should be a user message, got %s", in.Messages[2].Role)
	}
	if len(in.Messages[2].Content) != 2 {
		t.Fatalf("coalesced tool results = %d, want 2", len(in.Messages[2].Content))
	}
	if in.ToolConfig == nil || len(in.ToolConfig.Tools) != 1 {
		t.Fatalf("tools = %#v", in.ToolConfig)
	}
	spec, ok := in.ToolConfig.Tools[0].(*brtypes.ToolMemberToolSpec)
	if !ok || aws.ToString(spec.Value.Name) != "http_request" {
		t.Fatalf("tool spec = %#v", in.ToolConfig.Tools[0])
	}

	if resp.Content != "running a probe" {
		t.Errorf("content = %q", resp.Content)
	}
	if resp.StopReason != string(brtypes.StopReasonToolUse) {
		t.Errorf("stop = %q", resp.StopReason)
	}
	if resp.Usage.InputTokens != 11 || resp.Usage.OutputTokens != 7 {
		t.Errorf("usage = %+v", resp.Usage)
	}
	if resp.Usage.CacheReadInputTokens != 3 || resp.Usage.CacheCreationInputTokens != 2 {
		t.Errorf("cache usage = %+v", resp.Usage)
	}
	if len(resp.ToolCalls) != 1 {
		t.Fatalf("tool calls = %d", len(resp.ToolCalls))
	}
	tc := resp.ToolCalls[0]
	if tc.ID != "call_1" || tc.Name != "http_request" {
		t.Errorf("tool call = %+v", tc)
	}
	if !strings.Contains(tc.Arguments, "example.test") {
		t.Errorf("arguments = %s", tc.Arguments)
	}
}