package aws

import (
	"bytes"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/stephnangue/warden/framework"
	"github.com/stephnangue/warden/logger"
	"github.com/stephnangue/warden/provider/sdk/sigv4"
)

// createTestLogger creates a logger for testing that discards output
func createTestLogger() *logger.GatedLogger {
	config := &logger.Config{
		Level:   logger.TraceLevel,
		Format:  logger.DefaultFormat,
		Outputs: []io.Writer{io.Discard},
	}
	gl, _ := logger.NewGatedLogger(config, logger.GatedWriterConfig{
		Underlying:   io.Discard,
		InitialState: logger.GateOpen,
	})
	return gl
}

func TestExtractFromAuthHeader(t *testing.T) {
	tests := []struct {
		name              string
		authHeader        string
		expectedService   string
		expectedRegion    string
		expectedAccessKey string
		expectError       bool
	}{
		{
			name:              "valid S3 header",
			authHeader:        "AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/20130524/us-east-1/s3/aws4_request, SignedHeaders=host;x-amz-content-sha256;x-amz-date, Signature=abc123",
			expectedService:   "s3",
			expectedRegion:    "us-east-1",
			expectedAccessKey: "AKIAIOSFODNN7EXAMPLE",
			expectError:       false,
		},
		{
			name:              "valid EC2 header",
			authHeader:        "AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/20230101/eu-west-1/ec2/aws4_request, SignedHeaders=host;x-amz-date, Signature=def456",
			expectedService:   "ec2",
			expectedRegion:    "eu-west-1",
			expectedAccessKey: "AKIAIOSFODNN7EXAMPLE",
			expectError:       false,
		},
		{
			name:              "valid IAM header",
			authHeader:        "AWS4-HMAC-SHA256 Credential=ASIAXYZ123456789012/20230615/us-east-1/iam/aws4_request, SignedHeaders=host, Signature=xyz789",
			expectedService:   "iam",
			expectedRegion:    "us-east-1",
			expectedAccessKey: "ASIAXYZ123456789012",
			expectError:       false,
		},
		{
			name:        "empty header",
			authHeader:  "",
			expectError: true,
		},
		{
			name:        "invalid format",
			authHeader:  "Bearer token123",
			expectError: true,
		},
		{
			name:        "missing credential",
			authHeader:  "AWS4-HMAC-SHA256 SignedHeaders=host, Signature=abc",
			expectError: true,
		},
		{
			name:        "malformed credential",
			authHeader:  "AWS4-HMAC-SHA256 Credential=incomplete, SignedHeaders=host",
			expectError: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			service, region, accessKey, err := sigv4.ExtractFromAuthHeader(tt.authHeader)

			if tt.expectError {
				if err == nil {
					t.Error("Expected error, got nil")
				}
				return
			}

			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}
			if service != tt.expectedService {
				t.Errorf("service = %s, want %s", service, tt.expectedService)
			}
			if region != tt.expectedRegion {
				t.Errorf("region = %s, want %s", region, tt.expectedRegion)
			}
			if accessKey != tt.expectedAccessKey {
				t.Errorf("accessKey = %s, want %s", accessKey, tt.expectedAccessKey)
			}
		})
	}
}

func TestComputePayloadHash(t *testing.T) {
	tests := []struct {
		name     string
		body     []byte
		expected string
	}{
		{
			name:     "empty body",
			body:     []byte{},
			expected: "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
		},
		{
			name:     "simple text",
			body:     []byte("hello"),
			expected: "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824",
		},
		{
			name:     "json body",
			body:     []byte(`{"key":"value"}`),
			expected: "e43abcf3375244839c012f9633f95862d232a95b00d5bc7348b3098b9fed7f32",
		},
		{
			name:     "binary data",
			body:     []byte{0x00, 0x01, 0x02, 0x03},
			expected: "054edec1d0211f624fed0cbca9d4f9400b0e491c43742af2c5b0abebf0c990d8",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := sigv4.ComputePayloadHash(tt.body)
			if got != tt.expected {
				t.Errorf("sigv4.ComputePayloadHash() = %s, want %s", got, tt.expected)
			}
		})
	}
}

