internal/api/modules/amp/response_rewriter_test.go

package amp

import (
	"bytes"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/gin-gonic/gin"
)

// mockResponseWriter is a test double for gin.ResponseWriter
type mockResponseWriter struct {
	gin.ResponseWriter
	body       *bytes.Buffer
	headers    http.Header
	statusCode int
	flushed    bool
}

func newMockResponseWriter() *mockResponseWriter {
	return &mockResponseWriter{
		body:    &bytes.Buffer{},
		headers: make(http.Header),
	}
}

func (m *mockResponseWriter) Write(data []byte) (int, error) {
	return m.body.Write(data)
}

func (m *mockResponseWriter) WriteString(s string) (int, error) {
	return m.body.WriteString(s)
}

func (m *mockResponseWriter) Header() http.Header {
	return m.headers
}

func (m *mockResponseWriter) WriteHeader(code int) {
	m.statusCode = code
}

func (m *mockResponseWriter) Status() int {
	return m.statusCode
}

// Flush implements http.Flusher
func (m *mockResponseWriter) Flush() {
	m.flushed = true
}

// TestResponseRewriter_SplitJSONTokensAcrossChunks tests handling of JSON tokens split across chunk boundaries
func TestResponseRewriter_SplitJSONTokensAcrossChunks(t *testing.T) {
	tests := []struct {
		name          string
		originalModel string
		chunks        []string
		expected      string
	}{
		{
			name:          "model field split across chunks",
			originalModel: "original-model",
			chunks: []string{
				`{"mod`,
				`el": "mapped-model", "content": "test"}`,
			},
			expected: `{"model": "original-model", "content": "test"}`,
		},
		{
			name:          "modelVersion split across chunks",
			originalModel: "original-model",
			chunks: []string{
				`{"modelVers`,
				`ion": "mapped-version"}`,
			},
			expected: `{"modelVersion": "original-model"}`,
		},
		{
			name:          "multiple fields with splits",
			originalModel: "original-model",
			chunks: []string{
				`{"model": "mapped`,
				`-model", "modelVersion": "mapped`,
				`-version"}`,
			},
			expected: `{"model": "original-model", "modelVersion": "original-model"}`,
		},
		{
			name:          "empty chunks between data",
			originalModel: "original-model",
			chunks: []string{
				`{"model": "`,
				``, // Empty chunk
				`mapped-model"}`,
			},
			expected: `{"model": "original-model"}`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mock := newMockResponseWriter()
			rw := NewResponseRewriter(mock, tt.originalModel)

			// Simulate streaming response
			for _, chunk := range tt.chunks {
				rw.Write([]byte(chunk))
			}
			rw.Flush()

			result := mock.body.String()
			if result != tt.expected {
				t.Errorf("got %q, want %q", result, tt.expected)
			}
		})
	}
}

// TestResponseRewriter_InvalidMalformedJSON tests handling of invalid/malformed JSON
func TestResponseRewriter_InvalidMalformedJSON(t *testing.T) {
	tests := []struct {
		name          string
		originalModel string
		input         string
		shouldPass    bool // Whether the data should pass through
	}{
		{
			name:          "incomplete JSON",
			originalModel: "original",
			input:         `{"model": "mapped`,
			shouldPass:    true,
		},
		{
			name:          "invalid JSON structure",
			originalModel: "original",
			input:         `{"model": "mapped", invalid}`,
			shouldPass:    true,
		},
		{
			name:          "unclosed string",
			originalModel: "original",
			input:         `{"model": "mapped`,
			shouldPass:    true,
		},
		{
			name:          "binary garbage",
			originalModel: "original",
			input:         string([]byte{0x00, 0x01, 0x02, 0xFF}),
			shouldPass:    true,
		},
		{
			name:          "null bytes in JSON",
			originalModel: "original",
			input:         `{"model": "mapped\u0000model"}`,
			shouldPass:    true,
		},
		{
			name:          "deeply nested valid JSON",
			originalModel: "original",
			input:         `{"a": {"b": {"c": {"d": {"model": "mapped"}}}}}`,
			shouldPass:    true,
		},
		{
			name:          "valid JSON without model field",
			originalModel: "original",
			input:         `{"content": "test", "other": "data"}`,
			shouldPass:    true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mock := newMockResponseWriter()
			rw := NewResponseRewriter(mock, tt.originalModel)

			// For non-streaming, we buffer and flush
			rw.Write([]byte(tt.input))
			rw.Flush()

			result := mock.body.String()

			// Should not panic and should produce some output
			if result == "" && tt.shouldPass {
				t.Error("expected output but got empty string")
			}

			// For valid JSON with model field, verify replacement happened
			if strings.Contains(tt.input, `"model"`) && strings.Contains(tt.input, "mapped") {
				if strings.Contains(result, "mapped") && !strings.Contains(result, `"model": "original"`) {
					t.Logf("Model replacement may not have worked for: %s", tt.name)
				}
			}
		})
	}
}

