package nextgenrepoter

import (
	"context"
	"database/sql"
	"errors"
	"testing"
	"time"

	sqlmock "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/checkpoint"
	"github.com/tidbcloud/metering/pkg/common/config"
	"github.com/tidbcloud/metering/pkg/common/models"
	"go.uber.org/zap"
)

// mockTableManager implements tableManagerInterface for testing
type mockTableManager struct {
	queryResult []map[string]interface{}
	queryErr    error
	openErr     error
	lastSQL     string
	lastOptions *deltalake.QueryOptions
	openCalled  bool
}

func (m *mockTableManager) OpenTable(ctx context.Context) error {
	m.openCalled = true
	return m.openErr
}

func (m *mockTableManager) Query(ctx context.Context, sql string, options *deltalake.QueryOptions) ([]map[string]interface{}, error) {
	m.lastSQL = sql
	m.lastOptions = options
	return m.queryResult, m.queryErr
}

func TestNewNextGenRepoterLoader(t *testing.T) {
	logger := zap.NewNop()

	t.Run("success", func(t *testing.T) {
		originalFactory := newCheckpointFactory
		t.Cleanup(func() { newCheckpointFactory = originalFactory })

		mockManager := &mockCheckpointManager{}
		newCheckpointFactory = func(conf *config.ServiceConfig, log *zap.Logger) checkpointFactory {
			return &mockCheckpointFactory{manager: mockManager}
		}

		conf := &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{}}

		loader, err := NewNextGenRepoterLoader(conf, logger)
		require.NoError(t, err)

		impl, ok := loader.(*NextGenRepoterLoaderImpl)
		require.True(t, ok)
		assert.Equal(t, mockManager, impl.checkpointMgr)
		assert.NotNil(t, impl.duckFactory)
	})

	t.Run("missing_config", func(t *testing.T) {
		_, err := NewNextGenRepoterLoader(&config.ServiceConfig{}, logger)
		require.Error(t, err)
	})

	t.Run("checkpoint_error", func(t *testing.T) {
		originalFactory := newCheckpointFactory
		t.Cleanup(func() { newCheckpointFactory = originalFactory })

		newCheckpointFactory = func(conf *config.ServiceConfig, log *zap.Logger) checkpointFactory {
			return &mockCheckpointFactory{err: errors.New("failed")}
		}

		conf := &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{}}

		_, err := NewNextGenRepoterLoader(conf, logger)
		require.Error(t, err)
	})
}

func TestNextGenRepoterLoader_LoadCheckpoint(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		mgr := &mockCheckpointManager{checkpoint: 123}
		loader := &NextGenRepoterLoaderImpl{
			log:           zap.NewNop(),
			checkpointMgr: mgr,
		}

		checkpoint, err := loader.LoadCheckpoint(context.Background())
		require.NoError(t, err)
		assert.Equal(t, int64(123), checkpoint)
		assert.Equal(t, nextgenReporterTaskName, mgr.lastGetTask)
	})

	t.Run("error", func(t *testing.T) {
		mgr := &mockCheckpointManager{getErr: errors.New("fetch failed")}
		loader := &NextGenRepoterLoaderImpl{
			log:           zap.NewNop(),
			checkpointMgr: mgr,
		}

		_, err := loader.LoadCheckpoint(context.Background())
		require.Error(t, err)
		assert.Equal(t, nextgenReporterTaskName, mgr.lastGetTask)
	})
}

func TestNextGenRepoterLoader_SaveCheckpoint(t *testing.T) {
	t.Run("success", func(t *testing.T) {
		mgr := &mockCheckpointManager{}
		loader := &NextGenRepoterLoaderImpl{
			log:           zap.NewNop(),
			checkpointMgr: mgr,
		}

		err := loader.SaveCheckpoint(context.Background(), 456)
		require.NoError(t, err)
		assert.Equal(t, int64(456), mgr.checkpoint)
		assert.Equal(t, nextgenReporterTaskName, mgr.lastSaveTask)
		assert.Equal(t, int64(456), mgr.lastSavedValue)
	})

	t.Run("error", func(t *testing.T) {
		mgr := &mockCheckpointManager{saveErr: errors.New("persist failed")}
		loader := &NextGenRepoterLoaderImpl{
			log:           zap.NewNop(),
			checkpointMgr: mgr,
		}

		err := loader.SaveCheckpoint(context.Background(), 789)
		require.Error(t, err)
		assert.Equal(t, nextgenReporterTaskName, mgr.lastSaveTask)
		assert.Equal(t, int64(789), mgr.lastSavedValue)
	})
}

