package duckdb

import (
	"errors"
	"os"
	"path/filepath"
	"testing"

	"github.com/DATA-DOG/go-sqlmock"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	deltalake "github.com/tidbcloud/metering/delta-rs-golang"
	"github.com/tidbcloud/metering/pkg/common/config"
	"github.com/tidbcloud/metering/pkg/common/constants"
	"go.uber.org/zap/zaptest"
)

// setupTestDeltaTable creates a test Delta Lake table with sample data
func setupTestDeltaTable(t *testing.T) string {
	// Create a temporary directory for the test table
	tmpDir := t.TempDir()
	tablePath := filepath.Join(tmpDir, "test_table")

	// Create the directory for the Delta table
	err := os.MkdirAll(tablePath, 0755)
	require.NoError(t, err, "Failed to create Delta table directory")

	// Define the table schema
	columns := []deltalake.ColumnDefinition{
		{Name: "cost_time_stamp", Type: "long", Nullable: false},
		{Name: "hour_dimension", Type: "string", Nullable: false},
		{Name: "cost", Type: "double", Nullable: false},
	}

	// Create table options
	options := deltalake.CreateTableOptions{
		TablePath: tablePath,
		SaveMode:  deltalake.WriteModeOverwrite,
		Columns:   columns,
	}

	// Create the Delta Lake table
	table, err := deltalake.CreateDeltaLakeTable(options)
	require.NoError(t, err, "Failed to create Delta table")
	defer table.Close()

	t.Logf("Created test Delta table at: %s", tablePath)

	return tablePath
}

// setupAWSCredentials sets up fake AWS credentials for testing
func setupAWSCredentials(t *testing.T) func() {
	// Store original values
	originalAccessKey := os.Getenv("AWS_ACCESS_KEY_ID")
	originalSecretKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
	originalToken := os.Getenv("AWS_SESSION_TOKEN")

	// Set test credentials
	os.Setenv("AWS_ACCESS_KEY_ID", "test-access-key")
	os.Setenv("AWS_SECRET_ACCESS_KEY", "test-secret-key")
	os.Setenv("AWS_SESSION_TOKEN", "test-session-token")

	// Return cleanup function
	return func() {
		if originalAccessKey != "" {
			os.Setenv("AWS_ACCESS_KEY_ID", originalAccessKey)
		} else {
			os.Unsetenv("AWS_ACCESS_KEY_ID")
		}

		if originalSecretKey != "" {
			os.Setenv("AWS_SECRET_ACCESS_KEY", originalSecretKey)
		} else {
			os.Unsetenv("AWS_SECRET_ACCESS_KEY")
		}

		if originalToken != "" {
			os.Setenv("AWS_SESSION_TOKEN", originalToken)
		} else {
			os.Unsetenv("AWS_SESSION_TOKEN")
		}
	}
}

func setupEnvVars(t *testing.T, envVars map[string]string) func() {
	originals := make(map[string]*string, len(envVars))

	for key, value := range envVars {
		if current, exists := os.LookupEnv(key); exists {
			currentCopy := current
			originals[key] = &currentCopy
		} else {
			originals[key] = nil
		}

		require.NoError(t, os.Setenv(key, value))
	}

	return func() {
		for key, original := range originals {
			if original != nil {
				require.NoError(t, os.Setenv(key, *original))
			} else {
				require.NoError(t, os.Unsetenv(key))
			}
		}
	}
}

// TestDuckDBClient_NewDuckDBClient tests the constructor
func TestDuckDBClient_NewDuckDBClient(t *testing.T) {
	logger := zaptest.NewLogger(t)

	tests := []struct {
		name    string
		wantErr bool
	}{
		{
			name:    "should create DuckDBClient successfully",
			wantErr: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			client, err := NewDuckDBClient(logger)

			if tt.wantErr {
				assert.Error(t, err)
				assert.Nil(t, client)
			} else {
				assert.NoError(t, err)
				assert.NotNil(t, client)
				if client != nil {
					assert.NotNil(t, client.log)
					assert.NotNil(t, client.db)
					// Clean up
					client.Close()
				}
			}
		})
	}
}

// TestDuckDBClient_Close tests the Close method
func TestDuckDBClient_Close(t *testing.T) {
	logger := zaptest.NewLogger(t)

	tests := []struct {
		name    string
		setup   func() *DuckDBClient
		wantErr bool
	}{
		{
			name: "should close database connection successfully",
			setup: func() *DuckDBClient {
				client, _ := NewDuckDBClient(logger)
				return client
			},
			wantErr: false,
		},
		{
			name: "should handle nil database gracefully",
			setup: func() *DuckDBClient {
				return &DuckDBClient{
					log: logger,
					db:  nil,
				}
			},
			wantErr: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			client := tt.setup()
			err := client.Close()

			if tt.wantErr {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
			}
		})
	}
}