// TestResponseRewriter_MixedContentTypes tests handling of mixed content types
func TestResponseRewriter_MixedContentTypes(t *testing.T) {
	tests := []struct {
		name          string
		contentType   string
		originalModel string
		input         string
		isStreaming   bool
	}{
		{
			name:          "SSE stream",
			contentType:   "text/event-stream",
			originalModel: "original-model",
			input:         `data: {"model": "mapped-model", "content": "hello"}`,
			isStreaming:   true,
		},
		{
			name:          "JSON response",
			contentType:   "application/json",
			originalModel: "original-model",
			input:         `{"model": "mapped-model", "content": "hello"}`,
			isStreaming:   false,
		},
		{
			name:          "plain text",
			contentType:   "text/plain",
			originalModel: "original-model",
			input:         `model: mapped-model`,
			isStreaming:   false,
		},
		{
			name:          "octet stream",
			contentType:   "application/octet-stream",
			originalModel: "original-model",
			input:         `binary data`,
			isStreaming:   false,
		},
		{
			name:          "multipart form",
			contentType:   "multipart/form-data",
			originalModel: "original-model",
			input:         `--boundary\nContent-Type: application/json\n\n{"model": "mapped-model"}`,
			isStreaming:   false,
		},
		{
			name:          "JSON with stream in body",
			contentType:   "application/json",
			originalModel: "original-model",
			input:         `{"stream": true, "model": "mapped-model"}`,
			isStreaming:   false, // Content-Type determines streaming, not body
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mock := newMockResponseWriter()
			mock.headers.Set("Content-Type", tt.contentType)
			rw := NewResponseRewriter(mock, tt.originalModel)

			// First write triggers streaming detection
			rw.Write([]byte(tt.input))

			if rw.isStreaming != tt.isStreaming {
				t.Errorf("isStreaming = %v, want %v", rw.isStreaming, tt.isStreaming)
			}

			rw.Flush()

			// Verify output exists
			result := mock.body.String()
			if result == "" {
				t.Error("expected output but got empty string")
			}
		})
	}
}

// TestResponseRewriter_FallbackStrategy tests fallback when rewriting fails
func TestResponseRewriter_FallbackStrategy(t *testing.T) {
	tests := []struct {
		name          string
		originalModel string
		input         []string // Chunks for streaming
		expectError   bool
	}{
		{
			name:          "empty original model - pass through",
			originalModel: "",
			input:         []string{`{"model": "mapped-model"}`},
			expectError:   false,
		},
		{
			name:          "single chunk rewrite success",
			originalModel: "original-model",
			input:         []string{`{"model": "mapped-model"}`},
			expectError:   false,
		},
		{
			name:          "multiple chunks with partial data",
			originalModel: "original-model",
			input: []string{
				`{"model": "mapped`,
				`-model", "content": "test"}`,
			},
			expectError: false,
		},
		{
			name:          "chunk with only whitespace",
			originalModel: "original-model",
			input: []string{
				`   `,
				`{"model": "mapped-model"}`,
			},
			expectError: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mock := newMockResponseWriter()
			rw := NewResponseRewriter(mock, tt.originalModel)

			// Write all chunks
			for _, chunk := range tt.input {
				_, err := rw.Write([]byte(chunk))
				if err != nil && !tt.expectError {
					t.Errorf("unexpected error: %v", err)
				}
			}

			rw.Flush()

			// Should always produce output (fallback to pass-through)
			result := mock.body.String()
			if result == "" && len(tt.input) > 0 {
				t.Error("expected output but got empty string")
			}
		})
	}
}