func TestNextGenRepoterLoader_LoadMeteringDataInTimeRange(t *testing.T) {
	baseConf := &config.ServiceConfig{
		NextGenMeteringConfig: &config.NextGenMeteringConfig{DeltaTablePath: "s3://bucket/table"},
	}

	t.Run("empty_delta_path", func(t *testing.T) {
		loader := &NextGenRepoterLoaderImpl{
			conf: &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{}},
			log:  zap.NewNop(),
		}

		reports, err := loader.LoadMeteringDataInTimeRange(context.Background(), 1700000000, 1700003600)
		require.NoError(t, err)
		assert.Empty(t, reports)
	})

	t.Run("nil_factory_uses_default", func(t *testing.T) {
		original := defaultDuckDBFactory
		mockClient := &mockDuckDBClient{err: errors.New("query failure")}
		defaultDuckDBFactory = &mockDuckDBClientFactory{client: mockClient}
		t.Cleanup(func() { defaultDuckDBFactory = original })

		loader := &NextGenRepoterLoaderImpl{
			conf: baseConf,
			log:  zap.NewNop(),
		}

		_, err := loader.LoadMeteringDataInTimeRange(context.Background(), 1700000000, 1700003600)
		require.Error(t, err)
		assert.True(t, mockClient.closeCalled)
	})

	t.Run("factory_error", func(t *testing.T) {
		loader := &NextGenRepoterLoaderImpl{
			conf:        baseConf,
			log:         zap.NewNop(),
			duckFactory: &mockDuckDBClientFactory{err: errors.New("factory failure")},
		}

		_, err := loader.LoadMeteringDataInTimeRange(context.Background(), 1700000000, 1700003600)
		require.Error(t, err)
	})

	t.Run("query_error", func(t *testing.T) {
		mockClient := &mockDuckDBClient{err: errors.New("query failure")}
		loader := &NextGenRepoterLoaderImpl{
			conf:        baseConf,
			log:         zap.NewNop(),
			duckFactory: &mockDuckDBClientFactory{client: mockClient},
		}

		_, err := loader.LoadMeteringDataInTimeRange(context.Background(), 1700000000, 1700003600)
		require.Error(t, err)
		assert.True(t, mockClient.closeCalled)
	})

	t.Run("parse_error", func(t *testing.T) {
		db, mock, err := sqlmock.New()
		require.NoError(t, err)
		t.Cleanup(func() { require.NoError(t, mock.ExpectationsWereMet()) })
		t.Cleanup(func() { db.Close() })

		columns := []string{"scheme_version", "tenant_id", "cost_id", "cost_time_stamp", "unit", "unit_value", "cost", "labels", "created_at", "sku_type", "hour_dimension"}
		mockRows := sqlmock.NewRows(columns).
			AddRow("v1", "123", "cost-1", "1700000000", "bytes", "1", "100", "{}", "1700000000", models.NextGenSKUTypePublic, "2024010101").
			RowError(0, errors.New("row error"))

		mock.ExpectQuery("SELECT 1").WillReturnRows(mockRows)

		sqlRows, err := db.Query("SELECT 1")
		require.NoError(t, err)

		mockClient := &mockDuckDBClient{rows: sqlRows}
		loader := &NextGenRepoterLoaderImpl{
			conf:        baseConf,
			log:         zap.NewNop(),
			duckFactory: &mockDuckDBClientFactory{client: mockClient},
		}

		_, err = loader.LoadMeteringDataInTimeRange(context.Background(), 1700000000, 1700003600)
		require.Error(t, err)
		assert.True(t, mockClient.closeCalled)
	})

	t.Run("success", func(t *testing.T) {
		db, mock, err := sqlmock.New()
		require.NoError(t, err)
		t.Cleanup(func() { require.NoError(t, mock.ExpectationsWereMet()) })
		t.Cleanup(func() { db.Close() })

		columns := []string{"scheme_version", "tenant_id", "cost_id", "cost_time_stamp", "unit", "unit_value", "cost", "labels", "created_at", "sku_type", "hour_dimension"}
		mockRows := sqlmock.NewRows(columns).
			AddRow("v1", "123", "cost-1", "1700000000", "bytes", "1", "100", "{}", "1700000000", models.NextGenSKUTypePublic, "2024010101").
			AddRow("v1", "invalid", "cost-2", "invalid", "unit", "value", "bad", "{}", "abc", models.NextGenSKUTypePublic, "2024010102")

		mock.ExpectQuery("SELECT 1").WillReturnRows(mockRows)

		sqlRows, err := db.Query("SELECT 1")
		require.NoError(t, err)

		mockClient := &mockDuckDBClient{rows: sqlRows}
		loader := &NextGenRepoterLoaderImpl{
			conf:        baseConf,
			log:         zap.NewNop(),
			duckFactory: &mockDuckDBClientFactory{client: mockClient},
		}

		reports, err := loader.LoadMeteringDataInTimeRange(context.Background(), 1700000000, 1700003600)
		require.NoError(t, err)
		require.Len(t, reports, 1)
		report := reports[0]
		assert.Equal(t, "v1", report.SchemeVersion)
		assert.Equal(t, uint64(123), report.TenantID)
		assert.Equal(t, "cost-1", report.CostID)
		assert.Equal(t, int64(1700000000), report.CostTimestamp)
		assert.Equal(t, "bytes", report.Unit)
		assert.Equal(t, "1", report.UnitValue)
		assert.Equal(t, int64(100), report.Cost)
		assert.Equal(t, "{}", report.Labels)
		assert.Equal(t, int64(1700000000), report.CreatedAt)
		assert.Equal(t, models.NextGenSKUTypePublic, report.SKUType)
		assert.Equal(t, "2024010101", report.HourDimension)
		assert.True(t, mockClient.closeCalled)
		assert.Equal(t, "s3://bucket/table", mockClient.lastTablePath)
		assert.Contains(t, mockClient.lastQuery, "created_at > 1700000000")
		assert.Contains(t, mockClient.lastQuery, "sku_type = 'public'")
	})
}

