package main

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus/testutil"
)

// TestResponseValidation_EmptyJSONNonStreaming tests that empty JSON body on 2xx responses
// triggers retry up to MAX_RETRIES times, then returns 502 Bad Gateway.
func TestResponseValidation_EmptyJSONNonStreaming(t *testing.T) {
	testCases := []struct {
		name        string
		maxRetries  int
		wantCalls   int32 // Expected upstream calls: 1 initial + maxRetries
	}{
		{
			name:       "MAX_RETRIES=1",
			maxRetries: 1,
			wantCalls:  2, // 1 initial + 1 retry
		},
		{
			name:       "MAX_RETRIES=2",
			maxRetries: 2,
			wantCalls:  3, // 1 initial + 2 retries
		},
		{
			name:       "MAX_RETRIES=5",
			maxRetries: 5,
			wantCalls:  6, // 1 initial + 5 retries
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			var upstreamCallCount atomic.Int32

			// Create upstream that returns 200 with empty body every time
			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				upstreamCallCount.Add(1)
				w.WriteHeader(http.StatusOK)
				// Empty body - write nothing
			}))
			defer upstream.Close()

			handler := createTestProxyHandler(t, upstream.URL, tc.maxRetries)

			req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody()))
			req.Header.Set("Content-Type", "application/json")
			w := httptest.NewRecorder()

			start := time.Now()
			handler.ServeHTTP(w, req)
			duration := time.Since(start)

			// Verify 502 Bad Gateway is returned after exhausting retries
			if w.Code != http.StatusBadGateway {
				t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
			}

			// Verify exact upstream call count: 1 initial + maxRetries retries
			actualCalls := upstreamCallCount.Load()
			if actualCalls != tc.wantCalls {
				t.Errorf("Expected %d upstream calls (1 initial + %d retries), got %d",
					tc.wantCalls, tc.maxRetries, actualCalls)
			}

			// Verify exponential backoff took reasonable time
			// With maxRetries=N, backoff durations are: 1s, 2s, 4s, ...
			// For maxRetries=1: ~1s delay
			// For maxRetries=2: ~1s+2s=3s delay
			// For maxRetries=5: ~1s+2s+4s+8s+16s=31s delay
			// We allow 20% tolerance for test execution variance
			minExpectedDuration := time.Duration(1<<uint(tc.maxRetries-1)) * time.Second
			if tc.maxRetries > 0 && duration < minExpectedDuration {
				t.Logf("WARNING: Duration %v is less than expected minimum %v", duration, minExpectedDuration)
			}

			t.Logf("PASS: Empty JSON body triggered %d retries, %d upstream calls in %v, returned 502",
				tc.maxRetries, actualCalls, duration)
		})
	}
}

// TestResponseValidation_InvalidJSONNonStreaming tests that invalid JSON body on 2xx responses
// triggers retry up to MAX_RETRIES times, then returns 502 Bad Gateway.
func TestResponseValidation_InvalidJSONNonStreaming(t *testing.T) {
	testCases := []struct {
		name       string
		maxRetries int
		wantCalls  int32
		body       string
	}{
		{
			name:       "MAX_RETRIES=1 with malformed JSON",
			maxRetries: 1,
			wantCalls:  2,
			body:       "{invalid json",
		},
		{
			name:       "MAX_RETRIES=2 with truncated JSON",
			maxRetries: 2,
			wantCalls:  3,
			body:       `{"id": "msg_123", "content":`,
		},
		{
			name:       "MAX_RETRIES=3 with garbage data",
			maxRetries: 3,
			wantCalls:  4,
			body:       "not json at all!!!",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			var upstreamCallCount atomic.Int32

			// Create upstream that returns 200 with invalid JSON every time
			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				upstreamCallCount.Add(1)
				w.WriteHeader(http.StatusOK)
				w.Write([]byte(tc.body))
			}))
			defer upstream.Close()

			handler := createTestProxyHandler(t, upstream.URL, tc.maxRetries)

			req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody()))
			req.Header.Set("Content-Type", "application/json")
			w := httptest.NewRecorder()

			start := time.Now()
			handler.ServeHTTP(w, req)
			duration := time.Since(start)

			// Verify 502 Bad Gateway is returned
			if w.Code != http.StatusBadGateway {
				t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
			}

			// Verify exact upstream call count
			actualCalls := upstreamCallCount.Load()
			if actualCalls != tc.wantCalls {
				t.Errorf("Expected %d upstream calls (1 initial + %d retries), got %d",
					tc.wantCalls, tc.maxRetries, actualCalls)
			}

			t.Logf("PASS: Invalid JSON body triggered %d retries, %d upstream calls in %v, returned 502",
				tc.maxRetries, actualCalls, duration)
		})
	}
}