// TestDuckDBClient_QueryMaxTimestamp tests QueryMaxTimestamp method with real Delta tables
func TestDuckDBClient_QueryMaxTimestamp(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping Delta Lake integration test in short mode")
	}

	logger := zaptest.NewLogger(t)

	// Create a real DuckDB client
	client, err := NewDuckDBClient(logger)
	require.NoError(t, err)
	defer client.Close()

	t.Run("should return 0 for empty table", func(t *testing.T) {
		// Create an empty Delta table
		tablePath := setupTestDeltaTable(t)

		// Query max timestamp from empty table
		result, err := client.QueryMaxTimestamp(tablePath)
		assert.NoError(t, err)
		assert.Equal(t, int64(0), result)
	})

	t.Run("should handle non-existent table path gracefully", func(t *testing.T) {
		nonExistentPath := "/tmp/non_existent_delta_table"

		// Query should fail with appropriate error
		result, err := client.QueryMaxTimestamp(nonExistentPath)
		assert.Error(t, err)
		assert.Equal(t, int64(0), result)
		assert.Contains(t, err.Error(), "failed to query max timestamp")
	})
}

// TestDuckDBClient_QueryMaxTimestampWithPartitionFilter tests QueryMaxTimestampWithPartitionFilter method with real Delta tables
func TestDuckDBClient_QueryMaxTimestampWithPartitionFilter(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping Delta Lake integration test in short mode")
	}

	logger := zaptest.NewLogger(t)

	// Create a real DuckDB client
	client, err := NewDuckDBClient(logger)
	require.NoError(t, err)
	defer client.Close()

	t.Run("should return 0 for empty table with partition filter", func(t *testing.T) {
		// Create an empty Delta table
		tablePath := setupTestDeltaTable(t)

		// Query max timestamp from empty table with partition filter
		startTime := int64(1641081600) // 2022-01-02 00:00:00 UTC
		result, err := client.QueryMaxTimestampWithPartitionFilter(tablePath, startTime)
		assert.NoError(t, err)
		assert.Equal(t, int64(0), result)
	})

	t.Run("should handle non-existent table path gracefully with partition filter", func(t *testing.T) {
		nonExistentPath := "/tmp/non_existent_delta_table"
		startTime := int64(1641081600)

		// Query should fail with appropriate error
		result, err := client.QueryMaxTimestampWithPartitionFilter(nonExistentPath, startTime)
		assert.Error(t, err)
		assert.Equal(t, int64(0), result)
		assert.Contains(t, err.Error(), "failed to query max timestamp")
	})
}

// TestDuckDBClient_QueryData tests QueryData method
func TestDuckDBClient_QueryData(t *testing.T) {
	logger := zaptest.NewLogger(t)

	tests := []struct {
		name        string
		tablePath   string
		customSQL   string
		setupMock   func(sqlmock.Sqlmock)
		wantErr     bool
		errContains string
	}{
		{
			name:      "should execute custom query successfully",
			tablePath: "s3://test-bucket/table",
			customSQL: "SELECT * FROM {table_path} WHERE id > 100",
			setupMock: func(mock sqlmock.Sqlmock) {
				rows := sqlmock.NewRows([]string{"id", "name"}).
					AddRow(101, "test1").
					AddRow(102, "test2")
				mock.ExpectQuery("SELECT \\* FROM delta_scan").
					WillReturnRows(rows)
			},
			wantErr: false,
		},
		{
			name:      "should handle query execution error",
			tablePath: "s3://test-bucket/table",
			customSQL: "SELECT * FROM {table_path}",
			setupMock: func(mock sqlmock.Sqlmock) {
				mock.ExpectQuery("SELECT \\* FROM delta_scan").
					WillReturnError(errors.New("custom query failed"))
			},
			wantErr:     true,
			errContains: "failed to execute query",
		},
		{
			name:      "should handle credential error and return error when refresh fails",
			tablePath: "s3://test-bucket/table",
			customSQL: "SELECT COUNT(*) FROM {table_path}",
			setupMock: func(mock sqlmock.Sqlmock) {
				// First call fails with credential error
				mock.ExpectQuery("SELECT COUNT\\(\\*\\) FROM delta_scan").
					WillReturnError(errors.New("AuthenticationFailed"))
				// Since refresh fails in test environment, we expect error to be returned
			},
			wantErr:     true,
			errContains: "failed to execute query and refresh credentials",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create mock database
			db, mock, err := sqlmock.New()
			require.NoError(t, err)
			defer db.Close()

			// Setup mock expectations
			tt.setupMock(mock)

			// Create client with mock database
			client := &DuckDBClient{
				log:            logger,
				db:             db,
				vendor:         constants.VendorAWS,
				storageOptions: map[string]string{"region": "us-east-1"},
			}

			rows, err := client.QueryData(tt.tablePath, tt.customSQL)

			if tt.wantErr {
				assert.Error(t, err)
				if tt.errContains != "" {
					assert.Contains(t, err.Error(), tt.errContains)
				}
				assert.Nil(t, rows)
			} else {
				assert.NoError(t, err)
				assert.NotNil(t, rows)
				if rows != nil {
					rows.Close()
				}
			}

			// Verify all expectations were met
			assert.NoError(t, mock.ExpectationsWereMet())
		})
	}
}