func TestParseAWSDate(t *testing.T) {
	tests := []struct {
		name        string
		dateStr     string
		expectError bool
		expected    time.Time
	}{
		{
			name:        "valid date",
			dateStr:     "20230615T120000Z",
			expectError: false,
			expected:    time.Date(2023, 6, 15, 12, 0, 0, 0, time.UTC),
		},
		{
			name:        "midnight",
			dateStr:     "20230101T000000Z",
			expectError: false,
			expected:    time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC),
		},
		{
			name:        "end of day",
			dateStr:     "20231231T235959Z",
			expectError: false,
			expected:    time.Date(2023, 12, 31, 23, 59, 59, 0, time.UTC),
		},
		{
			name:        "invalid format",
			dateStr:     "2023-06-15T12:00:00Z",
			expectError: true,
		},
		{
			name:        "empty string",
			dateStr:     "",
			expectError: true,
		},
		{
			name:        "missing Z suffix",
			dateStr:     "20230615T120000",
			expectError: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := sigv4.ParseAWSDate(tt.dateStr)

			if tt.expectError {
				if err == nil {
					t.Error("Expected error, got nil")
				}
				return
			}

			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}
			if !got.Equal(tt.expected) {
				t.Errorf("sigv4.ParseAWSDate() = %v, want %v", got, tt.expected)
			}
		})
	}
}

func TestAwsBackend_ReadRequestBody(t *testing.T) {
	tests := []struct {
		name        string
		body        string
		maxBodySize int64
		expectError bool
		expected    string
	}{
		{
			name:        "normal body",
			body:        "test body content",
			maxBodySize: 1024,
			expectError: false,
			expected:    "test body content",
		},
		{
			name:        "empty body",
			body:        "",
			maxBodySize: 1024,
			expectError: false,
			expected:    "",
		},
		{
			name:        "body at limit",
			body:        "12345",
			maxBodySize: 10,
			expectError: false,
			expected:    "12345",
		},
		{
			name:        "body exceeds limit",
			body:        "12345678901",
			maxBodySize: 10,
			expectError: true,
		},
		{
			name:        "no limit set",
			body:        "any content",
			maxBodySize: 0,
			expectError: false,
			expected:    "any content",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodPost, "/test", bytes.NewBufferString(tt.body))

			got, err := sigv4.ReadRequestBody(req, tt.maxBodySize)

			if tt.expectError {
				if err == nil {
					t.Error("Expected error, got nil")
				}
				return
			}

			if err != nil {
				t.Fatalf("Unexpected error: %v", err)
			}
			if string(got) != tt.expected {
				t.Errorf("readRequestBody() = %s, want %s", string(got), tt.expected)
			}
		})
	}
}

func TestAwsBackend_ReadRequestBody_NilBody(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/test", nil)
	req.Body = nil

	got, err := sigv4.ReadRequestBody(req, 1024)

	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if got != nil {
		t.Errorf("Expected nil, got %v", got)
	}
}

func TestAwsBackend_RestoreRequestBody(t *testing.T) {
	tests := []struct {
		name      string
		bodyBytes []byte
	}{
		{
			name:      "normal body",
			bodyBytes: []byte("test content"),
		},
		{
			name:      "empty body",
			bodyBytes: []byte{},
		},
		{
			name:      "nil body",
			bodyBytes: nil,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req := httptest.NewRequest(http.MethodPost, "/test", nil)

			sigv4.RestoreRequestBody(req, tt.bodyBytes)

			if len(tt.bodyBytes) > 0 {
				// Body should be restored
				got, err := io.ReadAll(req.Body)
				if err != nil {
					t.Fatalf("Failed to read body: %v", err)
				}
				if string(got) != string(tt.bodyBytes) {
					t.Errorf("Body = %s, want %s", string(got), string(tt.bodyBytes))
				}
				if req.ContentLength != int64(len(tt.bodyBytes)) {
					t.Errorf("ContentLength = %d, want %d", req.ContentLength, len(tt.bodyBytes))
				}
			} else {
				// Body should be nil
				if req.Body != nil {
					t.Error("Expected nil body")
				}
				if req.ContentLength != 0 {
					t.Errorf("ContentLength = %d, want 0", req.ContentLength)
				}
			}
		})
	}
}