// TestResponseValidation_EmptyStreaming tests that streaming responses opening with zero bytes
// triggers retry up to MAX_RETRIES times, then returns 502 Bad Gateway.
func TestResponseValidation_EmptyStreaming(t *testing.T) {
	testCases := []struct {
		name       string
		maxRetries int
		wantCalls  int32
	}{
		{
			name:       "MAX_RETRIES=1",
			maxRetries: 1,
			wantCalls:  2,
		},
		{
			name:       "MAX_RETRIES=2",
			maxRetries: 2,
			wantCalls:  3,
		},
		{
			name:       "MAX_RETRIES=4",
			maxRetries: 4,
			wantCalls:  5,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			var upstreamCallCount atomic.Int32

			// Create upstream that returns 200 with empty streaming response
			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				upstreamCallCount.Add(1)
				w.Header().Set("Content-Type", "text/event-stream")
				w.WriteHeader(http.StatusOK)
				// Empty streaming - write no data
			}))
			defer upstream.Close()

			handler := createTestProxyHandler(t, upstream.URL, tc.maxRetries)

			req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody()))
			req.Header.Set("Content-Type", "application/json")
			w := httptest.NewRecorder()

			start := time.Now()
			handler.ServeHTTP(w, req)
			duration := time.Since(start)

			// Verify 502 Bad Gateway is returned
			if w.Code != http.StatusBadGateway {
				t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
			}

			// Verify exact upstream call count
			actualCalls := upstreamCallCount.Load()
			if actualCalls != tc.wantCalls {
				t.Errorf("Expected %d upstream calls (1 initial + %d retries), got %d",
					tc.wantCalls, tc.maxRetries, actualCalls)
			}

			// Verify exponential backoff took reasonable time
			minExpectedDuration := time.Duration(1<<uint(tc.maxRetries-1)) * time.Second
			if tc.maxRetries > 0 && duration < minExpectedDuration {
				t.Logf("WARNING: Duration %v is less than expected minimum %v", duration, minExpectedDuration)
			}

			t.Logf("PASS: Empty streaming triggered %d retries, %d upstream calls in %v, returned 502",
				tc.maxRetries, actualCalls, duration)
		})
	}
}

// TestResponseValidation_MetricsEmptyJSON tests that zai_proxy_upstream_errors_total
// metric increments with correct error_type="truncated_response" for empty JSON retries.
func TestResponseValidation_MetricsEmptyJSON(t *testing.T) {
	// Reset metrics before test
	retryAttempts.Reset()
	upstreamErrors.Reset()

	maxRetries := 2
	var upstreamCallCount atomic.Int32

	// Create upstream that returns 200 with empty body
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upstreamCallCount.Add(1)
		w.WriteHeader(http.StatusOK)
		// Empty body
	}))
	defer upstream.Close()

	handler := createTestProxyHandler(t, upstream.URL, maxRetries)

	req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody()))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()

	handler.ServeHTTP(w, req)

	// Verify zai_proxy_upstream_errors_total metric increments
	// Each failed attempt (except the final one that returns 502) increments the metric
	// With maxRetries=2, we have: initial attempt fails (error count 1), retry 1 fails (error count 2), retry 2 fails (error count 3, then returns 502)
	// So we expect 3 error increments
	expectedErrorCount := float64(maxRetries + 1)
	errorMetric := upstreamErrors.WithLabelValues("truncated_response", "test")
	errorCount := testutil.ToFloat64(errorMetric)
	if errorCount != expectedErrorCount {
		t.Errorf("Expected zai_proxy_upstream_errors_total{{error_type=\"truncated_response\"}} = %f, got %f",
			expectedErrorCount, errorCount)
	}

	// Verify zai_proxy_retry_attempts_total metric increments
	// Each retry increments the retry attempts metric
	// With maxRetries=2, we expect 2 retry attempts (not 3 - the final attempt doesn't retry, it returns 502)
	expectedRetryCount := float64(maxRetries)
	retryMetric := retryAttempts.WithLabelValues("truncated_response", "test")
	retryCount := testutil.ToFloat64(retryMetric)
	if retryCount != expectedRetryCount {
		t.Errorf("Expected zai_proxy_retry_attempts_total{{reason=\"truncated_response\"}} = %f, got %f",
			expectedRetryCount, retryCount)
	}

	// Verify upstream call count
	expectedCalls := int32(maxRetries + 1)
	actualCalls := upstreamCallCount.Load()
	if actualCalls != expectedCalls {
		t.Errorf("Expected %d upstream calls, got %d", expectedCalls, actualCalls)
	}

	// Verify 502 returned
	if w.Code != http.StatusBadGateway {
		t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
	}

	t.Logf("PASS: Metrics verified - upstream_errors=%f, retry_attempts=%f, calls=%d, status=%d",
		errorCount, retryCount, actualCalls, w.Code)
}