// TestDuckDBClient_ConfigureForStorageOptions tests ConfigureForStorageOptions method
func TestDuckDBClient_ConfigureForStorageOptions(t *testing.T) {
	logger := zaptest.NewLogger(t)

	tests := []struct {
		name           string
		vendor         string
		storageOptions map[string]string
		setupMock      func(sqlmock.Sqlmock)
		wantErr        bool
		errContains    string
		envVars        map[string]string
	}{
		{
			name:           "should handle nil storage options",
			vendor:         constants.VendorAWS,
			storageOptions: nil,
			setupMock:      func(mock sqlmock.Sqlmock) {},
			wantErr:        false,
		},
		{
			name:   "should configure AWS access successfully",
			vendor: constants.VendorAWS,
			storageOptions: map[string]string{
				"region": "us-west-2",
			},
			setupMock: func(mock sqlmock.Sqlmock) {
				// AWS configuration will fail in test environment without real credentials
			},
			wantErr:     true,
			errContains: "failed to set s3_region",
		},
		{
			name:   "should configure Alicloud access successfully",
			vendor: constants.VendorAlicloud,
			storageOptions: map[string]string{
				"BucketEndpoint": "https://oss-cn-hangzhou.aliyuncs.com",
				"region":         "cn-hangzhou",
				"bucket":         "test-bucket",
			},
			envVars: map[string]string{
				constants.AlicloudAccessKeyIDEnv:     "test-ak",
				constants.AlicloudAccessKeySecretEnv: "test-sk",
			},
			setupMock: func(mock sqlmock.Sqlmock) {
				mock.ExpectExec("(?s)CREATE OR REPLACE SECRET alicloud_secret \\(.+ENDPOINT 'test-bucket\\.oss-cn-hangzhou\\.aliyuncs\\.com'").
					WillReturnResult(sqlmock.NewResult(0, 0))
			},
			wantErr: false,
		},
		{
			name:   "should fail Alicloud configuration when required options missing",
			vendor: constants.VendorAlicloud,
			storageOptions: map[string]string{
				"region": "cn-hangzhou",
			},
			envVars: map[string]string{
				constants.AlicloudAccessKeyIDEnv:     "test-ak",
				constants.AlicloudAccessKeySecretEnv: "test-sk",
			},
			setupMock:   func(mock sqlmock.Sqlmock) {},
			wantErr:     true,
			errContains: "BucketEndpoint is required",
		},
		{
			name:   "should fail Alicloud configuration when region missing",
			vendor: constants.VendorAlicloud,
			storageOptions: map[string]string{
				"BucketEndpoint": "https://oss-cn-hangzhou.aliyuncs.com",
				"bucket":         "test-bucket",
			},
			envVars: map[string]string{
				constants.AlicloudAccessKeyIDEnv:     "test-ak",
				constants.AlicloudAccessKeySecretEnv: "test-sk",
			},
			setupMock:   func(mock sqlmock.Sqlmock) {},
			wantErr:     true,
			errContains: "region or AWS_REGION is required",
		},
		{
			name:   "should fail Alicloud configuration when bucket missing",
			vendor: constants.VendorAlicloud,
			storageOptions: map[string]string{
				"BucketEndpoint": "https://oss-cn-hangzhou.aliyuncs.com",
				"region":         "cn-hangzhou",
			},
			envVars: map[string]string{
				constants.AlicloudAccessKeyIDEnv:     "test-ak",
				constants.AlicloudAccessKeySecretEnv: "test-sk",
			},
			setupMock:   func(mock sqlmock.Sqlmock) {},
			wantErr:     true,
			errContains: "bucket is required",
		},
		{
			name:   "should fail Alicloud configuration when endpoint invalid",
			vendor: constants.VendorAlicloud,
			storageOptions: map[string]string{
				"BucketEndpoint": "ht!tp://invalid-endpoint",
				"region":         "cn-hangzhou",
				"bucket":         "test-bucket",
			},
			envVars: map[string]string{
				constants.AlicloudAccessKeyIDEnv:     "test-ak",
				constants.AlicloudAccessKeySecretEnv: "test-sk",
			},
			setupMock:   func(mock sqlmock.Sqlmock) {},
			wantErr:     true,
			errContains: "invalid BucketEndpoint",
		},
		{
			name:   "should fail Alicloud configuration when env missing",
			vendor: constants.VendorAlicloud,
			storageOptions: map[string]string{
				"BucketEndpoint": "https://oss-cn-hangzhou.aliyuncs.com",
				"region":         "cn-hangzhou",
				"bucket":         "test-bucket",
			},
			setupMock:   func(mock sqlmock.Sqlmock) {},
			wantErr:     true,
			errContains: constants.AlicloudAccessKeyIDEnv,
		},
		{
			name:   "should configure GCP access successfully",
			vendor: constants.VendorGCP,
			storageOptions: map[string]string{
				"project": "test-project",
			},
			setupMock: func(mock sqlmock.Sqlmock) {
				mock.ExpectExec("CREATE OR REPLACE SECRET gcs_secret").WillReturnResult(sqlmock.NewResult(0, 0))
			},
			wantErr: false,
		},
		{
			name:   "should configure Azure access (requires real Azure credentials)",
			vendor: constants.VendorAzure,
			storageOptions: map[string]string{
				constants.AzureStorageAccountNameEnv: "test-account",
			},
			setupMock: func(mock sqlmock.Sqlmock) {
				// configureAzureAccess calls azidentity which can't be mocked with sqlmock.
			},
			wantErr: true, // Will fail without real Azure credentials
		},
		{
			name:   "should handle unsupported vendor gracefully",
			vendor: "unsupported",
			storageOptions: map[string]string{
				"key": "value",
			},
			setupMock: func(mock sqlmock.Sqlmock) {},
			wantErr:   false,
		},
		{
			name:   "should handle AWS configuration error",
			vendor: constants.VendorAWS,
			storageOptions: map[string]string{
				"region": "us-west-2",
			},
			setupMock: func(mock sqlmock.Sqlmock) {
				mock.ExpectExec("SET s3_region").WillReturnError(errors.New("failed to set region"))
			},
			wantErr:     true,
			errContains: "failed to set s3_region",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cleanup := func() {}
			if tt.envVars != nil {
				cleanup = setupEnvVars(t, tt.envVars)
			}
			defer cleanup()

			// Create mock database
			db, mock, err := sqlmock.New()
			require.NoError(t, err)
			defer db.Close()

			// Setup mock expectations
			tt.setupMock(mock)

			// Create client with mock database
			client := &DuckDBClient{
				log: logger,
				db:  db,
			}

			err = client.ConfigureForStorageOptions(tt.vendor, tt.storageOptions)

			if tt.wantErr {
				assert.Error(t, err)
				if tt.errContains != "" {
					assert.Contains(t, err.Error(), tt.errContains)
				}
			} else {
				assert.NoError(t, err)
				// Only check vendor if storage options are provided
				if tt.storageOptions != nil {
					assert.Equal(t, tt.vendor, client.vendor)
				}
				assert.Equal(t, tt.storageOptions, client.storageOptions)
			}

			// Verify all expectations were met
			assert.NoError(t, mock.ExpectationsWereMet())
		})
	}
}

