package audit

import (
	"bytes"
	"io"
	"strings"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestLimitWriter_Write(t *testing.T) {
	tests := []struct {
		name          string
		limit         int64
		writes        []string
		wantNs        []int
		wantBuf       string
		wantTruncated bool
	}{
		{
			name:          "under limit",
			limit:         100,
			writes:        []string{"hello world"},
			wantNs:        []int{len("hello world")},
			wantBuf:       "hello world",
			wantTruncated: false,
		},
		{
			name:          "exactly at limit",
			limit:         10,
			writes:        []string{"0123456789"},
			wantNs:        []int{10},
			wantBuf:       "0123456789",
			wantTruncated: false,
		},
		{
			name:          "exceeds limit",
			limit:         5,
			writes:        []string{"hello world"},
			wantNs:        []int{len("hello world")},
			wantBuf:       "hello",
			wantTruncated: true,
		},
		{
			name:          "multiple writes",
			limit:         10,
			writes:        []string{"hello ", "world!"},
			wantNs:        []int{6, 6},
			wantBuf:       "hello worl",
			wantTruncated: true,
		},
		{
			name:          "writes after limit",
			limit:         5,
			writes:        []string{"12345", "extra"},
			wantNs:        []int{5, 5},
			wantBuf:       "12345",
			wantTruncated: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			buf := &bytes.Buffer{}
			lw := newLimitWriter(buf, tt.limit)

			gotNs := make([]int, 0, len(tt.writes))
			for _, data := range tt.writes {
				n, err := lw.Write([]byte(data))
				require.NoError(t, err)
				gotNs = append(gotNs, n)
			}

			assert.Equal(t, tt.wantNs, gotNs)
			assert.Equal(t, tt.wantBuf, buf.String())
			assert.Equal(t, tt.wantTruncated, lw.truncated)
		})
	}
}

func TestBodyWrapper_ReadAllAndAuditCapture(t *testing.T) {
	tests := []struct {
		name    string
		content string
		maxSize int
	}{
		{
			name:    "small body",
			content: "small body content",
			maxSize: 256,
		},
		{
			name:    "empty body",
			content: "",
			maxSize: 256,
		},
		{
			name:    "large body truncates audit",
			content: strings.Repeat("x", 1000),
			maxSize: 256,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			body := io.NopCloser(strings.NewReader(tt.content))
			wrapper := NewBodyWrapper(body, tt.maxSize)

			data, err := io.ReadAll(wrapper)
			require.NoError(t, err)
			assert.Equal(t, tt.content, string(data))

			wantAuditLen := len(tt.content)
			if wantAuditLen > tt.maxSize {
				wantAuditLen = tt.maxSize
			}
			assert.Equal(t, tt.content[:wantAuditLen], string(wrapper.AuditBytes()))
			assert.Equal(t, len(tt.content) > tt.maxSize, wrapper.Truncated())
		})
	}
}

func TestBodyWrapper_DownstreamGetsFullBody(t *testing.T) {
	// This is the critical test: ensure downstream gets full body
	// even when audit captures only a portion
	largeContent := strings.Repeat("A", 10000)
	body := io.NopCloser(strings.NewReader(largeContent))
	wrapper := NewBodyWrapper(body, 256)

	// Simulate downstream reading in chunks (like httputil.ReverseProxy does)
	var downstream bytes.Buffer
	buf := make([]byte, 1024)
	for {
		n, err := wrapper.Read(buf)
		if n > 0 {
			downstream.Write(buf[:n])
		}
		if err == io.EOF {
			break
		}
		require.NoError(t, err)
	}

	// Verify downstream got the FULL body
	assert.Equal(t, 10000, downstream.Len())
	assert.Equal(t, largeContent, downstream.String())

	// But audit only got the truncated portion
	assert.Equal(t, 256, len(wrapper.AuditBytes()))
}

func TestBodyWrapper_Close(t *testing.T) {
	closeCount := 0
	body := &mockReadCloser{
		Reader: strings.NewReader("test"),
		closeFunc: func() error {
			closeCount++
			return nil
		},
	}

	wrapper := NewBodyWrapper(body, 256)

	// Read some data
	_, _ = io.ReadAll(wrapper)

	// Close should call underlying Close
	err := wrapper.Close()
	require.NoError(t, err)
	assert.Equal(t, 1, closeCount)
}

func TestBodyWrapper_AuditBytesBeforeFullRead(t *testing.T) {
	content := "partial read test"
	body := io.NopCloser(strings.NewReader(content))
	wrapper := NewBodyWrapper(body, 256)

	// Read only 7 bytes
	buf := make([]byte, 7)
	n, err := wrapper.Read(buf)
	require.NoError(t, err)
	assert.Equal(t, 7, n)

	// AuditBytes should return what's been read so far
	auditData := wrapper.AuditBytes()
	assert.Equal(t, "partial", string(auditData))
}

// mockReadCloser is a test helper
type mockReadCloser struct {
	io.Reader
	closeFunc func() error
}

func (m *mockReadCloser) Close() error {
	if m.closeFunc != nil {
		return m.closeFunc()
	}
	return nil
}