// TestResponseValidation_MetricsInvalidJSON tests that zai_proxy_upstream_errors_total
// metric increments with correct error_type="truncated_response" for invalid JSON retries.
func TestResponseValidation_MetricsInvalidJSON(t *testing.T) {
	// Reset metrics before test
	retryAttempts.Reset()
	upstreamErrors.Reset()

	maxRetries := 3
	var upstreamCallCount atomic.Int32

	// Create upstream that returns 200 with invalid JSON
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upstreamCallCount.Add(1)
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("{this is not valid json"))
	}))
	defer upstream.Close()

	handler := createTestProxyHandler(t, upstream.URL, maxRetries)

	req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody()))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()

	handler.ServeHTTP(w, req)

	// Verify upstream errors metric
	expectedErrorCount := float64(maxRetries + 1)
	errorMetric := upstreamErrors.WithLabelValues("truncated_response", "test")
	errorCount := testutil.ToFloat64(errorMetric)
	if errorCount != expectedErrorCount {
		t.Errorf("Expected zai_proxy_upstream_errors_total{{error_type=\"truncated_response\"}} = %f, got %f",
			expectedErrorCount, errorCount)
	}

	// Verify retry attempts metric
	expectedRetryCount := float64(maxRetries)
	retryMetric := retryAttempts.WithLabelValues("truncated_response", "test")
	retryCount := testutil.ToFloat64(retryMetric)
	if retryCount != expectedRetryCount {
		t.Errorf("Expected zai_proxy_retry_attempts_total{{reason=\"truncated_response\"}} = %f, got %f",
			expectedRetryCount, retryCount)
	}

	// Verify upstream call count
	expectedCalls := int32(maxRetries + 1)
	actualCalls := upstreamCallCount.Load()
	if actualCalls != expectedCalls {
		t.Errorf("Expected %d upstream calls, got %d", expectedCalls, actualCalls)
	}

	// Verify 502 returned
	if w.Code != http.StatusBadGateway {
		t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
	}

	t.Logf("PASS: Metrics verified - upstream_errors=%f, retry_attempts=%f, calls=%d, status=%d",
		errorCount, retryCount, actualCalls, w.Code)
}

// TestResponseValidation_MetricsEmptyStreaming tests that zai_proxy_upstream_errors_total
// metric increments with correct error_type="empty_streaming" for empty streaming retries.
func TestResponseValidation_MetricsEmptyStreaming(t *testing.T) {
	// Reset metrics before test
	retryAttempts.Reset()
	upstreamErrors.Reset()

	maxRetries := 2
	var upstreamCallCount atomic.Int32

	// Create upstream that returns 200 with empty streaming response
	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		upstreamCallCount.Add(1)
		w.Header().Set("Content-Type", "text/event-stream")
		w.WriteHeader(http.StatusOK)
		// Empty streaming response
	}))
	defer upstream.Close()

	handler := createTestProxyHandler(t, upstream.URL, maxRetries)

	req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody()))
	req.Header.Set("Content-Type", "application/json")
	w := httptest.NewRecorder()

	handler.ServeHTTP(w, req)

	// Verify upstream errors metric with error_type="empty_streaming"
	expectedErrorCount := float64(maxRetries + 1)
	errorMetric := upstreamErrors.WithLabelValues("empty_streaming", "test")
	errorCount := testutil.ToFloat64(errorMetric)
	if errorCount != expectedErrorCount {
		t.Errorf("Expected zai_proxy_upstream_errors_total{{error_type=\"empty_streaming\"}} = %f, got %f",
			expectedErrorCount, errorCount)
	}

	// Verify retry attempts metric with reason="empty_streaming"
	expectedRetryCount := float64(maxRetries)
	retryMetric := retryAttempts.WithLabelValues("empty_streaming", "test")
	retryCount := testutil.ToFloat64(retryMetric)
	if retryCount != expectedRetryCount {
		t.Errorf("Expected zai_proxy_retry_attempts_total{{reason=\"empty_streaming\"}} = %f, got %f",
			expectedRetryCount, retryCount)
	}

	// Verify upstream call count
	expectedCalls := int32(maxRetries + 1)
	actualCalls := upstreamCallCount.Load()
	if actualCalls != expectedCalls {
		t.Errorf("Expected %d upstream calls, got %d", expectedCalls, actualCalls)
	}

	// Verify 502 returned
	if w.Code != http.StatusBadGateway {
		t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
	}

	t.Logf("PASS: Metrics verified - upstream_errors=%f, retry_attempts=%f, calls=%d, status=%d",
		errorCount, retryCount, actualCalls, w.Code)
}

