package mcp

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"strings"
	"testing"
	"time"
)

// ============================================================
// Test RequestBuilder Basic Features
// ============================================================

func TestRequestBuilder_BasicUsage(t *testing.T) {
	request, err := NewRequestBuilder().
		WithSystemPrompt("You are helpful").
		WithUserPrompt("Hello").
		Build()

	if err != nil {
		t.Fatalf("Build should not error: %v", err)
	}

	if len(request.Messages) != 2 {
		t.Errorf("expected 2 messages, got %d", len(request.Messages))
	}

	if request.Messages[0].Role != "system" {
		t.Errorf("first message should be system, got %s", request.Messages[0].Role)
	}

	if request.Messages[1].Role != "user" {
		t.Errorf("second message should be user, got %s", request.Messages[1].Role)
	}
}

func TestRequestBuilder_EmptyMessages(t *testing.T) {
	_, err := NewRequestBuilder().Build()

	if err == nil {
		t.Error("Build should error when no messages")
	}

	if err.Error() != "at least one message is required" {
		t.Errorf("unexpected error: %v", err)
	}
}

// ============================================================
// Test Message Building Methods
// ============================================================

func TestRequestBuilder_MultipleMessages(t *testing.T) {
	request := NewRequestBuilder().
		AddSystemMessage("You are helpful").
		AddUserMessage("What is Go?").
		AddAssistantMessage("Go is a programming language").
		AddUserMessage("Tell me more").
		MustBuild()

	if len(request.Messages) != 4 {
		t.Fatalf("expected 4 messages, got %d", len(request.Messages))
	}

	expectedRoles := []string{"system", "user", "assistant", "user"}
	for i, expected := range expectedRoles {
		if request.Messages[i].Role != expected {
			t.Errorf("message %d: expected role %s, got %s", i, expected, request.Messages[i].Role)
		}
	}
}

func TestRequestBuilder_AddConversationHistory(t *testing.T) {
	history := []Message{
		NewUserMessage("Previous question"),
		NewAssistantMessage("Previous answer"),
	}

	request := NewRequestBuilder().
		AddConversationHistory(history).
		AddUserMessage("New question").
		MustBuild()

	if len(request.Messages) != 3 {
		t.Fatalf("expected 3 messages, got %d", len(request.Messages))
	}
}

// ============================================================
// Test Parameter Control Methods
// ============================================================

func TestRequestBuilder_WithTemperature(t *testing.T) {
	request := NewRequestBuilder().
		WithUserPrompt("Hello").
		WithTemperature(0.8).
		MustBuild()

	if request.Temperature == nil {
		t.Fatal("Temperature should be set")
	}

	if *request.Temperature != 0.8 {
		t.Errorf("expected temperature 0.8, got %f", *request.Temperature)
	}
}

func TestRequestBuilder_WithMaxTokens(t *testing.T) {
	request := NewRequestBuilder().
		WithUserPrompt("Hello").
		WithMaxTokens(2000).
		MustBuild()

	if request.MaxTokens == nil {
		t.Fatal("MaxTokens should be set")
	}

	if *request.MaxTokens != 2000 {
		t.Errorf("expected maxTokens 2000, got %d", *request.MaxTokens)
	}
}

func TestRequestBuilder_WithTopP(t *testing.T) {
	request := NewRequestBuilder().
		WithUserPrompt("Hello").
		WithTopP(0.9).
		MustBuild()

	if request.TopP == nil {
		t.Fatal("TopP should be set")
	}

	if *request.TopP != 0.9 {
		t.Errorf("expected topP 0.9, got %f", *request.TopP)
	}
}

func TestRequestBuilder_WithPenalties(t *testing.T) {
	request := NewRequestBuilder().
		WithUserPrompt("Hello").
		WithFrequencyPenalty(0.5).
		WithPresencePenalty(0.6).
		MustBuild()

	if request.FrequencyPenalty == nil || *request.FrequencyPenalty != 0.5 {
		t.Error("FrequencyPenalty should be 0.5")
	}

	if request.PresencePenalty == nil || *request.PresencePenalty != 0.6 {
		t.Error("PresencePenalty should be 0.6")
	}
}

func TestRequestBuilder_WithStopSequences(t *testing.T) {
	request := NewRequestBuilder().
		WithUserPrompt("Hello").
		WithStopSequences([]string{"STOP", "END"}).
		MustBuild()

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

	if request.Stop[0] != "STOP" || request.Stop[1] != "END" {
		t.Error("stop sequences not set correctly")
	}
}