func TestNextGenRepoterLoader_LoadMeteringDataInTimeRangeWithDeltaRS(t *testing.T) {
	baseConf := &config.ServiceConfig{
		NextGenMeteringConfig: &config.NextGenMeteringConfig{DeltaTablePath: "s3://bucket/table"},
	}

	t.Run("empty_delta_path", func(t *testing.T) {
		loader := &NextGenRepoterLoaderImpl{
			conf: &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{}},
			log:  zap.NewNop(),
		}

		reports, err := loader.LoadMeteringDataInTimeRangeWithDeltaRS(context.Background(), 1700000000, 1700003600)
		require.NoError(t, err)
		assert.Empty(t, reports)
	})

	t.Run("nil_table_manager", func(t *testing.T) {
		loader := &NextGenRepoterLoaderImpl{
			conf:         baseConf,
			log:          zap.NewNop(),
			tableManager: nil, // tableManager not set
		}

		_, err := loader.LoadMeteringDataInTimeRangeWithDeltaRS(context.Background(), 1700000000, 1700003600)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "table manager not initialized")
	})

	t.Run("open_table_error", func(t *testing.T) {
		mockTM := &mockTableManager{openErr: errors.New("failed to open table")}
		loader := &NextGenRepoterLoaderImpl{
			conf:         baseConf,
			log:          zap.NewNop(),
			tableManager: mockTM,
		}

		_, err := loader.LoadMeteringDataInTimeRangeWithDeltaRS(context.Background(), 1700000000, 1700003600)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "failed to open delta table")
		assert.True(t, mockTM.openCalled)
	})

	t.Run("query_error", func(t *testing.T) {
		mockTM := &mockTableManager{queryErr: errors.New("query failed")}
		loader := &NextGenRepoterLoaderImpl{
			conf:         baseConf,
			log:          zap.NewNop(),
			tableManager: mockTM,
		}

		_, err := loader.LoadMeteringDataInTimeRangeWithDeltaRS(context.Background(), 1700000000, 1700003600)
		require.Error(t, err)
		assert.Contains(t, mockTM.lastSQL, "created_at > 1700000000")
	})

	t.Run("success", func(t *testing.T) {
		rows := []map[string]interface{}{
			{
				"scheme_version":  "v1",
				"tenant_id":       float64(321),
				"cost_id":         "cost-1",
				"cost_time_stamp": "1700000000",
				"unit":            "bytes",
				"unit_value":      "2",
				"cost":            int64(200),
				"labels":          "{}",
				"created_at":      int64(1700000010),
				"sku_type":        models.NextGenSKUTypePublic,
				"hour_dimension":  "2024010101",
			},
			{
				"tenant_id": int64(0),
				"cost_id":   "",
			},
		}

		mockTM := &mockTableManager{queryResult: rows}
		loader := &NextGenRepoterLoaderImpl{
			conf:         baseConf,
			log:          zap.NewNop(),
			tableManager: mockTM,
		}

		reports, err := loader.LoadMeteringDataInTimeRangeWithDeltaRS(context.Background(), 1700000000, 1700003600)
		require.NoError(t, err)
		require.Len(t, reports, 1)

		report := reports[0]
		assert.Equal(t, "v1", report.SchemeVersion)
		assert.Equal(t, uint64(321), report.TenantID)
		assert.Equal(t, "cost-1", report.CostID)
		assert.Equal(t, int64(1700000000), report.CostTimestamp)
		assert.Equal(t, "bytes", report.Unit)
		assert.Equal(t, "2", report.UnitValue)
		assert.Equal(t, int64(200), report.Cost)
		assert.Equal(t, "{}", report.Labels)
		assert.Equal(t, int64(1700000010), report.CreatedAt)
		assert.Equal(t, models.NextGenSKUTypePublic, report.SKUType)
		assert.Equal(t, "2024010101", report.HourDimension)
		assert.Contains(t, mockTM.lastSQL, "metering_table")
	})
}