// TestResponseValidation_ValidResponsePassesThrough tests that valid responses
// (both non-streaming and streaming) pass through without triggering retries.
func TestResponseValidation_ValidResponsePassesThrough(t *testing.T) {
	t.Run("valid non-streaming response", func(t *testing.T) {
		var upstreamCallCount atomic.Int32

		upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			upstreamCallCount.Add(1)
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusOK)
			json.NewEncoder(w).Encode(map[string]interface{}{
				"id":      "msg_test123",
				"type":    "message",
				"role":    "assistant",
				"content": []map[string]string{{"type": "text", "text": "Hello!"}},
				"usage":   map[string]int{"input_tokens": 10, "output_tokens": 5},
			})
		}))
		defer upstream.Close()

		handler := createTestProxyHandler(t, upstream.URL, 3)

		req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody()))
		req.Header.Set("Content-Type", "application/json")
		w := httptest.NewRecorder()

		handler.ServeHTTP(w, req)

		// Verify success response returned
		if w.Code != http.StatusOK {
			t.Errorf("Expected status 200 OK, got %d", w.Code)
		}

		// Verify response body is valid JSON
		body := w.Body.String()
		if !json.Valid([]byte(body)) {
			t.Error("Response body is not valid JSON")
		}

		// Verify exactly 1 upstream call (no retries)
		if upstreamCallCount.Load() != 1 {
			t.Errorf("Expected 1 upstream call (no retry), got %d", upstreamCallCount.Load())
		}

		t.Logf("PASS: Valid non-streaming response passed through without retry")
	})

	t.Run("valid streaming response", func(t *testing.T) {
		var upstreamCallCount atomic.Int32

		upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			upstreamCallCount.Add(1)
			w.Header().Set("Content-Type", "text/event-stream")
			w.WriteHeader(http.StatusOK)
			// Send some SSE data
			w.Write([]byte("data: {\"type\":\"message_start\"}\n\n"))
			w.Write([]byte("data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"Hello\"}}\n\n"))
			w.Write([]byte("data: [DONE]\n\n"))
		}))
		defer upstream.Close()

		handler := createTestProxyHandler(t, upstream.URL, 3)

		req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody()))
		req.Header.Set("Content-Type", "application/json")
		w := httptest.NewRecorder()

		handler.ServeHTTP(w, req)

		// Verify success response returned
		if w.Code != http.StatusOK {
			t.Errorf("Expected status 200 OK, got %d", w.Code)
		}

		// Verify exactly 1 upstream call (no retries)
		if upstreamCallCount.Load() != 1 {
			t.Errorf("Expected 1 upstream call (no retry), got %d", upstreamCallCount.Load())
		}

		t.Logf("PASS: Valid streaming response passed through without retry")
	})
}

