package checkpoint

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"
	"go.uber.org/zap/zaptest"

	"github.com/tidbcloud/metering/pkg/common/config"
	"github.com/tidbcloud/metering/pkg/common/constants"
)

// ==== Mock Implementations ====

// MockMeteringService is a mock for MeteringServiceInterface
type MockMeteringService struct {
	GetTableVersionFunc                      func(ctx context.Context) (int64, error)
	QueryMaxTimestampWithPartitionFilterFunc func(ctx context.Context, partitionTime int64) (int64, error)
}

func (m *MockMeteringService) GetTableVersion(ctx context.Context) (int64, error) {
	if m.GetTableVersionFunc != nil {
		return m.GetTableVersionFunc(ctx)
	}
	return 0, errors.New("mock not configured")
}

func (m *MockMeteringService) QueryMaxTimestampWithPartitionFilter(ctx context.Context, partitionTime int64) (int64, error) {
	if m.QueryMaxTimestampWithPartitionFilterFunc != nil {
		return m.QueryMaxTimestampWithPartitionFilterFunc(ctx, partitionTime)
	}
	return 0, errors.New("mock not configured")
}

// MockDuckDBClient is a mock for DuckDBClientInterface
type MockDuckDBClient struct {
	QueryMaxTimestampWithPartitionFilterFunc func(tablePath string, partitionTime int64) (int64, error)
	CloseFunc                                func() error
}

func (m *MockDuckDBClient) QueryMaxTimestampWithPartitionFilter(tablePath string, partitionTime int64) (int64, error) {
	if m.QueryMaxTimestampWithPartitionFilterFunc != nil {
		return m.QueryMaxTimestampWithPartitionFilterFunc(tablePath, partitionTime)
	}
	return 0, errors.New("mock not configured")
}

func (m *MockDuckDBClient) Close() error {
	if m.CloseFunc != nil {
		return m.CloseFunc()
	}
	return nil
}

// MockDuckDBClientFactory is a mock factory for creating DuckDB clients
type MockDuckDBClientFactory struct {
	NewDuckDBClientWithConfigFunc func(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error)
}

func (m *MockDuckDBClientFactory) NewDuckDBClientWithConfig(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error) {
	if m.NewDuckDBClientWithConfigFunc != nil {
		return m.NewDuckDBClientWithConfigFunc(conf, log)
	}
	return nil, errors.New("mock factory not configured")
}

// ==== Unit Tests ====

func TestNewDeltaLakeCheckpointHelper(t *testing.T) {
	logger := zaptest.NewLogger(t)
	conf := &config.ServiceConfig{
		Vendor: "aws",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			DeltaTablePath: "s3://test-bucket/test-table",
		},
	}

	helper := NewDeltaLakeCheckpointHelper(conf, logger)

	assert.NotNil(t, helper)
	assert.Equal(t, conf, helper.conf)
	assert.Equal(t, logger, helper.log)
	assert.NotNil(t, helper.duckDBFactory)
}

// Test edge cases
func TestDeltaLakeCheckpointHelper_EdgeCases(t *testing.T) {
	logger := zaptest.NewLogger(t)

	t.Run("nil config", func(t *testing.T) {
		helper := NewDeltaLakeCheckpointHelper(nil, logger)
		assert.NotNil(t, helper)
		assert.Nil(t, helper.conf)
		assert.Equal(t, logger, helper.log)
	})

	t.Run("empty checkpoint data", func(t *testing.T) {
		conf := &config.ServiceConfig{}
		helper := NewDeltaLakeCheckpointHelper(conf, logger)

		emptyCheckpoint := &CheckpointData{}
		shouldRecover, version, err := helper.ShouldRecoverFromTable(context.Background(), nil, emptyCheckpoint)

		require.NoError(t, err)
		assert.False(t, shouldRecover)
		assert.Equal(t, int64(0), version)
	})

	t.Run("mock client close error", func(t *testing.T) {
		conf := &config.ServiceConfig{}
		closeError := errors.New("close error")

		mockFactory := &MockDuckDBClientFactory{
			NewDuckDBClientWithConfigFunc: func(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error) {
				return &MockDuckDBClient{
					QueryMaxTimestampWithPartitionFilterFunc: func(tablePath string, partitionTime int64) (int64, error) {
						return 100, nil
					},
					CloseFunc: func() error {
						return closeError
					},
				}, nil
			},
		}

		helper := NewDeltaLakeCheckpointHelperWithFactory(conf, logger, mockFactory)

		// Even if close fails, the main operation should succeed
		timestamp, err := helper.GetLatestCostTimestamp(context.Background(), nil, "test-path", 12345)

		require.NoError(t, err)
		assert.Equal(t, int64(100), timestamp)
	})
}