func TestConvertToInt64(t *testing.T) {
	loader := &NextGenRepoterLoaderImpl{log: zap.NewNop()}

	tests := []struct {
		name    string
		input   interface{}
		want    int64
		wantErr bool
	}{
		{name: "int64", input: int64(42), want: 42},
		{name: "int", input: int(7), want: 7},
		{name: "int32", input: int32(11), want: 11},
		{name: "float64", input: float64(12.3), want: 12},
		{name: "float32", input: float32(22.7), want: 22},
		{name: "string", input: "99", want: 99},
		{name: "invalid_string", input: "abc", wantErr: true},
		{name: "unsupported", input: true, wantErr: true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := loader.convertToInt64(tt.input)
			if tt.wantErr {
				require.Error(t, err)
				return
			}

			require.NoError(t, err)
			assert.Equal(t, tt.want, got)
		})
	}
}

func TestConvertToUint64(t *testing.T) {
	loader := &NextGenRepoterLoaderImpl{log: zap.NewNop()}

	tests := []struct {
		name    string
		input   interface{}
		want    uint64
		wantErr bool
	}{
		{name: "uint64", input: uint64(5), want: 5},
		{name: "int64", input: int64(8), want: 8},
		{name: "int", input: int(6), want: 6},
		{name: "int32", input: int32(18), want: 18},
		{name: "uint32", input: uint32(12), want: 12},
		{name: "float64", input: float64(15.9), want: 15},
		{name: "float32", input: float32(19.4), want: 19},
		{name: "string", input: "77", want: 77},
		{name: "negative", input: int64(-1), wantErr: true},
		{name: "negative_float", input: float64(-3.2), wantErr: true},
		{name: "negative_float32", input: float32(-2.5), wantErr: true},
		{name: "invalid_string", input: "bad", wantErr: true},
		{name: "unsupported", input: []byte("10"), wantErr: true},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, err := loader.convertToUint64(tt.input)
			if tt.wantErr {
				require.Error(t, err)
				return
			}

			require.NoError(t, err)
			assert.Equal(t, tt.want, got)
		})
	}
}