// TestResponseValidation_VariousInvalidBodies tests various forms of invalid response bodies
// to ensure they all trigger retries correctly.
func TestResponseValidation_VariousInvalidBodies(t *testing.T) {
	testCases := []struct {
		name  string
		body  string
		valid bool
	}{
		{
			name:  "empty body",
			body:  "",
			valid: false,
		},
		{
			name:  "single opening brace",
			body:  "{",
			valid: false,
		},
		{
			name:  "unclosed object",
			body:  `{"key": "value"`,
			valid: false,
		},
		{
			name:  "unclosed array",
			body:  `["item1", "item2"`,
			valid: false,
		},
		{
			name:  "random text",
			body:  "this is not json at all",
			valid: false,
		},
		{
			name:  "partial JSON with trailing comma",
			body:  `{"id": "msg123", "content": "text",}`,
			valid: false,
		},
		{
			name:  "valid JSON object",
			body:  `{"id": "msg123", "content": "hello"}`,
			valid: true,
		},
		{
			name:  "valid JSON array",
			body:  `[{"id": "msg123"}, {"id": "msg456"}]`,
			valid: true,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			var upstreamCallCount atomic.Int32
			maxRetries := 1

			upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				upstreamCallCount.Add(1)
				w.WriteHeader(http.StatusOK)
				if len(tc.body) > 0 {
					w.Write([]byte(tc.body))
				}
			}))
			defer upstream.Close()

			handler := createTestProxyHandler(t, upstream.URL, maxRetries)

			req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createNonStreamingRequestBody()))
			req.Header.Set("Content-Type", "application/json")
			w := httptest.NewRecorder()

			handler.ServeHTTP(w, req)

			if tc.valid {
				// Valid JSON should pass through
				if w.Code != http.StatusOK {
					t.Errorf("Expected status 200 OK for valid JSON, got %d", w.Code)
				}
				if upstreamCallCount.Load() != 1 {
					t.Errorf("Expected 1 upstream call for valid JSON (no retry), got %d", upstreamCallCount.Load())
				}
			} else {
				// Invalid JSON should trigger retry and return 502
				if w.Code != http.StatusBadGateway {
					t.Errorf("Expected status 502 Bad Gateway for invalid JSON, got %d", w.Code)
				}
				expectedCalls := int32(maxRetries + 1)
				if upstreamCallCount.Load() != expectedCalls {
					t.Errorf("Expected %d upstream calls for invalid JSON (1 initial + %d retries), got %d",
						expectedCalls, maxRetries, upstreamCallCount.Load())
				}
			}

			t.Logf("PASS: %s - status=%d, calls=%d", tc.name, w.Code, upstreamCallCount.Load())
		})
	}
}

// TestResponseValidation_StreamingZeroBytesPeek tests that the streaming zero-byte detection
// works by peeking at the first chunk of the response body.
func TestResponseValidation_StreamingZeroBytesPeek(t *testing.T) {
	t.Run("zero bytes on first peek", func(t *testing.T) {
		var upstreamCallCount atomic.Int32
		maxRetries := 1

		upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			upstreamCallCount.Add(1)
			w.Header().Set("Content-Type", "text/event-stream")
			w.WriteHeader(http.StatusOK)
			// Close immediately without writing anything
		}))
		defer upstream.Close()

		handler := createTestProxyHandler(t, upstream.URL, maxRetries)

		req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody()))
		req.Header.Set("Content-Type", "application/json")
		w := httptest.NewRecorder()

		handler.ServeHTTP(w, req)

		// Verify retry was triggered and 502 returned
		if w.Code != http.StatusBadGateway {
			t.Errorf("Expected status 502 Bad Gateway, got %d", w.Code)
		}
		expectedCalls := int32(maxRetries + 1)
		if upstreamCallCount.Load() != expectedCalls {
			t.Errorf("Expected %d upstream calls, got %d", expectedCalls, upstreamCallCount.Load())
		}

		t.Logf("PASS: Zero bytes on first peek triggered retry, %d calls, status 502", upstreamCallCount.Load())
	})

	t.Run("non-zero bytes on first peek (no retry)", func(t *testing.T) {
		var upstreamCallCount atomic.Int32

		upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			upstreamCallCount.Add(1)
			w.Header().Set("Content-Type", "text/event-stream")
			w.WriteHeader(http.StatusOK)
			// Write at least one byte immediately
			w.Write([]byte("d"))
		}))
		defer upstream.Close()

		handler := createTestProxyHandler(t, upstream.URL, 3)

		req := httptest.NewRequest("POST", "/v1/messages", strings.NewReader(createStreamingRequestBody()))
		req.Header.Set("Content-Type", "application/json")
		w := httptest.NewRecorder()

		handler.ServeHTTP(w, req)

		// Verify no retry - response passes through
		if w.Code != http.StatusOK {
			t.Errorf("Expected status 200 OK, got %d", w.Code)
		}
		if upstreamCallCount.Load() != 1 {
			t.Errorf("Expected 1 upstream call (no retry), got %d", upstreamCallCount.Load())
		}

		t.Logf("PASS: Non-zero bytes on first peek, no retry, status 200")
	})
}