// TestDuckDBClient_isCredentialError tests isCredentialError method
func TestDuckDBClient_isCredentialError(t *testing.T) {
	logger := zaptest.NewLogger(t)
	client := &DuckDBClient{log: logger}

	tests := []struct {
		name     string
		err      error
		expected bool
	}{
		{
			name:     "should return false for nil error",
			err:      nil,
			expected: false,
		},
		{
			name:     "should detect expired token error",
			err:      errors.New("The provided token has expired."),
			expected: true,
		},
		{
			name:     "should detect authentication failed error",
			err:      errors.New("AuthenticationFailed"),
			expected: true,
		},
		{
			name:     "should detect invalid access key error",
			err:      errors.New("InvalidAccessKeyId"),
			expected: true,
		},
		{
			name:     "should detect HTTP 403 error",
			err:      errors.New("HTTP 403 Forbidden"),
			expected: true,
		},
		{
			name:     "should detect access denied error",
			err:      errors.New("Access Denied"),
			expected: true,
		},
		{
			name:     "should return false for non-credential error",
			err:      errors.New("Connection timeout"),
			expected: false,
		},
		{
			name:     "should return false for general SQL error",
			err:      errors.New("syntax error"),
			expected: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := client.isCredentialError(tt.err)
			assert.Equal(t, tt.expected, result)
		})
	}
}