// Test direct implementation with refactored interfaces
func TestDeltaLakeCheckpointHelper_DirectImplementation(t *testing.T) {
	logger := zaptest.NewLogger(t)

	t.Run("GetTableVersion with mock service", func(t *testing.T) {
		conf := &config.ServiceConfig{
			Vendor: "aws",
		}

		helper := NewDeltaLakeCheckpointHelper(conf, logger)

		// Create mock service
		mockService := &MockMeteringService{
			GetTableVersionFunc: func(ctx context.Context) (int64, error) {
				return 42, nil
			},
		}

		version, err := helper.GetTableVersion(context.Background(), mockService)

		require.NoError(t, err)
		assert.Equal(t, int64(42), version)
	})

	t.Run("GetTableVersion with service error", func(t *testing.T) {
		conf := &config.ServiceConfig{}
		helper := NewDeltaLakeCheckpointHelper(conf, logger)

		// Create mock service that returns error
		mockService := &MockMeteringService{
			GetTableVersionFunc: func(ctx context.Context) (int64, error) {
				return 0, errors.New("service error")
			},
		}

		version, err := helper.GetTableVersion(context.Background(), mockService)

		assert.Error(t, err)
		assert.Equal(t, int64(0), version)
		assert.Contains(t, err.Error(), "failed to get table version")
	})

	t.Run("ShouldRecoverFromTable with empty path", func(t *testing.T) {
		conf := &config.ServiceConfig{}
		helper := NewDeltaLakeCheckpointHelper(conf, logger)

		checkpointData := &CheckpointData{
			TaskName:          "test-task",
			Checkpoint:        12345,
			DeltaTablePath:    "", // Empty path
			DeltaTableVersion: 10,
		}

		shouldRecover, version, err := helper.ShouldRecoverFromTable(context.Background(), nil, checkpointData)

		require.NoError(t, err)
		assert.False(t, shouldRecover)
		assert.Equal(t, int64(0), version)
	})

	t.Run("ShouldRecoverFromTable with service error", func(t *testing.T) {
		conf := &config.ServiceConfig{}
		helper := NewDeltaLakeCheckpointHelper(conf, logger)

		mockService := &MockMeteringService{
			GetTableVersionFunc: func(ctx context.Context) (int64, error) {
				return 0, errors.New("version error")
			},
		}

		checkpointData := &CheckpointData{
			TaskName:          "test-task",
			Checkpoint:        12345,
			DeltaTablePath:    "s3://test-bucket/test-table",
			DeltaTableVersion: 10,
		}

		shouldRecover, version, err := helper.ShouldRecoverFromTable(context.Background(), mockService, checkpointData)

		assert.Error(t, err)
		assert.False(t, shouldRecover)
		assert.Equal(t, int64(0), version)
		assert.Contains(t, err.Error(), "failed to get current table version")
	})

	t.Run("ShouldRecoverFromTable with recovery scenarios", func(t *testing.T) {
		conf := &config.ServiceConfig{}
		helper := NewDeltaLakeCheckpointHelper(conf, logger)

		testCases := []struct {
			name            string
			currentVersion  int64
			checkpointVer   int64
			expectedRecover bool
		}{
			{"current version higher - recovery needed", 15, 10, true},
			{"current version same - no recovery needed", 10, 10, false},
			{"current version lower - no recovery needed", 8, 10, false},
		}

		for _, tc := range testCases {
			t.Run(tc.name, func(t *testing.T) {
				mockService := &MockMeteringService{
					GetTableVersionFunc: func(ctx context.Context) (int64, error) {
						return tc.currentVersion, nil
					},
				}

				checkpointData := &CheckpointData{
					TaskName:          "test-task",
					Checkpoint:        1640995200,
					DeltaTablePath:    "s3://test-bucket/test-table",
					DeltaTableVersion: tc.checkpointVer,
				}

				shouldRecover, version, err := helper.ShouldRecoverFromTable(context.Background(), mockService, checkpointData)

				require.NoError(t, err)
				assert.Equal(t, tc.expectedRecover, shouldRecover)
				assert.Equal(t, tc.currentVersion, version)
			})
		}
	})

	t.Run("GetLatestCostTimestamp uses metering service for Alicloud", func(t *testing.T) {
		conf := &config.ServiceConfig{
			Vendor: constants.VendorAlicloud,
			NextGenMeteringConfig: &config.NextGenMeteringConfig{
				DeltaTablePath: "oss://test-bucket/test-table",
			},
		}

		factoryCalled := false

		mockFactory := &MockDuckDBClientFactory{
			NewDuckDBClientWithConfigFunc: func(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error) {
				factoryCalled = true
				return nil, errors.New("duckdb should not be used for Alicloud")
			},
		}

		helper := NewDeltaLakeCheckpointHelperWithFactory(conf, logger, mockFactory)

		mockService := &MockMeteringService{
			QueryMaxTimestampWithPartitionFilterFunc: func(ctx context.Context, partitionTime int64) (int64, error) {
				return 1641081600, nil
			},
		}

		timestamp, err := helper.GetLatestCostTimestamp(context.Background(), mockService, "oss://test-bucket/test-table", 12345)

		require.NoError(t, err)
		assert.Equal(t, int64(1641081600), timestamp)
		assert.False(t, factoryCalled)
	})

	t.Run("GetLatestCostTimestamp with Alicloud vendor and nil service", func(t *testing.T) {
		conf := &config.ServiceConfig{
			Vendor: constants.VendorAlicloud,
			NextGenMeteringConfig: &config.NextGenMeteringConfig{
				DeltaTablePath: "oss://test-bucket/test-table",
			},
		}

		helper := NewDeltaLakeCheckpointHelper(conf, logger)

		timestamp, err := helper.GetLatestCostTimestamp(context.Background(), nil, "oss://test-bucket/test-table", 12345)

		assert.Error(t, err)
		assert.Equal(t, int64(0), timestamp)
		assert.Contains(t, err.Error(), "metering service is required")
	})

	t.Run("GetLatestCostTimestamp with mock factory success", func(t *testing.T) {
		conf := &config.ServiceConfig{}

		mockFactory := &MockDuckDBClientFactory{
			NewDuckDBClientWithConfigFunc: func(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error) {
				return &MockDuckDBClient{
					QueryMaxTimestampWithPartitionFilterFunc: func(tablePath string, partitionTime int64) (int64, error) {
						return 1641081600, nil
					},
					CloseFunc: func() error {
						return nil
					},
				}, nil
			},
		}

		helper := NewDeltaLakeCheckpointHelperWithFactory(conf, logger, mockFactory)

		timestamp, err := helper.GetLatestCostTimestamp(context.Background(), nil, "s3://test-bucket/test-table", 12345)

		require.NoError(t, err)
		assert.Equal(t, int64(1641081600), timestamp)
	})

	t.Run("GetLatestCostTimestamp with factory error", func(t *testing.T) {
		conf := &config.ServiceConfig{}

		mockFactory := &MockDuckDBClientFactory{
			NewDuckDBClientWithConfigFunc: func(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error) {
				return nil, errors.New("factory error")
			},
		}

		helper := NewDeltaLakeCheckpointHelperWithFactory(conf, logger, mockFactory)

		timestamp, err := helper.GetLatestCostTimestamp(context.Background(), nil, "s3://test-bucket/test-table", 12345)

		assert.Error(t, err)
		assert.Equal(t, int64(0), timestamp)
		assert.Contains(t, err.Error(), "failed to create DuckDB client")
	})

	t.Run("GetLatestCostTimestamp with query error", func(t *testing.T) {
		conf := &config.ServiceConfig{}

		mockFactory := &MockDuckDBClientFactory{
			NewDuckDBClientWithConfigFunc: func(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error) {
				return &MockDuckDBClient{
					QueryMaxTimestampWithPartitionFilterFunc: func(tablePath string, partitionTime int64) (int64, error) {
						return 0, errors.New("query error")
					},
					CloseFunc: func() error {
						return nil
					},
				}, nil
			},
		}

		helper := NewDeltaLakeCheckpointHelperWithFactory(conf, logger, mockFactory)

		timestamp, err := helper.GetLatestCostTimestamp(context.Background(), nil, "s3://test-bucket/test-table", 12345)

		assert.Error(t, err)
		assert.Equal(t, int64(0), timestamp)
		assert.Contains(t, err.Error(), "failed to query max timestamp")
	})

	t.Run("RecoverCheckpointFromTable success", func(t *testing.T) {
		conf := &config.ServiceConfig{}

		mockFactory := &MockDuckDBClientFactory{
			NewDuckDBClientWithConfigFunc: func(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error) {
				return &MockDuckDBClient{
					QueryMaxTimestampWithPartitionFilterFunc: func(tablePath string, partitionTime int64) (int64, error) {
						return 1641081600, nil
					},
					CloseFunc: func() error {
						return nil
					},
				}, nil
			},
		}

		helper := NewDeltaLakeCheckpointHelperWithFactory(conf, logger, mockFactory)

		checkpointData := &CheckpointData{
			TaskName:          "test-task",
			Checkpoint:        1640995200,
			DeltaTablePath:    "s3://test-bucket/test-table",
			DeltaTableVersion: 10,
		}

		recoveredData, err := helper.RecoverCheckpointFromTable(context.Background(), nil, 15, checkpointData)

		require.NoError(t, err)
		assert.NotNil(t, recoveredData)
		assert.Equal(t, "test-task", recoveredData.TaskName)
		assert.Equal(t, int64(1641081600), recoveredData.Checkpoint)
		assert.Equal(t, int64(15), recoveredData.DeltaTableVersion)
		assert.Equal(t, "s3://test-bucket/test-table", recoveredData.DeltaTablePath)
		assert.True(t, recoveredData.UpdatedAt > 0)
		assert.True(t, recoveredData.UpdatedAt <= time.Now().Unix())
	})

	t.Run("RecoverCheckpointFromTable uses metering service for Alicloud", func(t *testing.T) {
		conf := &config.ServiceConfig{
			Vendor: constants.VendorAlicloud,
			NextGenMeteringConfig: &config.NextGenMeteringConfig{
				DeltaTablePath: "oss://test-bucket/test-table",
			},
		}

		factoryCalled := false

		mockFactory := &MockDuckDBClientFactory{
			NewDuckDBClientWithConfigFunc: func(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error) {
				factoryCalled = true
				return nil, errors.New("duckdb should not be used for Alicloud")
			},
		}

		helper := NewDeltaLakeCheckpointHelperWithFactory(conf, logger, mockFactory)

		mockService := &MockMeteringService{
			QueryMaxTimestampWithPartitionFilterFunc: func(ctx context.Context, partitionTime int64) (int64, error) {
				return 1641081600, nil
			},
		}

		checkpointData := &CheckpointData{
			TaskName:          "test-task",
			Checkpoint:        1640995200,
			DeltaTablePath:    "oss://test-bucket/test-table",
			DeltaTableVersion: 10,
		}

		recoveredData, err := helper.RecoverCheckpointFromTable(context.Background(), mockService, 15, checkpointData)

		require.NoError(t, err)
		assert.NotNil(t, recoveredData)
		assert.Equal(t, int64(1641081600), recoveredData.Checkpoint)
		assert.Equal(t, int64(15), recoveredData.DeltaTableVersion)
		assert.False(t, factoryCalled)
	})

	t.Run("RecoverCheckpointFromTable with timestamp error", func(t *testing.T) {
		conf := &config.ServiceConfig{}

		mockFactory := &MockDuckDBClientFactory{
			NewDuckDBClientWithConfigFunc: func(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error) {
				return &MockDuckDBClient{
					QueryMaxTimestampWithPartitionFilterFunc: func(tablePath string, partitionTime int64) (int64, error) {
						return 0, errors.New("timestamp query failed")
					},
					CloseFunc: func() error {
						return nil
					},
				}, nil
			},
		}

		helper := NewDeltaLakeCheckpointHelperWithFactory(conf, logger, mockFactory)

		checkpointData := &CheckpointData{
			TaskName:          "test-task",
			Checkpoint:        1640995200,
			DeltaTablePath:    "s3://test-bucket/test-table",
			DeltaTableVersion: 10,
		}

		recoveredData, err := helper.RecoverCheckpointFromTable(context.Background(), nil, 15, checkpointData)

		assert.Error(t, err)
		assert.Nil(t, recoveredData)
		assert.Contains(t, err.Error(), "failed to get latest timestamp from table")
	})

	t.Run("Constructor with factory", func(t *testing.T) {
		conf := &config.ServiceConfig{
			Vendor: "aws",
		}

		mockFactory := &MockDuckDBClientFactory{}
		helper := NewDeltaLakeCheckpointHelperWithFactory(conf, logger, mockFactory)

		assert.NotNil(t, helper)
		assert.Equal(t, conf, helper.conf)
		assert.Equal(t, logger, helper.log)
		assert.Equal(t, mockFactory, helper.duckDBFactory)
	})
}