// ============================================================
// Test Tool/Function Calling
// ============================================================

func TestRequestBuilder_AddTool(t *testing.T) {
	tool := Tool{
		Type: "function",
		Function: FunctionDef{
			Name:        "get_weather",
			Description: "Get weather",
			Parameters: map[string]any{
				"type": "object",
				"properties": map[string]any{
					"location": map[string]any{"type": "string"},
				},
			},
		},
	}

	request := NewRequestBuilder().
		WithUserPrompt("What's the weather?").
		AddTool(tool).
		WithToolChoice("auto").
		MustBuild()

	if len(request.Tools) != 1 {
		t.Fatalf("expected 1 tool, got %d", len(request.Tools))
	}

	if request.Tools[0].Function.Name != "get_weather" {
		t.Error("tool not added correctly")
	}

	if request.ToolChoice != "auto" {
		t.Error("tool choice not set correctly")
	}
}

func TestRequestBuilder_AddFunction(t *testing.T) {
	params := map[string]any{
		"type": "object",
		"properties": map[string]any{
			"city": map[string]any{"type": "string"},
		},
	}

	request := NewRequestBuilder().
		WithUserPrompt("Hello").
		AddFunction("get_weather", "Get current weather", params).
		MustBuild()

	if len(request.Tools) != 1 {
		t.Fatalf("expected 1 tool, got %d", len(request.Tools))
	}

	if request.Tools[0].Type != "function" {
		t.Error("tool type should be function")
	}

	if request.Tools[0].Function.Name != "get_weather" {
		t.Error("function name not set correctly")
	}
}

// ============================================================
// Test Convenience Methods
// ============================================================

func TestRequestBuilder_ForChat(t *testing.T) {
	request := ForChat().
		WithUserPrompt("Hello").
		MustBuild()

	if request.Temperature == nil {
		t.Fatal("ForChat should set temperature")
	}

	if *request.Temperature != 0.7 {
		t.Errorf("ForChat should set temperature to 0.7, got %f", *request.Temperature)
	}

	if request.MaxTokens == nil {
		t.Fatal("ForChat should set maxTokens")
	}

	if *request.MaxTokens != 2000 {
		t.Errorf("ForChat should set maxTokens to 2000, got %d", *request.MaxTokens)
	}
}

func TestRequestBuilder_ForCodeGeneration(t *testing.T) {
	request := ForCodeGeneration().
		WithUserPrompt("Generate code").
		MustBuild()

	if request.Temperature == nil || *request.Temperature != 0.2 {
		t.Error("ForCodeGeneration should set low temperature")
	}

	if request.TopP == nil || *request.TopP != 0.1 {
		t.Error("ForCodeGeneration should set low topP")
	}
}

func TestRequestBuilder_ForCreativeWriting(t *testing.T) {
	request := ForCreativeWriting().
		WithUserPrompt("Write a story").
		MustBuild()

	if request.Temperature == nil || *request.Temperature != 1.2 {
		t.Error("ForCreativeWriting should set high temperature")
	}

	if request.PresencePenalty == nil || *request.PresencePenalty != 0.6 {
		t.Error("ForCreativeWriting should set presence penalty")
	}

	if request.FrequencyPenalty == nil || *request.FrequencyPenalty != 0.5 {
		t.Error("ForCreativeWriting should set frequency penalty")
	}
}

// ============================================================
// Test CallWithRequest Integration
// ============================================================

func TestClient_CallWithRequest_Success(t *testing.T) {
	mockHTTP := NewMockHTTPClient()
	mockHTTP.SetSuccessResponse("Builder response")
	mockLogger := NewMockLogger()

	client := NewClient(
		WithHTTPClient(mockHTTP.ToHTTPClient()),
		WithLogger(mockLogger),
		WithAPIKey("sk-test-key"),
	)

	request := NewRequestBuilder().
		WithSystemPrompt("You are helpful").
		WithUserPrompt("Hello").
		WithTemperature(0.8).
		MustBuild()

	result, err := client.CallWithRequest(request)

	if err != nil {
		t.Fatalf("should not error: %v", err)
	}

	if result != "Builder response" {
		t.Errorf("expected 'Builder response', got '%s'", result)
	}

	// Verify request body
	requests := mockHTTP.GetRequests()
	if len(requests) != 1 {
		t.Fatalf("expected 1 request, got %d", len(requests))
	}

	// Parse request body to verify parameters
	var body map[string]interface{}
	decoder := json.NewDecoder(requests[0].Body)
	if err := decoder.Decode(&body); err != nil {
		t.Fatalf("failed to decode request body: %v", err)
	}

	// Verify temperature
	if body["temperature"] != 0.8 {
		t.Errorf("expected temperature 0.8, got %v", body["temperature"])
	}

	// Verify messages
	messages, ok := body["messages"].([]interface{})
	if !ok || len(messages) != 2 {
		t.Error("messages not correctly formatted")
	}
}