// TestDuckDBClient_refreshCredentialChain tests refreshCredentialChain method
func TestDuckDBClient_refreshCredentialChain(t *testing.T) {
	logger := zaptest.NewLogger(t)

	tests := []struct {
		name           string
		vendor         string
		storageOptions map[string]string
		setupMock      func(sqlmock.Sqlmock)
		wantErr        bool
		errContains    string
		envVars        map[string]string
	}{
		{
			name:           "should return error when no vendor configured",
			vendor:         "",
			storageOptions: nil,
			setupMock:      func(mock sqlmock.Sqlmock) {},
			wantErr:        true,
			errContains:    "no vendor or storage options configured",
		},
		{
			name:   "should refresh AWS credentials successfully",
			vendor: constants.VendorAWS,
			storageOptions: map[string]string{
				"region": "us-east-1",
			},
			setupMock: func(mock sqlmock.Sqlmock) {
				// AWS refresh will fail in test environment without real credentials
			},
			wantErr:     true,
			errContains: "failed to set s3_region",
		},
		{
			name:   "should refresh Alicloud credentials successfully",
			vendor: constants.VendorAlicloud,
			storageOptions: map[string]string{
				"BucketEndpoint": "https://oss-cn-hangzhou.aliyuncs.com",
				"region":         "cn-hangzhou",
				"bucket":         "test-bucket",
			},
			envVars: map[string]string{
				constants.AlicloudAccessKeyIDEnv:     "test-ak",
				constants.AlicloudAccessKeySecretEnv: "test-sk",
			},
			setupMock: func(mock sqlmock.Sqlmock) {
				mock.ExpectExec("(?s)CREATE OR REPLACE SECRET alicloud_secret \\(.+ENDPOINT 'test-bucket\\.oss-cn-hangzhou\\.aliyuncs\\.com'").
					WillReturnResult(sqlmock.NewResult(0, 0))
			},
			wantErr: false,
		},
		{
			name:   "should refresh GCP credentials successfully",
			vendor: constants.VendorGCP,
			storageOptions: map[string]string{
				"project": "test-project",
			},
			setupMock: func(mock sqlmock.Sqlmock) {
				mock.ExpectExec("CREATE OR REPLACE SECRET gcs_secret").WillReturnResult(sqlmock.NewResult(0, 0))
			},
			wantErr: false,
		},
		{
			name:   "should return error for unsupported vendor",
			vendor: "unsupported",
			storageOptions: map[string]string{
				"key": "value",
			},
			setupMock:   func(mock sqlmock.Sqlmock) {},
			wantErr:     true,
			errContains: "unsupported vendor for credential refresh",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cleanup := func() {}
			if tt.envVars != nil {
				cleanup = setupEnvVars(t, tt.envVars)
			}
			defer cleanup()

			// Create mock database
			db, mock, err := sqlmock.New()
			require.NoError(t, err)
			defer db.Close()

			// Setup mock expectations
			tt.setupMock(mock)

			// Create client with mock database
			client := &DuckDBClient{
				log:            logger,
				db:             db,
				vendor:         tt.vendor,
				storageOptions: tt.storageOptions,
			}

			err = client.refreshCredentialChain()

			if tt.wantErr {
				assert.Error(t, err)
				if tt.errContains != "" {
					assert.Contains(t, err.Error(), tt.errContains)
				}
			} else {
				assert.NoError(t, err)
			}

			// Verify all expectations were met
			assert.NoError(t, mock.ExpectationsWereMet())
		})
	}
}

// TestDuckDBClient_NewDuckDBClientWithConfig tests NewDuckDBClientWithConfig function
func TestDuckDBClient_NewDuckDBClientWithConfig(t *testing.T) {
	logger := zaptest.NewLogger(t)

	tests := []struct {
		name    string
		config  *config.ServiceConfig
		wantErr bool
	}{
		{
			name: "should create client with basic config",
			config: &config.ServiceConfig{
				Vendor: constants.VendorAWS,
			},
			wantErr: false,
		},
		{
			name: "should create client with NextGen config with AWS",
			config: &config.ServiceConfig{
				Vendor: constants.VendorAWS,
				NextGenMeteringConfig: &config.NextGenMeteringConfig{
					StorageOptions: map[string]string{
						"region": "us-east-1",
					},
				},
			},
			wantErr: true, // AWS credentials not available in CI
		},
		{
			name: "should create client with nil NextGen config",
			config: &config.ServiceConfig{
				Vendor:                constants.VendorGCP,
				NextGenMeteringConfig: nil,
			},
			wantErr: false,
		},
		{
			name: "should create client with Azure config and extract container from path",
			config: &config.ServiceConfig{
				Vendor: constants.VendorAzure,
				NextGenMeteringConfig: &config.NextGenMeteringConfig{
					StorageOptions: map[string]string{
						constants.AzureStorageAccountNameEnv: "testaccount",
					},
					DeltaTablePath: "az://testaccount.blob.core.windows.net/mycontainer/path/to/table",
				},
			},
			wantErr: true, // Azure credentials not available in CI
		},
		{
			name: "should create client with Azure config and explicit container name",
			config: &config.ServiceConfig{
				Vendor: constants.VendorAzure,
				NextGenMeteringConfig: &config.NextGenMeteringConfig{
					StorageOptions: map[string]string{
						constants.AzureStorageAccountNameEnv: "testaccount",
						constants.AzureContainerNameKey:      "explicit-container",
					},
					DeltaTablePath: "az://testaccount.blob.core.windows.net/other-container/path",
				},
			},
			wantErr: true, // Azure credentials not available in CI
		},
		{
			name: "should create client with GCP config",
			config: &config.ServiceConfig{
				Vendor: constants.VendorGCP,
				NextGenMeteringConfig: &config.NextGenMeteringConfig{
					StorageOptions: map[string]string{
						"bucket": "test-bucket",
					},
				},
			},
			wantErr: true, // GCP credential_chain requires real credentials
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			client, err := NewDuckDBClientWithConfig(tt.config, logger)

			if tt.wantErr {
				if err != nil {
					assert.Error(t, err)
					assert.Nil(t, client)
				} else {
					// Cloud credentials may be available locally
					assert.NotNil(t, client)
					if client != nil {
						client.Close()
					}
				}
			} else {
				assert.NoError(t, err)
				assert.NotNil(t, client)
				if client != nil {
					assert.Equal(t, logger, client.log)
					// Clean up
					client.Close()
				}
			}
		})
	}
}