// Test additional edge cases to improve coverage
func TestDeltaLakeCheckpointHelper_AdditionalCoverage(t *testing.T) {
	logger := zaptest.NewLogger(t)

	t.Run("ShouldRecoverFromTable_with_GetTableVersion_error", func(t *testing.T) {
		mockMeteringService := &MockMeteringService{
			GetTableVersionFunc: func(ctx context.Context) (int64, error) {
				return 0, errors.New("table version query error")
			},
		}

		helper := &DeltaLakeCheckpointHelper{
			log:           logger,
			conf:          &config.ServiceConfig{},
			duckDBFactory: &MockDuckDBClientFactory{},
		}

		checkpointData := &CheckpointData{
			DeltaTablePath:    "s3://test-bucket/test-table",
			DeltaTableVersion: 10,
		}

		shouldRecover, version, err := helper.ShouldRecoverFromTable(context.Background(), mockMeteringService, checkpointData)

		assert.Error(t, err)
		assert.False(t, shouldRecover)
		assert.Equal(t, int64(0), version)
		assert.Contains(t, err.Error(), "failed to get current table version")
	})

	t.Run("RecoverCheckpointFromTable_with_GetLatestCostTimestamp_error", func(t *testing.T) {
		mockDuckDBClient := &MockDuckDBClient{
			QueryMaxTimestampWithPartitionFilterFunc: func(tablePath string, partitionTime int64) (int64, error) {
				return 0, errors.New("timestamp query error")
			},
			CloseFunc: func() error {
				return nil
			},
		}

		mockFactory := &MockDuckDBClientFactory{
			NewDuckDBClientWithConfigFunc: func(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error) {
				return mockDuckDBClient, nil
			},
		}

		helper := &DeltaLakeCheckpointHelper{
			log:           logger,
			conf:          &config.ServiceConfig{},
			duckDBFactory: mockFactory,
		}

		checkpointData := &CheckpointData{
			TaskName:          "test-task",
			DeltaTablePath:    "s3://test-bucket/test-table",
			DeltaTableVersion: 10,
			Checkpoint:        1640995200,
		}

		recovered, err := helper.RecoverCheckpointFromTable(context.Background(), nil, 15, checkpointData)

		assert.Error(t, err)
		assert.Nil(t, recovered)
		assert.Contains(t, err.Error(), "failed to get latest timestamp from table")
	})
}