func TestClient_CallWithRequest_AttachesRequestContextToHTTP(t *testing.T) {
	type contextKey string
	const key contextKey = "stage"
	ctx := context.WithValue(context.Background(), key, "planner")

	mockHTTP := NewMockHTTPClient()
	mockHTTP.ResponseFunc = func(req *http.Request) (*http.Response, error) {
		if req.Context().Value(key) != "planner" {
			t.Fatalf("expected HTTP request to inherit mcp.Request context")
		}
		return &http.Response{
			StatusCode: http.StatusOK,
			Body:       io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"ok"}}]}`)),
			Header:     make(http.Header),
		}, nil
	}

	client := NewClient(
		WithHTTPClient(mockHTTP.ToHTTPClient()),
		WithLogger(NewMockLogger()),
		WithAPIKey("sk-test-key"),
	)
	request := NewRequestBuilder().WithUserPrompt("Hello").MustBuild()
	request.Ctx = ctx

	result, err := client.CallWithRequest(request)
	if err != nil {
		t.Fatalf("should not error: %v", err)
	}
	if result != "ok" {
		t.Fatalf("expected ok, got %q", result)
	}
}

func TestClient_CallWithRequest_RetrySleepStopsWhenContextCancelled(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	mockHTTP := NewMockHTTPClient()
	mockHTTP.SetNetworkError(io.EOF)
	client := NewClient(
		WithHTTPClient(mockHTTP.ToHTTPClient()),
		WithLogger(NewMockLogger()),
		WithAPIKey("sk-test-key"),
		WithMaxRetries(2),
		WithRetryWaitBase(time.Hour),
	)
	request := NewRequestBuilder().WithUserPrompt("Hello").MustBuild()
	request.Ctx = ctx

	start := time.Now()
	_, err := client.CallWithRequest(request)
	if err == nil || !strings.Contains(err.Error(), "context canceled") {
		t.Fatalf("expected context canceled during retry wait, got %v", err)
	}
	if elapsed := time.Since(start); elapsed > 500*time.Millisecond {
		t.Fatalf("retry sleep did not respect context cancellation, elapsed=%v", elapsed)
	}
	if got := len(mockHTTP.GetRequests()); got != 1 {
		t.Fatalf("expected no retry after context cancellation, got %d requests", got)
	}
}

func TestClient_CallWithRequest_RetriesUpstreamEmptyOutput(t *testing.T) {
	mockHTTP := NewMockHTTPClient()
	attempts := 0
	mockHTTP.ResponseFunc = func(req *http.Request) (*http.Response, error) {
		attempts++
		if attempts == 1 {
			body := `{"error":{"code":"upstream_empty_output","message":"Upstream model returned empty output.","type":"rate_limit_error"}}`
			return &http.Response{
				StatusCode: http.StatusTooManyRequests,
				Body:       io.NopCloser(strings.NewReader(body)),
				Header:     make(http.Header),
			}, nil
		}
		return &http.Response{
			StatusCode: http.StatusOK,
			Body:       io.NopCloser(strings.NewReader(`{"choices":[{"message":{"content":"ok after retry"}}]}`)),
			Header:     make(http.Header),
		}, nil
	}

	client := NewClient(
		WithHTTPClient(mockHTTP.ToHTTPClient()),
		WithLogger(NewMockLogger()),
		WithAPIKey("sk-test-key"),
		WithMaxRetries(2),
		WithRetryWaitBase(time.Millisecond),
	)
	request := NewRequestBuilder().WithUserPrompt("Hello").MustBuild()

	result, err := client.CallWithRequest(request)
	if err != nil {
		t.Fatalf("should retry upstream empty output and succeed: %v", err)
	}
	if result != "ok after retry" {
		t.Fatalf("expected retry result, got %q", result)
	}
	if attempts != 2 {
		t.Fatalf("expected 2 attempts, got %d", attempts)
	}
}

func TestClient_CallWithRequest_MultiRound(t *testing.T) {
	mockHTTP := NewMockHTTPClient()
	mockHTTP.SetSuccessResponse("Multi-round response")
	mockLogger := NewMockLogger()

	client := NewClient(
		WithHTTPClient(mockHTTP.ToHTTPClient()),
		WithLogger(mockLogger),
		WithAPIKey("sk-test-key"),
	)

	// Build multi-round conversation
	request := NewRequestBuilder().
		AddSystemMessage("You are a trading advisor").
		AddUserMessage("Analyze BTC").
		AddAssistantMessage("BTC is bullish").
		AddUserMessage("What about entry point?").
		WithTemperature(0.3).
		MustBuild()

	result, err := client.CallWithRequest(request)

	if err != nil {
		t.Fatalf("should not error: %v", err)
	}

	if result != "Multi-round response" {
		t.Errorf("expected 'Multi-round response', got '%s'", result)
	}

	// Verify request body contains all messages
	requests := mockHTTP.GetRequests()
	var body map[string]interface{}
	json.NewDecoder(requests[0].Body).Decode(&body)

	messages := body["messages"].([]interface{})
	if len(messages) != 4 {
		t.Errorf("expected 4 messages in request, got %d", len(messages))
	}
}

func TestClient_CallWithRequest_WithTools(t *testing.T) {
	mockHTTP := NewMockHTTPClient()
	mockHTTP.SetSuccessResponse("Tool response")
	mockLogger := NewMockLogger()

	client := NewClient(
		WithHTTPClient(mockHTTP.ToHTTPClient()),
		WithLogger(mockLogger),
		WithAPIKey("sk-test-key"),
	)

	request := NewRequestBuilder().
		WithUserPrompt("What's the weather in Beijing?").
		AddFunction("get_weather", "Get weather", map[string]any{
			"type": "object",
			"properties": map[string]any{
				"location": map[string]any{"type": "string"},
			},
		}).
		WithToolChoice("auto").
		MustBuild()

	_, err := client.CallWithRequest(request)

	if err != nil {
		t.Fatalf("should not error: %v", err)
	}

	// Verify request body contains tools
	requests := mockHTTP.GetRequests()
	var body map[string]interface{}
	json.NewDecoder(requests[0].Body).Decode(&body)

	tools, ok := body["tools"].([]interface{})
	if !ok || len(tools) == 0 {
		t.Error("tools should be present in request")
	}

	toolChoice, ok := body["tool_choice"].(string)
	if !ok || toolChoice != "auto" {
		t.Error("tool_choice should be 'auto'")
	}
}

func TestClient_CallWithRequest_NoAPIKey(t *testing.T) {
	client := NewClient()

	request := NewRequestBuilder().
		WithUserPrompt("Hello").
		MustBuild()

	_, err := client.CallWithRequest(request)

	if err == nil {
		t.Error("should error when API key not set")
	}

	if err.Error() != "AI API key not set, please call SetAPIKey first" {
		t.Errorf("unexpected error: %v", err)
	}
}

func TestClient_CallWithRequest_UsesClientModel(t *testing.T) {
	mockHTTP := NewMockHTTPClient()
	mockHTTP.SetSuccessResponse("Response")
	mockLogger := NewMockLogger()

	client := NewClient(
		WithDeepSeekConfig("sk-test-key"),
		WithHTTPClient(mockHTTP.ToHTTPClient()),
		WithLogger(mockLogger),
	)

	// Request does not set model, should use Client's model
	request := NewRequestBuilder().
		WithUserPrompt("Hello").
		MustBuild()

	if request.Model != "" {
		t.Error("request.Model should be empty initially")
	}

	client.CallWithRequest(request)

	// Verify DeepSeek's model is used
	requests := mockHTTP.GetRequests()
	var body map[string]interface{}
	json.NewDecoder(requests[0].Body).Decode(&body)

	if body["model"] != DefaultDeepSeekModel {
		t.Errorf("expected model %s, got %v", DefaultDeepSeekModel, body["model"])
	}
}