// TestDuckDBClient_MemoryConfiguration tests memory configuration
func TestDuckDBClient_MemoryConfiguration(t *testing.T) {
	logger := zaptest.NewLogger(t)

	tests := []struct {
		name        string
		envVars     map[string]string
		setupClient func() (*DuckDBClient, sqlmock.Sqlmock, error)
		postSetup   func(*DuckDBClient) error
		wantErr     bool
	}{
		{
			name: "should configure memory with environment variables",
			envVars: map[string]string{
				"POD_MEMORY_LIMITS_MB":     "4096",
				"DUCKDB_MEMORY_PERCENTAGE": "70",
				"DUCKDB_SKIP_CA_BUNDLE":    "true",
			},
			setupClient: func() (*DuckDBClient, sqlmock.Sqlmock, error) {
				db, mock, err := sqlmock.New()
				if err != nil {
					return nil, nil, err
				}

				mock.ExpectExec("SET memory_limit").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("INSTALL httpfs; LOAD httpfs;").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("INSTALL azure; LOAD azure;").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("INSTALL cache_httpfs FROM community; LOAD cache_httpfs; SET cache_httpfs_cache_directory='.*'; SET cache_httpfs_type='on_disk';").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("SET enable_object_cache=true; SET enable_http_metadata_cache=true; SET enable_external_file_cache=true; SET parquet_metadata_cache=true;").WillReturnResult(sqlmock.NewResult(0, 0))

				client := &DuckDBClient{
					log: logger,
					db:  db,
				}
				return client, mock, nil
			},
			wantErr: false,
		},
		{
			name: "should use defaults for invalid environment variables",
			envVars: map[string]string{
				"POD_MEMORY_LIMITS_MB":     "invalid",
				"DUCKDB_MEMORY_PERCENTAGE": "150",
				"DUCKDB_SKIP_CA_BUNDLE":    "true",
			},
			setupClient: func() (*DuckDBClient, sqlmock.Sqlmock, error) {
				db, mock, err := sqlmock.New()
				if err != nil {
					return nil, nil, err
				}

				mock.ExpectExec("SET memory_limit").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("INSTALL httpfs; LOAD httpfs;").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("INSTALL azure; LOAD azure;").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("INSTALL cache_httpfs FROM community; LOAD cache_httpfs; SET cache_httpfs_cache_directory='.*'; SET cache_httpfs_type='on_disk';").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("SET enable_object_cache=true; SET enable_http_metadata_cache=true; SET enable_external_file_cache=true; SET parquet_metadata_cache=true;").WillReturnResult(sqlmock.NewResult(0, 0))

				client := &DuckDBClient{
					log: logger,
					db:  db,
				}
				return client, mock, nil
			},
			wantErr: false,
		},
		{
			name:    "should load delta extension from storage option path",
			envVars: map[string]string{"DUCKDB_SKIP_CA_BUNDLE": "true"},
			setupClient: func() (*DuckDBClient, sqlmock.Sqlmock, error) {
				db, mock, err := sqlmock.New()
				if err != nil {
					return nil, nil, err
				}

				mock.ExpectExec("SET memory_limit").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("INSTALL httpfs; LOAD httpfs;").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("INSTALL azure; LOAD azure;").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("INSTALL cache_httpfs FROM community; LOAD cache_httpfs; SET cache_httpfs_cache_directory='.*'; SET cache_httpfs_type='on_disk';").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("SET enable_object_cache=true; SET enable_http_metadata_cache=true; SET enable_external_file_cache=true; SET parquet_metadata_cache=true;").WillReturnResult(sqlmock.NewResult(0, 0))
				mock.ExpectExec("LOAD '/opt/extensions/delta\\.duckdb_extension';").WillReturnResult(sqlmock.NewResult(0, 0))

				client := &DuckDBClient{
					log: logger,
					db:  db,
					storageOptions: map[string]string{
						constants.StorageOptionPatchDeltaSo: "/opt/extensions/delta.duckdb_extension",
					},
				}
				return client, mock, nil
			},
			postSetup: func(c *DuckDBClient) error {
				return c.installAdditionalExtensions(c.storageOptions)
			},
			wantErr: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Set environment variables
			for key, value := range tt.envVars {
				os.Setenv(key, value)
			}
			defer func() {
				for key := range tt.envVars {
					os.Unsetenv(key)
				}
			}()

			client, mock, err := tt.setupClient()
			require.NoError(t, err)
			defer client.db.Close()

			err = client.setupDuckDB()

			if tt.wantErr {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
				if tt.postSetup != nil {
					err = tt.postSetup(client)
					if tt.wantErr {
						assert.Error(t, err)
					} else {
						assert.NoError(t, err)
					}
				}
			}

			// Verify all expectations were met
			assert.NoError(t, mock.ExpectationsWereMet())
		})
	}
} // TestDuckDBClient_ConfigureForStorageOptionsReal tests AWS configuration with real DuckDB and environment variables
func TestDuckDBClient_ConfigureForStorageOptionsReal(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping AWS configuration integration test in short mode")
	}

	logger := zaptest.NewLogger(t)

	t.Run("should configure AWS access with environment credentials", func(t *testing.T) {
		// Setup AWS environment credentials
		cleanup := setupAWSCredentials(t)
		defer cleanup()

		// Create a real DuckDB client
		client, err := NewDuckDBClient(logger)
		require.NoError(t, err)
		defer client.Close()

		// Configure AWS storage options
		storageOptions := map[string]string{
			"region": "us-east-1",
		}

		err = client.ConfigureForStorageOptions(constants.VendorAWS, storageOptions)
		assert.NoError(t, err, "Should configure AWS access successfully with environment credentials")

		// Verify client state
		assert.Equal(t, constants.VendorAWS, client.vendor)
		assert.Equal(t, storageOptions, client.storageOptions)
	})

	t.Run("should handle nil storage options gracefully", func(t *testing.T) {
		// Create a real DuckDB client
		client, err := NewDuckDBClient(logger)
		require.NoError(t, err)
		defer client.Close()

		// Configure with nil storage options
		err = client.ConfigureForStorageOptions(constants.VendorAWS, nil)
		assert.NoError(t, err, "Should handle nil storage options gracefully")
	})
}