// TestBuildQuery tests the hour_dimension logic in the buildQuery method
func TestBuildQuery(t *testing.T) {
	// Create a NextGenRepoterLoaderImpl instance for testing
	loader := &NextGenRepoterLoaderImpl{}

	tests := []struct {
		name                 string
		startTime            int64
		endTime              int64 // Keep for test data context, but not used in buildQuery
		expectedStartHourDim string
		expectedContainsSQL  []string
	}{
		{
			name:                 "query_within_same_hour",
			startTime:            1704067200, // 2024-01-01 00:00:00 UTC
			endTime:              1704070800, // 2024-01-01 01:00:00 UTC
			expectedStartHourDim: "2024010100",
			expectedContainsSQL: []string{
				"hour_dimension >= '2024010100'",
				"created_at > 1704067200",
			},
		},
		{
			name:                 "query_across_multiple_hours",
			startTime:            1704067200, // 2024-01-01 00:00:00 UTC
			endTime:              1704081600, // 2024-01-01 04:00:00 UTC
			expectedStartHourDim: "2024010100",
			expectedContainsSQL: []string{
				"hour_dimension >= '2024010100'",
				"created_at > 1704067200",
			},
		},
		{
			name:                 "non_hour_boundary_start_and_end",
			startTime:            1704067800, // 2024-01-01 00:10:00 UTC
			endTime:              1704070200, // 2024-01-01 00:50:00 UTC
			expectedStartHourDim: "2024010100",
			expectedContainsSQL: []string{
				"hour_dimension >= '2024010100'",
				"created_at > 1704067800",
			},
		},
		{
			name:                 "cross_day_query",
			startTime:            1704063600, // 2023-12-31 23:00:00 UTC
			endTime:              1704070800, // 2024-01-01 01:00:00 UTC
			expectedStartHourDim: "2023123123",
			expectedContainsSQL: []string{
				"hour_dimension >= '2023123123'",
				"created_at > 1704063600",
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			query := loader.buildQuery(tt.startTime)

			// Verify that the query contains all expected SQL fragments
			for _, expectedSQL := range tt.expectedContainsSQL {
				assert.Contains(t, query, expectedSQL,
					"query should contain '%s'", expectedSQL)
			}

			// Verify basic query structure
			assert.Contains(t, query, "SELECT", "query should be a SELECT statement")
			assert.Contains(t, query, "FROM {table_path}", "query should contain table path placeholder")
			assert.Contains(t, query, "WHERE", "query should contain WHERE condition")
			assert.Contains(t, query, "ORDER BY created_at", "query should be ordered by created_at")

			t.Logf("Generated query for %s:\n%s", tt.name, query)
		})
	}
}

// TestHourDimensionCalculation tests the hour_dimension calculation logic
func TestHourDimensionCalculation(t *testing.T) {
	tests := []struct {
		name                  string
		timestamp             int64
		expectedHourDim       string
		expectedTruncatedTime string
	}{
		{
			name:                  "exact_hour_time",
			timestamp:             1704067200, // 2024-01-01 00:00:00 UTC
			expectedHourDim:       "2024010100",
			expectedTruncatedTime: "2024-01-01 00:00:00",
		},
		{
			name:                  "non_exact_hour_time",
			timestamp:             1704067800, // 2024-01-01 00:10:00 UTC
			expectedHourDim:       "2024010100",
			expectedTruncatedTime: "2024-01-01 00:00:00",
		},
		{
			name:                  "before_midnight",
			timestamp:             1704063600, // 2023-12-31 23:00:00 UTC
			expectedHourDim:       "2023123123",
			expectedTruncatedTime: "2023-12-31 23:00:00",
		},
		{
			name:                  "year_end_cross_year",
			timestamp:             1704067199, // 2023-12-31 23:59:59 UTC
			expectedHourDim:       "2023123123",
			expectedTruncatedTime: "2023-12-31 23:00:00",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Simulate the time calculation logic in buildQuery
			timeObj := time.Unix(tt.timestamp, 0).UTC()
			truncatedHour := timeObj.Truncate(time.Hour)
			hourDimension := truncatedHour.Format("2006010215")

			assert.Equal(t, tt.expectedHourDim, hourDimension,
				"hour_dimension calculation is incorrect")
			assert.Equal(t, tt.expectedTruncatedTime,
				truncatedHour.Format("2006-01-02 15:04:05"),
				"truncated time is incorrect")

			t.Logf("Timestamp %d (%s) -> hour_dimension: %s, truncated: %s",
				tt.timestamp,
				time.Unix(tt.timestamp, 0).UTC().Format("2006-01-02 15:04:05"),
				hourDimension,
				truncatedHour.Format("2006-01-02 15:04:05"))
		})
	}
}