func TestAwsBackend_NormalizeRequest_HopByHopHeaders(t *testing.T) {
	b := &awsBackend{
		StreamingBackend: &framework.StreamingBackend{
			Logger: createTestLogger(),
		},
	}

	req := httptest.NewRequest(http.MethodPost, "/test", nil)

	// Add hop-by-hop headers
	req.Header.Set("Connection", "keep-alive")
	req.Header.Set("Keep-Alive", "timeout=5")
	req.Header.Set("Transfer-Encoding", "chunked")
	req.Header.Set("Upgrade", "h2c")

	// Add proxy headers
	req.Header.Set("X-Forwarded-For", "192.168.1.1")
	req.Header.Set("X-Forwarded-Host", "original.host.com")
	req.Header.Set("X-Real-Ip", "10.0.0.1")

	// Add headers that should be kept
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-Amz-Date", "20230615T120000Z")
	req.Header.Set("Authorization", "AWS4-HMAC-SHA256 ...")

	body := []byte("test body")
	result := b.normalizeRequest(req, body)

	// Body should pass through unchanged (no aws-chunked)
	if string(result) != "test body" {
		t.Errorf("expected body unchanged, got %q", string(result))
	}

	// Verify hop-by-hop headers are removed
	for _, h := range []string{"Connection", "Keep-Alive", "Transfer-Encoding", "Upgrade"} {
		if req.Header.Get(h) != "" {
			t.Errorf("Header %s should have been removed", h)
		}
	}

	// Verify proxy headers are removed
	for _, h := range []string{"X-Forwarded-For", "X-Forwarded-Host", "X-Real-Ip"} {
		if req.Header.Get(h) != "" {
			t.Errorf("Header %s should have been removed", h)
		}
	}

	// Verify other headers are kept
	if req.Header.Get("Content-Type") != "application/json" {
		t.Error("Content-Type header should be preserved")
	}
	if req.Header.Get("X-Amz-Date") != "20230615T120000Z" {
		t.Error("X-Amz-Date header should be preserved")
	}
	if req.Header.Get("Authorization") == "" {
		t.Error("Authorization header should be preserved")
	}
}

func TestAwsBackend_NormalizeRequest_ConnectionListed(t *testing.T) {
	b := &awsBackend{
		StreamingBackend: &framework.StreamingBackend{
			Logger: createTestLogger(),
		},
	}

	req := httptest.NewRequest(http.MethodPost, "/test", nil)

	// Connection header lists additional headers to remove
	req.Header.Set("Connection", "Custom-Header, Another-Header")
	req.Header.Set("Custom-Header", "value1")
	req.Header.Set("Another-Header", "value2")
	req.Header.Set("Keep-Header", "should-stay")

	b.normalizeRequest(req, []byte{})

	// Connection should be removed (it's a hop-by-hop header)
	if req.Header.Get("Connection") != "" {
		t.Error("Connection header should have been removed")
	}

	// Other headers should remain
	if req.Header.Get("Keep-Header") != "should-stay" {
		t.Error("Keep-Header should be preserved")
	}
}