// TestExtractContainerFromAzurePath tests the extractContainerFromAzurePath helper
func TestExtractContainerFromAzurePath(t *testing.T) {
	tests := []struct {
		name     string
		path     string
		expected string
	}{
		{
			name:     "should extract container from az:// path",
			path:     "az://myaccount.blob.core.windows.net/mycontainer/path/to/data",
			expected: "mycontainer",
		},
		{
			name:     "should extract container from https:// path",
			path:     "https://myaccount.blob.core.windows.net/mycontainer/path/to/data",
			expected: "mycontainer",
		},
		{
			name:     "should extract container with no sub-path",
			path:     "az://myaccount.blob.core.windows.net/mycontainer",
			expected: "mycontainer",
		},
		{
			name:     "should return empty for empty path",
			path:     "",
			expected: "",
		},
		{
			name:     "should return empty for root path",
			path:     "az://myaccount.blob.core.windows.net/",
			expected: "",
		},
		{
			name:     "should return empty for host-only path",
			path:     "az://myaccount.blob.core.windows.net",
			expected: "",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := extractContainerFromAzurePath(tt.path)
			assert.Equal(t, tt.expected, result)
		})
	}
}

// TestDuckDBClient_ConfigureAzureAccess tests Azure configuration using access_token.
// Since configureAzureAccess calls azidentity.NewDefaultAzureCredential which requires
// real Azure credentials, we test error paths and configuration validation here.
func TestDuckDBClient_ConfigureAzureAccess(t *testing.T) {
	logger := zaptest.NewLogger(t)

	t.Run("should error when account name is missing", func(t *testing.T) {
		db, _, err := sqlmock.New()
		require.NoError(t, err)
		defer db.Close()

		client := &DuckDBClient{
			log: logger,
			db:  db,
		}

		err = client.configureAzureAccess(map[string]string{})
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "AZURE_STORAGE_ACCOUNT_NAME")
	})

	t.Run("should error when account name is only whitespace", func(t *testing.T) {
		db, _, err := sqlmock.New()
		require.NoError(t, err)
		defer db.Close()

		client := &DuckDBClient{
			log: logger,
			db:  db,
		}

		err = client.configureAzureAccess(map[string]string{
			constants.AzureStorageAccountNameEnv: "   ",
		})
		assert.Error(t, err)
		assert.Contains(t, err.Error(), "AZURE_STORAGE_ACCOUNT_NAME")
	})

	t.Run("should error when Azure credentials are not available", func(t *testing.T) {
		// Clear all Azure-related env vars to ensure no credentials are available
		cleanup := setupEnvVars(t, map[string]string{
			"AZURE_TENANT_ID":     "",
			"AZURE_CLIENT_ID":     "",
			"AZURE_CLIENT_SECRET": "",
			"AZURE_USERNAME":      "",
			"AZURE_PASSWORD":      "",
		})
		defer cleanup()
		os.Unsetenv("AZURE_TENANT_ID")
		os.Unsetenv("AZURE_CLIENT_ID")
		os.Unsetenv("AZURE_CLIENT_SECRET")
		os.Unsetenv("AZURE_USERNAME")
		os.Unsetenv("AZURE_PASSWORD")

		db, _, err := sqlmock.New()
		require.NoError(t, err)
		defer db.Close()

		client := &DuckDBClient{
			log: logger,
			db:  db,
		}

		err = client.configureAzureAccess(map[string]string{
			constants.AzureStorageAccountNameEnv: "test-account",
		})
		assert.Error(t, err)
	})
}