// TestPartitionPruningLogic tests the partition pruning logic
func TestPartitionPruningLogic(t *testing.T) {
	tests := []struct {
		name               string
		startTime          int64
		endTime            int64
		expectedPartitions []string
		description        string
	}{
		{
			name:               "single_hour_partition",
			startTime:          1704067200, // 2024-01-01 00:00:00 UTC
			endTime:            1704070800, // 2024-01-01 01:00:00 UTC
			expectedPartitions: []string{"2024010100"},
			description:        "should only access one partition",
		},
		{
			name:               "multiple_hour_partitions",
			startTime:          1704067200, // 2024-01-01 00:00:00 UTC
			endTime:            1704081600, // 2024-01-01 04:00:00 UTC
			expectedPartitions: []string{"2024010100", "2024010101", "2024010102", "2024010103"},
			description:        "should access 4 consecutive partitions",
		},
		{
			name:               "cross_day_partitions",
			startTime:          1704063600, // 2023-12-31 23:00:00 UTC
			endTime:            1704074400, // 2024-01-01 02:00:00 UTC
			expectedPartitions: []string{"2023123123", "2024010100", "2024010101"},
			description:        "should span partitions across two days",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Simulate partition calculation logic
			startTime := time.Unix(tt.startTime, 0).UTC()
			endTime := time.Unix(tt.endTime, 0).UTC()

			startHour := startTime.Truncate(time.Hour)
			endHour := endTime.Truncate(time.Hour)
			if endTime.After(endHour) {
				endHour = endHour.Add(time.Hour)
			}

			// Calculate the actual partitions that will be accessed
			var actualPartitions []string
			current := startHour
			for current.Before(endHour) {
				partition := current.Format("2006010215")
				actualPartitions = append(actualPartitions, partition)
				current = current.Add(time.Hour)
			}

			assert.Equal(t, tt.expectedPartitions, actualPartitions,
				"partition calculation is incorrect: %s", tt.description)

			t.Logf("Test: %s", tt.name)
			t.Logf("Time range: %s to %s",
				startTime.Format("2006-01-02 15:04:05"),
				endTime.Format("2006-01-02 15:04:05"))
			t.Logf("Partitions accessed: %v", actualPartitions)
			t.Logf("Description: %s", tt.description)
		})
	}
}

type mockCheckpointFactory struct {
	manager checkpoint.NextGenCheckpoint
	err     error
}

func (m *mockCheckpointFactory) CreateCheckpointManager() (checkpoint.NextGenCheckpoint, error) {
	if m.err != nil {
		return nil, m.err
	}
	return m.manager, nil
}

type mockCheckpointManager struct {
	checkpoint     int64
	getErr         error
	saveErr        error
	lastGetTask    string
	lastSaveTask   string
	lastSavedValue int64
}

func (m *mockCheckpointManager) GetCheckpoint(ctx context.Context, taskName string) (int64, error) {
	m.lastGetTask = taskName
	if m.getErr != nil {
		return 0, m.getErr
	}
	return m.checkpoint, nil
}

func (m *mockCheckpointManager) SaveCheckpoint(ctx context.Context, taskName string, checkpoint int64) error {
	m.lastSaveTask = taskName
	m.lastSavedValue = checkpoint
	if m.saveErr != nil {
		return m.saveErr
	}
	m.checkpoint = checkpoint
	return nil
}

func (m *mockCheckpointManager) GetCheckpointData(ctx context.Context, taskName string) (*checkpoint.CheckpointData, error) {
	return nil, nil
}

func (m *mockCheckpointManager) SaveCheckpointData(ctx context.Context, data *checkpoint.CheckpointData) error {
	return nil
}

type mockDuckDBClientFactory struct {
	client duckDBClient
	err    error
}

func (m *mockDuckDBClientFactory) NewDuckDBClientWithConfig(conf *config.ServiceConfig, log *zap.Logger) (duckDBClient, error) {
	if m.err != nil {
		return nil, m.err
	}
	return m.client, nil
}

type mockDuckDBClient struct {
	rows          *sql.Rows
	err           error
	closeCalled   bool
	lastTablePath string
	lastQuery     string
}

func (m *mockDuckDBClient) QueryData(tablePath, customSQL string) (*sql.Rows, error) {
	m.lastTablePath = tablePath
	m.lastQuery = customSQL
	if m.err != nil {
		return nil, m.err
	}
	return m.rows, nil
}

func (m *mockDuckDBClient) Close() error {
	m.closeCalled = true
	return nil
}