// TestResponseRewriter_SSEEdgeCases tests Server-Sent Events edge cases
func TestResponseRewriter_SSEEdgeCases(t *testing.T) {
	tests := []struct {
		name          string
		originalModel string
		chunks        []string
		expected      string
	}{
		{
			name:          "SSE with data prefix",
			originalModel: "original-model",
			chunks: []string{
				`data: {"model": "mapped-model"}`,
			},
			expected: `data: {"model": "original-model"}`,
		},
		{
			name:          "SSE multiple data lines",
			originalModel: "original-model",
			chunks: []string{
				"data: {\"model\": \"mapped-model\"}\n\ndata: {\"model\": \"mapped-model2\"}",
			},
			expected: "data: {\"model\": \"original-model\"}\n\ndata: {\"model\": \"original-model\"}",
		},
		{
			name:          "SSE with event and id",
			originalModel: "original-model",
			chunks: []string{
				"id: 1\nevent: message\ndata: {\"model\": \"mapped-model\"}",
			},
			expected: "id: 1\nevent: message\ndata: {\"model\": \"original-model\"}",
		},
		{
			name:          "SSE data split across chunks",
			originalModel: "original-model",
			chunks: []string{
				`data: {"mod`,
				`el": "mapped-model"}`,
			},
			expected: `data: {"model": "original-model"}`,
		},
		{
			name:          "SSE non-JSON data",
			originalModel: "original-model",
			chunks: []string{
				`data: plain text message`,
			},
			expected: `data: plain text message`,
		},
		{
			name:          "SSE empty data line",
			originalModel: "original-model",
			chunks: []string{
				`data:`,
			},
			expected: `data:`,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mock := newMockResponseWriter()
			mock.headers.Set("Content-Type", "text/event-stream")
			rw := NewResponseRewriter(mock, tt.originalModel)

			for _, chunk := range tt.chunks {
				rw.Write([]byte(chunk))
			}

			result := mock.body.String()
			if result != tt.expected {
				t.Errorf("got %q, want %q", result, tt.expected)
			}
		})
	}
}

// TestResponseRewriter_ThinkingBlockSuppression tests the thinking block suppression feature
func TestResponseRewriter_ThinkingBlockSuppression(t *testing.T) {
	tests := []struct {
		name          string
		originalModel string
		input         string
		shouldFilter  bool
	}{
		{
			name:          "tool_use with thinking - should filter",
			originalModel: "original-model",
			input: `{
				"content": [
					{"type": "thinking", "thinking": "secret"},
					{"type": "tool_use", "name": "calculator"}
				]
			}`,
			shouldFilter: true,
		},
		{
			name:          "thinking without tool_use - should not filter",
			originalModel: "original-model",
			input: `{
				"content": [
					{"type": "thinking", "thinking": "not secret"}
				]
			}`,
			shouldFilter: false,
		},
		{
			name:          "no content field",
			originalModel: "original-model",
			input:         `{"model": "mapped-model"}`,
			shouldFilter:  false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mock := newMockResponseWriter()
			rw := NewResponseRewriter(mock, tt.originalModel)

			rw.Write([]byte(tt.input))
			rw.Flush()

			result := mock.body.String()

			// Check if thinking blocks were filtered when they should be
			hasThinking := strings.Contains(result, `"type": "thinking"`)
			hasToolUse := strings.Contains(tt.input, `"type": "tool_use"`)

			if tt.shouldFilter && hasThinking {
				t.Error("thinking blocks should have been filtered but are present")
			}

			if !tt.shouldFilter && hasToolUse && !hasThinking {
				// This might be OK if the thinking was legitimately removed
				t.Logf("Note: thinking blocks were removed even though no tool_use detected")
			}
		})
	}
}