func TestDuckDBClient_ConfigureCABundle(t *testing.T) {
	logger := zaptest.NewLogger(t)

	t.Run("should skip when DUCKDB_SKIP_CA_BUNDLE is set", func(t *testing.T) {
		cleanup := setupEnvVars(t, map[string]string{
			"DUCKDB_SKIP_CA_BUNDLE": "true",
		})
		defer cleanup()

		db, mock, err := sqlmock.New()
		require.NoError(t, err)
		defer db.Close()

		// No SQL expectations needed since CA bundle is skipped

		client := &DuckDBClient{
			log: logger,
			db:  db,
		}

		err = client.configureCABundle()
		assert.NoError(t, err)
		assert.NoError(t, mock.ExpectationsWereMet())
	})

	t.Run("should use DUCKDB_CA_CERT_FILE env override", func(t *testing.T) {
		// Create a temp file to use as CA bundle
		tmpFile, err := os.CreateTemp("", "ca-cert-*.crt")
		require.NoError(t, err)
		tmpFile.WriteString("dummy CA cert")
		tmpFile.Close()
		defer os.Remove(tmpFile.Name())

		cleanup := setupEnvVars(t, map[string]string{
			"DUCKDB_CA_CERT_FILE": tmpFile.Name(),
		})
		defer cleanup()

		db, mock, err := sqlmock.New()
		require.NoError(t, err)
		defer db.Close()

		mock.ExpectExec("SET ca_cert_file=").WillReturnResult(sqlmock.NewResult(0, 0))
		mock.ExpectExec("SET azure_transport_option_ca_info=").WillReturnResult(sqlmock.NewResult(0, 0))
		mock.ExpectExec("SET azure_transport_option_type=").WillReturnResult(sqlmock.NewResult(0, 0))

		client := &DuckDBClient{
			log: logger,
			db:  db,
		}

		err = client.configureCABundle()
		assert.NoError(t, err)
		assert.NoError(t, mock.ExpectationsWereMet())
	})
}

func TestDuckDBClient_ExportAndSetAWSCredentials(t *testing.T) {
	logger := zaptest.NewLogger(t)

	t.Run("should error when aws cli is not available", func(t *testing.T) {
		origPath := os.Getenv("PATH")
		os.Setenv("PATH", "")
		defer os.Setenv("PATH", origPath)

		client := &DuckDBClient{log: logger}
		err := client.exportAndSetAWSCredentials("test-profile")
		assert.Error(t, err)
	})

	t.Run("should successfully export and set credentials", func(t *testing.T) {
		// Create a fake aws CLI script
		tmpDir := t.TempDir()
		fakeAWS := tmpDir + "/aws"
		err := os.WriteFile(fakeAWS, []byte("#!/bin/sh\necho '{\"AccessKeyId\":\"AKIAIOSFODNN7EXAMPLE\",\"SecretAccessKey\":\"wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY\",\"SessionToken\":\"test-session-token\",\"Expiration\":\"2024-01-01T00:00:00Z\"}'"), 0755)
		require.NoError(t, err)

		origPath := os.Getenv("PATH")
		origAccessKey := os.Getenv("AWS_ACCESS_KEY_ID")
		origSecretKey := os.Getenv("AWS_SECRET_ACCESS_KEY")
		origSessionToken := os.Getenv("AWS_SESSION_TOKEN")
		os.Setenv("PATH", tmpDir)
		defer func() {
			os.Setenv("PATH", origPath)
			if origAccessKey != "" {
				os.Setenv("AWS_ACCESS_KEY_ID", origAccessKey)
			} else {
				os.Unsetenv("AWS_ACCESS_KEY_ID")
			}
			if origSecretKey != "" {
				os.Setenv("AWS_SECRET_ACCESS_KEY", origSecretKey)
			} else {
				os.Unsetenv("AWS_SECRET_ACCESS_KEY")
			}
			if origSessionToken != "" {
				os.Setenv("AWS_SESSION_TOKEN", origSessionToken)
			} else {
				os.Unsetenv("AWS_SESSION_TOKEN")
			}
		}()

		client := &DuckDBClient{log: logger}
		err = client.exportAndSetAWSCredentials("test-profile")
		assert.NoError(t, err)
		assert.Equal(t, "AKIAIOSFODNN7EXAMPLE", os.Getenv("AWS_ACCESS_KEY_ID"))
		assert.Equal(t, "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", os.Getenv("AWS_SECRET_ACCESS_KEY"))
		assert.Equal(t, "test-session-token", os.Getenv("AWS_SESSION_TOKEN"))
	})
}

func TestDuckDBClient_GetAWSCredentialsFromSDK(t *testing.T) {
	logger := zaptest.NewLogger(t)

	t.Run("should retrieve credentials from SDK", func(t *testing.T) {
		client := &DuckDBClient{log: logger}
		creds, err := client.getAWSCredentialsFromSDK("us-east-1", "")
		if err != nil {
			assert.Nil(t, creds)
		} else {
			assert.NotNil(t, creds)
			assert.NotEmpty(t, creds.AccessKeyID)
		}
	})

	t.Run("should error with invalid role ARN", func(t *testing.T) {
		client := &DuckDBClient{log: logger}
		creds, err := client.getAWSCredentialsFromSDK("us-east-1", "arn:aws:iam::000000000000:role/nonexistent-role")
		assert.Error(t, err)
		assert.Nil(t, creds)
	})
}

func TestDuckDBClient_SetupDuckDB(t *testing.T) {
	logger := zaptest.NewLogger(t)

	t.Run("should setup DuckDB successfully", func(t *testing.T) {
		realClient, err := NewDuckDBClient(logger)
		require.NoError(t, err)
		defer realClient.Close()

		err = realClient.setupDuckDB()
		assert.NoError(t, err)
	})
}