func TestAwsBackend_NormalizeRequest_AWSChunked(t *testing.T) {
	b := &awsBackend{
		StreamingBackend: &framework.StreamingBackend{
			Logger: createTestLogger(),
		},
	}

	req := httptest.NewRequest(http.MethodPut, "/test", nil)
	req.Header.Set("Content-Encoding", "aws-chunked")
	req.Header.Set("X-Amz-Decoded-Content-Length", "5")
	req.Header.Set("X-Amz-Content-Sha256", "STREAMING-AWS4-HMAC-SHA256-PAYLOAD")
	req.Header.Set("Accept-Encoding", "gzip")
	req.Header.Set("Content-Length", "100")
	req.Header.Set("Expect", "100-continue")

	// aws-chunked encoded body for "hello"
	chunkedBody := []byte("5;chunk-signature=abc\r\nhello\r\n0;chunk-signature=def\r\n\r\n")

	result := b.normalizeRequest(req, chunkedBody)

	// Body should be decoded
	if string(result) != "hello" {
		t.Errorf("expected decoded body 'hello', got %q", string(result))
	}

	// Streaming headers should be removed
	if req.Header.Get("Content-Encoding") != "" {
		t.Error("Content-Encoding should be removed")
	}
	if req.Header.Get("X-Amz-Decoded-Content-Length") != "" {
		t.Error("X-Amz-Decoded-Content-Length should be removed")
	}
	if req.Header.Get("X-Amz-Content-Sha256") != "" {
		t.Error("X-Amz-Content-Sha256 should be removed")
	}
	if req.Header.Get("Accept-Encoding") != "" {
		t.Error("Accept-Encoding should be removed")
	}
	if req.Header.Get("Content-Length") != "" {
		t.Error("Content-Length should be removed")
	}
	if req.Header.Get("Expect") != "" {
		t.Error("Expect should be removed")
	}
}

func TestAwsBackend_NormalizeRequest_TrailerHeaders(t *testing.T) {
	b := &awsBackend{
		StreamingBackend: &framework.StreamingBackend{
			Logger: createTestLogger(),
		},
	}

	req := httptest.NewRequest(http.MethodPut, "/test", nil)
	req.Header.Set("X-Amz-Trailer", "x-amz-checksum-crc32")
	req.Header.Set("X-Amz-Sdk-Checksum-Algorithm", "CRC32")
	req.Header.Set("X-Amz-Checksum-Algorithm", "CRC32")
	req.Header.Set("X-Amz-Checksum-Crc32", "abc123")
	req.Header.Set("Content-Type", "application/octet-stream")

	b.normalizeRequest(req, []byte("data"))

	// Trailer-related headers should be removed
	if req.Header.Get("X-Amz-Trailer") != "" {
		t.Error("X-Amz-Trailer should be removed")
	}
	if req.Header.Get("X-Amz-Sdk-Checksum-Algorithm") != "" {
		t.Error("X-Amz-Sdk-Checksum-Algorithm should be removed")
	}
	if req.Header.Get("X-Amz-Checksum-Algorithm") != "" {
		t.Error("X-Amz-Checksum-Algorithm should be removed")
	}
	if req.Header.Get("X-Amz-Checksum-Crc32") != "" {
		t.Error("X-Amz-Checksum-Crc32 should be removed")
	}

	// Non-trailer headers should be kept
	if req.Header.Get("Content-Type") != "application/octet-stream" {
		t.Error("Content-Type should be preserved")
	}
}

func BenchmarkExtractFromAuthHeader(b *testing.B) {
	authHeader := "AWS4-HMAC-SHA256 Credential=AKIAIOSFODNN7EXAMPLE/20130524/us-east-1/s3/aws4_request, SignedHeaders=host;x-amz-content-sha256;x-amz-date, Signature=abc123"

	b.ResetTimer()
	for b.Loop() {
		sigv4.ExtractFromAuthHeader(authHeader)
	}
}

func BenchmarkComputePayloadHash(b *testing.B) {
	body := []byte("test body content for hashing benchmark")

	b.ResetTimer()
	for b.Loop() {
		sigv4.ComputePayloadHash(body)
	}
}

func BenchmarkParseAWSDate(b *testing.B) {
	dateStr := "20230615T120000Z"

	b.ResetTimer()
	for b.Loop() {
		sigv4.ParseAWSDate(dateStr)
	}
}

func BenchmarkNormalizeRequest(b *testing.B) {
	backend := &awsBackend{StreamingBackend: &framework.StreamingBackend{Logger: createTestLogger()}}

	b.ResetTimer()
	for b.Loop() {
		req := httptest.NewRequest(http.MethodPost, "/test", nil)
		req.Header.Set("Connection", "keep-alive")
		req.Header.Set("Keep-Alive", "timeout=5")
		req.Header.Set("X-Forwarded-For", "192.168.1.1")
		req.Header.Set("Content-Type", "application/json")

		backend.normalizeRequest(req, []byte{})
	}
}