// TestResponseRewriter_ConcurrentWrites tests concurrent write operations
func TestResponseRewriter_ConcurrentWrites(t *testing.T) {
	gin.SetMode(gin.TestMode)

	mock := newMockResponseWriter()
	mock.headers.Set("Content-Type", "application/json")
	rw := NewResponseRewriter(mock, "original-model")

	// Write data that simulates concurrent chunks
	chunks := []string{
		`{"model": "`,
		`mapped`,
		`-model", "`,
		`content": "`,
		`test"}`,
	}

	for _, chunk := range chunks {
		rw.Write([]byte(chunk))
	}
	rw.Flush()

	result := mock.body.String()
	expected := `{"model": "original-model", "content": "test"}`
	if result != expected {
		t.Errorf("got %q, want %q", result, expected)
	}
}

// TestResponseRewriter_LargeResponse tests handling of large responses
func TestResponseRewriter_LargeResponse(t *testing.T) {
	mock := newMockResponseWriter()
	rw := NewResponseRewriter(mock, "original-model")

	// Create a large JSON response
	largeContent := strings.Repeat("a", 100000)
	input := `{"model": "mapped-model", "content": "` + largeContent + `"}`

	rw.Write([]byte(input))
	rw.Flush()

	result := mock.body.String()

	// Should contain the original model
	if !strings.Contains(result, `"model": "original-model"`) {
		t.Error("model replacement failed for large response")
	}

	// Should still contain the large content
	if !strings.Contains(result, largeContent) {
		t.Error("large content was lost")
	}
}

// TestResponseRewriter_EmptyAndNilInputs tests edge cases with empty/nil inputs
func TestResponseRewriter_EmptyAndNilInputs(t *testing.T) {
	tests := []struct {
		name  string
		input []byte
	}{
		{
			name:  "nil input",
			input: nil,
		},
		{
			name:  "empty input",
			input: []byte{},
		},
		{
			name:  "whitespace only",
			input: []byte("   \n\t  "),
		},
		{
			name:  "single brace",
			input: []byte("{"),
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mock := newMockResponseWriter()
			rw := NewResponseRewriter(mock, "original-model")

			// Should not panic
			rw.Write(tt.input)
			rw.Flush()

			// Result should be safe
			_ = mock.body.String()
		})
	}
}

// TestNewResponseRewriter tests the constructor
func TestNewResponseRewriter(t *testing.T) {
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)

	rw := NewResponseRewriter(c.Writer, "test-model")

	if rw == nil {
		t.Fatal("NewResponseRewriter returned nil")
	}

	if rw.originalModel != "test-model" {
		t.Errorf("originalModel = %q, want %q", rw.originalModel, "test-model")
	}

	if rw.body == nil {
		t.Error("body buffer is nil")
	}
}

// TestResponseRewriter_FlushBehavior tests Flush method behavior
func TestResponseRewriter_FlushBehavior(t *testing.T) {
	tests := []struct {
		name       string
		isStreaming bool
		writeData  string
		expectFlush bool
	}{
		{
			name:       "flush non-streaming",
			isStreaming: false,
			writeData:  `{"model": "mapped"}`,
			expectFlush: true,
		},
		{
			name:       "flush streaming",
			isStreaming: true,
			writeData:  `data: {"model": "mapped"}`,
			expectFlush: true,
		},
		{
			name:       "flush empty body",
			isStreaming: false,
			writeData:  "",
			expectFlush: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			mock := newMockResponseWriter()
			if tt.isStreaming {
				mock.headers.Set("Content-Type", "text/event-stream")
			}
			rw := NewResponseRewriter(mock, "original")

			if tt.writeData != "" {
				rw.Write([]byte(tt.writeData))
			}

			// Reset flushed flag
			mock.flushed = false

			rw.Flush()

			if tt.expectFlush && !mock.flushed {
				t.Error("expected Flush to be called on underlying writer")
			}
		})
	}
}