package nextgentasks

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

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/metering/pkg/common/checkpoint"
	"github.com/tidbcloud/metering/pkg/common/config"
	"github.com/tidbcloud/metering/pkg/common/constants"
	"github.com/tidbcloud/metering/pkg/common/models"
	"github.com/tidbcloud/metering/test/mockcheckpoint"
	"go.uber.org/mock/gomock"
	"go.uber.org/zap"
	"go.uber.org/zap/zaptest"
)

type getVersionResponse struct {
	value int64
	err   error
}

type fakeMeteringService struct {
	tableExistsResponses     []bool
	tableExistsErr           error
	tableExistsCalls         int
	createTableErr           error
	createTableCalls         int
	processMinuteRecords     map[int64][]models.BillingReportV2Deltalake
	processMinuteErrors      map[int64]error
	processMinuteCalls       []int64
	writeDeltaTableRecords   []models.BillingReportV2Deltalake
	writeDeltaTableErr       error
	writeDeltaTableCalls     int
	getTableVersionResponses []getVersionResponse
	getMaxTimestampResponses []getVersionResponse
	closeErr                 error
	closed                   bool
}

func (f *fakeMeteringService) TableExists(ctx context.Context) (bool, error) {
	f.tableExistsCalls++
	if f.tableExistsErr != nil {
		return false, f.tableExistsErr
	}
	if len(f.tableExistsResponses) == 0 {
		return true, nil
	}
	result := f.tableExistsResponses[0]
	f.tableExistsResponses = f.tableExistsResponses[1:]
	return result, nil
}

func (f *fakeMeteringService) CreateDeltaTable() error {
	f.createTableCalls++
	return f.createTableErr
}

func (f *fakeMeteringService) ProcessMinute(ctx context.Context, timestamp int64) ([]models.BillingReportV2Deltalake, error) {
	f.processMinuteCalls = append(f.processMinuteCalls, timestamp)
	if err := f.processMinuteErrors[timestamp]; err != nil {
		return nil, err
	}
	if records, ok := f.processMinuteRecords[timestamp]; ok {
		return records, nil
	}
	// Return empty slice by default
	return []models.BillingReportV2Deltalake{}, nil
}

func (f *fakeMeteringService) WriteToDeltaTable(records []models.BillingReportV2Deltalake, saveMode string) error {
	f.writeDeltaTableCalls++
	f.writeDeltaTableRecords = append(f.writeDeltaTableRecords, records...)
	return f.writeDeltaTableErr
}

func (f *fakeMeteringService) Close() error {
	f.closed = true
	return f.closeErr
}

func (f *fakeMeteringService) GetTableVersion(ctx context.Context) (int64, error) {
	if len(f.getTableVersionResponses) == 0 {
		return 0, nil
	}
	resp := f.getTableVersionResponses[0]
	f.getTableVersionResponses = f.getTableVersionResponses[1:]
	if resp.err != nil {
		return 0, resp.err
	}
	return resp.value, nil
}

func (f *fakeMeteringService) QueryMaxTimestampWithPartitionFilter(ctx context.Context, partitionTime int64) (int64, error) {
	if len(f.getMaxTimestampResponses) == 0 {
		return 0, nil
	}
	resp := f.getMaxTimestampResponses[0]
	f.getMaxTimestampResponses = f.getMaxTimestampResponses[1:]
	if resp.err != nil {
		return 0, resp.err
	}
	return resp.value, nil
}

type shouldRecoverResponse struct {
	shouldRecover bool
	version       int64
	err           error
}

type recoverResponse struct {
	data *checkpoint.CheckpointData
	err  error
}

type fakeDeltaHelper struct {
	shouldRecoverResponses []shouldRecoverResponse
	recoverResponses       []recoverResponse
	shouldRecoverCalls     int
	recoverCalls           int
}

func (f *fakeDeltaHelper) ShouldRecoverFromTable(ctx context.Context, _ checkpoint.MeteringServiceInterface, _ *checkpoint.CheckpointData) (bool, int64, error) {
	f.shouldRecoverCalls++
	if len(f.shouldRecoverResponses) == 0 {
		return false, 0, nil
	}
	resp := f.shouldRecoverResponses[0]
	f.shouldRecoverResponses = f.shouldRecoverResponses[1:]
	return resp.shouldRecover, resp.version, resp.err
}

func (f *fakeDeltaHelper) RecoverCheckpointFromTable(ctx context.Context, _ checkpoint.MeteringServiceInterface, currentTableVersion int64, checkpointData *checkpoint.CheckpointData) (*checkpoint.CheckpointData, error) {
	f.recoverCalls++
	if len(f.recoverResponses) == 0 {
		return &checkpoint.CheckpointData{DeltaTableVersion: currentTableVersion}, nil
	}
	resp := f.recoverResponses[0]
	f.recoverResponses = f.recoverResponses[1:]
	if resp.data == nil {
		return nil, resp.err
	}
	return resp.data, resp.err
}

func overrideFactories(t *testing.T, ms meteringService, msErr error, helper deltaLakeHelper, nowFunc func() time.Time) func() {
	t.Helper()
	originalMetering := newNextGenMeteringService
	originalHelper := newDeltaLakeCheckpointHelper
	originalNow := timeNow

	if ms != nil || msErr != nil {
		newNextGenMeteringService = func(*config.ServiceConfig, *zap.Logger) (meteringService, error) {
			return ms, msErr
		}
	} else {
		newNextGenMeteringService = originalMetering
	}

	if helper != nil {
		newDeltaLakeCheckpointHelper = func(*config.ServiceConfig, *zap.Logger) deltaLakeHelper {
			return helper
		}
	} else {
		newDeltaLakeCheckpointHelper = originalHelper
	}

	if nowFunc != nil {
		timeNow = nowFunc
	} else {
		timeNow = originalNow
	}

	return func() {
		newNextGenMeteringService = originalMetering
		newDeltaLakeCheckpointHelper = originalHelper
		timeNow = originalNow
	}
}

func resetGlobalDeltaLakeLock() {
	globalDeltaLakeLock = nil
}

type fakeCheckpointFactory struct {
	manager checkpoint.NextGenCheckpoint
	err     error
}

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

func overrideCheckpointFactory(t *testing.T, factory checkpointFactory) func() {
	t.Helper()
	original := newCheckpointFactory
	newCheckpointFactory = func(*config.ServiceConfig, *zap.Logger) checkpointFactory {
		return factory
	}
	return func() {
		newCheckpointFactory = original
	}
}

func TestNextGenTask_Run_NoCheckpoint_NoWork(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	resetGlobalDeltaLakeLock()
	mockCheckpoint := mockcheckpoint.NewMockNextGenCheckpoint(ctrl)
	InitGlobalDeltaLakeLock(zaptest.NewLogger(t))

	fixedNow := time.Date(2025, time.September, 25, 10, 10, 0, 0, time.UTC)
	delayMinutes := int64(5)
	rounded := time.Date(fixedNow.Year(), fixedNow.Month(), fixedNow.Day(), fixedNow.Hour(), fixedNow.Minute(), 0, 0, time.UTC)
	endTime := rounded.Add(-time.Duration(delayMinutes) * time.Minute).Unix()

	checkpointData := &checkpoint.CheckpointData{
		Checkpoint:        0,
		TaskName:          "nextgen-metering",
		DeltaTablePath:    "",
		DeltaTableVersion: 0,
	}
	mockCheckpoint.EXPECT().GetCheckpointData(gomock.Any(), "nextgen-metering").Return(checkpointData, nil)
	mockCheckpoint.EXPECT().SaveCheckpointData(gomock.Any(), gomock.Any()).Times(0)

	conf := &config.ServiceConfig{
		Vendor: "aws",
		Bucket: "test-bucket",
		Region: "us-west-2",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			StartTime:    endTime,
			DelayMinutes: int(delayMinutes),
		},
	}

	fakeMeter := &fakeMeteringService{}
	fakeHelper := &fakeDeltaHelper{}
	restore := overrideFactories(t, fakeMeter, nil, fakeHelper, func() time.Time { return fixedNow })
	defer restore()

	task, err := NewNextGenTaskWithCheckpointManager(conf, zaptest.NewLogger(t), mockCheckpoint)
	require.NoError(t, err)

	err = task.Run(context.Background())
	require.NoError(t, err)

	assert.True(t, fakeMeter.closed)
	assert.Empty(t, fakeMeter.processMinuteCalls)
	assert.Equal(t, 0, fakeHelper.shouldRecoverCalls)
	assert.Equal(t, 0, fakeHelper.recoverCalls)
}

func TestNextGenTask_Run_CreatesTableAndProcesses(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	resetGlobalDeltaLakeLock()
	mockCheckpoint := mockcheckpoint.NewMockNextGenCheckpoint(ctrl)
	InitGlobalDeltaLakeLock(zaptest.NewLogger(t))

	fixedNow := time.Date(2025, time.September, 25, 10, 10, 0, 0, time.UTC)
	startTime := fixedNow.Add(-8 * time.Minute).Unix()
	deltaPath := "s3://delta/table"

	checkpointData := &checkpoint.CheckpointData{
		Checkpoint:        0,
		TaskName:          "nextgen-metering",
		DeltaTablePath:    "",
		DeltaTableVersion: 0,
	}
	mockCheckpoint.EXPECT().GetCheckpointData(gomock.Any(), "nextgen-metering").Return(checkpointData, nil)

	expectedTimestamps := []int64{startTime, startTime + 60, startTime + 120}
	finalTableVersion := int64(41) // Version after batch write
	saveIndex := 0
	// In batch processing mode, only save checkpoint once with the last timestamp
	mockCheckpoint.EXPECT().SaveCheckpointData(gomock.Any(), gomock.AssignableToTypeOf(&checkpoint.CheckpointData{})).
		Times(1).
		DoAndReturn(func(_ context.Context, data *checkpoint.CheckpointData) error {
			require.Equal(t, expectedTimestamps[len(expectedTimestamps)-1], data.Checkpoint) // Last timestamp
			require.Equal(t, finalTableVersion, data.DeltaTableVersion)
			require.Equal(t, deltaPath, data.DeltaTablePath)
			saveIndex++
			return nil
		})

	fakeMeter := &fakeMeteringService{
		tableExistsResponses: []bool{false},
		processMinuteRecords: map[int64][]models.BillingReportV2Deltalake{
			expectedTimestamps[0]: {{SchemeVersion: "v1"}}, // Add some test data
			expectedTimestamps[1]: {{SchemeVersion: "v1"}},
			expectedTimestamps[2]: {{SchemeVersion: "v1"}},
		},
		getTableVersionResponses: []getVersionResponse{
			{value: finalTableVersion}, // Version after batch write
		},
	}
	fakeHelper := &fakeDeltaHelper{
		shouldRecoverResponses: []shouldRecoverResponse{{shouldRecover: false, version: 40}},
	}
	restore := overrideFactories(t, fakeMeter, nil, fakeHelper, func() time.Time { return fixedNow })
	defer restore()

	conf := &config.ServiceConfig{
		Vendor: "aws",
		Bucket: "test-bucket",
		Region: "us-west-2",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			StartTime:      startTime,
			DelayMinutes:   5,
			DeltaTablePath: deltaPath,
		},
	}

	task, err := NewNextGenTaskWithCheckpointManager(conf, zaptest.NewLogger(t), mockCheckpoint)
	require.NoError(t, err)

	err = task.Run(context.Background())
	require.NoError(t, err)
	require.Equal(t, 1, saveIndex) // Only save checkpoint once after batch write

	assert.Equal(t, expectedTimestamps, fakeMeter.processMinuteCalls)
	assert.Equal(t, 1, fakeMeter.createTableCalls)
	assert.Equal(t, 1, fakeMeter.writeDeltaTableCalls) // Only one batch write
	assert.True(t, fakeMeter.closed)
	assert.Equal(t, 1, fakeHelper.shouldRecoverCalls)
	assert.Equal(t, 0, fakeHelper.recoverCalls)
}

func TestNextGenTask_Run_MeteringFailure(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	resetGlobalDeltaLakeLock()
	mockCheckpoint := mockcheckpoint.NewMockNextGenCheckpoint(ctrl)
	InitGlobalDeltaLakeLock(zaptest.NewLogger(t))

	fixedNow := time.Date(2025, time.September, 25, 10, 10, 0, 0, time.UTC)
	startTime := fixedNow.Add(-6 * time.Minute).Unix()
	deltaPath := "s3://delta/table"

	checkpointData := &checkpoint.CheckpointData{
		Checkpoint:        0,
		TaskName:          "nextgen-metering",
		DeltaTablePath:    "",
		DeltaTableVersion: 0,
	}
	mockCheckpoint.EXPECT().GetCheckpointData(gomock.Any(), "nextgen-metering").Return(checkpointData, nil)
	mockCheckpoint.EXPECT().SaveCheckpointData(gomock.Any(), gomock.Any()).Times(0)

	fakeMeter := &fakeMeteringService{
		tableExistsResponses: []bool{true},
		processMinuteErrors: map[int64]error{
			startTime: errors.New("process minute failed"),
		},
	}
	restore := overrideFactories(t, fakeMeter, nil, &fakeDeltaHelper{}, func() time.Time { return fixedNow })
	defer restore()

	conf := &config.ServiceConfig{
		Vendor: "aws",
		Bucket: "test-bucket",
		Region: "us-west-2",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			StartTime:      startTime,
			DelayMinutes:   5,
			DeltaTablePath: deltaPath,
		},
	}

	task, err := NewNextGenTaskWithCheckpointManager(conf, zaptest.NewLogger(t), mockCheckpoint)
	require.NoError(t, err)

	err = task.Run(context.Background())
	require.Error(t, err)
	assert.Contains(t, err.Error(), "process minute failed")
	assert.Equal(t, []int64{startTime}, fakeMeter.processMinuteCalls)
	assert.True(t, fakeMeter.closed)
}

func TestNextGenTask_Run_RecoverFromTable(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	resetGlobalDeltaLakeLock()
	mockCheckpoint := mockcheckpoint.NewMockNextGenCheckpoint(ctrl)
	InitGlobalDeltaLakeLock(zaptest.NewLogger(t))

	fixedNow := time.Date(2025, time.September, 25, 10, 10, 0, 0, time.UTC)
	deltaPath := "s3://delta/table"
	initialCheckpoint := fixedNow.Add(-15 * time.Minute).Unix()
	recoveredCheckpoint := fixedNow.Add(-9 * time.Minute).Unix()

	checkpointData := &checkpoint.CheckpointData{
		Checkpoint:        initialCheckpoint,
		TaskName:          "nextgen-metering",
		DeltaTablePath:    deltaPath,
		DeltaTableVersion: 5,
	}
	mockCheckpoint.EXPECT().GetCheckpointData(gomock.Any(), "nextgen-metering").Return(checkpointData, nil)

	expectedTimestamps := []int64{recoveredCheckpoint + 60, recoveredCheckpoint + 120, recoveredCheckpoint + 180}
	finalTableVersion := int64(21) // Version after batch write
	recoveredData := &checkpoint.CheckpointData{
		TaskName:          "nextgen-metering",
		Checkpoint:        recoveredCheckpoint,
		DeltaTablePath:    deltaPath,
		DeltaTableVersion: 7,
	}

	// In batch processing mode, we save checkpoint twice:
	// 1. When recovering from table (recoveredCheckpoint with version 7)
	// 2. After batch write completes (last timestamp with final version)
	expectedCheckpoints := []int64{recoveredCheckpoint, expectedTimestamps[len(expectedTimestamps)-1]}
	expectedAllVersions := []int64{recoveredData.DeltaTableVersion, finalTableVersion}
	saveIndex := 0
	mockCheckpoint.EXPECT().SaveCheckpointData(gomock.Any(), gomock.AssignableToTypeOf(&checkpoint.CheckpointData{})).
		Times(2). // Changed from 1 to 2
		DoAndReturn(func(_ context.Context, data *checkpoint.CheckpointData) error {
			require.Equal(t, expectedCheckpoints[saveIndex], data.Checkpoint)
			require.Equal(t, expectedAllVersions[saveIndex], data.DeltaTableVersion)
			require.Equal(t, deltaPath, data.DeltaTablePath)
			saveIndex++
			return nil
		})

	fakeMeter := &fakeMeteringService{
		processMinuteRecords: map[int64][]models.BillingReportV2Deltalake{
			expectedTimestamps[0]: {{SchemeVersion: "v1"}}, // Add some test data
			expectedTimestamps[1]: {{SchemeVersion: "v1"}},
			expectedTimestamps[2]: {{SchemeVersion: "v1"}},
		},
		getTableVersionResponses: []getVersionResponse{
			{value: finalTableVersion}, // Version after batch write
		},
	}
	fakeHelper := &fakeDeltaHelper{
		shouldRecoverResponses: []shouldRecoverResponse{{shouldRecover: true, version: recoveredData.DeltaTableVersion}},
		recoverResponses:       []recoverResponse{{data: recoveredData}},
	}
	restore := overrideFactories(t, fakeMeter, nil, fakeHelper, func() time.Time { return fixedNow })
	defer restore()

	conf := &config.ServiceConfig{
		Vendor: "aws",
		Bucket: "test-bucket",
		Region: "us-west-2",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			StartTime:      initialCheckpoint,
			DelayMinutes:   5,
			DeltaTablePath: deltaPath,
		},
	}

	task, err := NewNextGenTaskWithCheckpointManager(conf, zaptest.NewLogger(t), mockCheckpoint)
	require.NoError(t, err)

	err = task.Run(context.Background())
	require.NoError(t, err)
	require.Equal(t, 2, saveIndex) // Save checkpoint twice: recovery + batch completion

	assert.Equal(t, expectedTimestamps, fakeMeter.processMinuteCalls)
	assert.Equal(t, 1, fakeMeter.writeDeltaTableCalls) // Only one batch write
	assert.Equal(t, 1, fakeHelper.shouldRecoverCalls)
	assert.Equal(t, 1, fakeHelper.recoverCalls)
	assert.True(t, fakeMeter.closed)
}

func TestNextGenTask_Run_TableExistsError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	resetGlobalDeltaLakeLock()
	mockCheckpoint := mockcheckpoint.NewMockNextGenCheckpoint(ctrl)
	InitGlobalDeltaLakeLock(zaptest.NewLogger(t))

	fixedNow := time.Date(2025, time.September, 25, 10, 10, 0, 0, time.UTC)
	startTime := fixedNow.Add(-7 * time.Minute).Unix()
	deltaPath := "s3://delta/table"

	checkpointData := &checkpoint.CheckpointData{
		Checkpoint:        0,
		TaskName:          "nextgen-metering",
		DeltaTablePath:    deltaPath,
		DeltaTableVersion: 0,
	}
	mockCheckpoint.EXPECT().GetCheckpointData(gomock.Any(), "nextgen-metering").Return(checkpointData, nil)
	mockCheckpoint.EXPECT().SaveCheckpointData(gomock.Any(), gomock.Any()).Times(0)

	tableErr := errors.New("failed to check table")
	fakeMeter := &fakeMeteringService{
		tableExistsErr: tableErr,
	}
	restore := overrideFactories(t, fakeMeter, nil, &fakeDeltaHelper{}, func() time.Time { return fixedNow })
	defer restore()

	conf := &config.ServiceConfig{
		Vendor: "aws",
		Bucket: "test-bucket",
		Region: "us-west-2",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			StartTime:      startTime,
			DelayMinutes:   5,
			DeltaTablePath: deltaPath,
		},
	}

	task, err := NewNextGenTaskWithCheckpointManager(conf, zaptest.NewLogger(t), mockCheckpoint)
	require.NoError(t, err)

	err = task.Run(context.Background())
	require.Error(t, err)
	assert.EqualError(t, err, tableErr.Error())
	assert.Equal(t, 1, fakeMeter.tableExistsCalls)
	assert.True(t, fakeMeter.closed)
}

func TestNextGenTask_Run_FailsFastWhenExistingTableCannotBeOpened(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	resetGlobalDeltaLakeLock()
	mockCheckpoint := mockcheckpoint.NewMockNextGenCheckpoint(ctrl)
	InitGlobalDeltaLakeLock(zaptest.NewLogger(t))

	fixedNow := time.Date(2025, time.September, 25, 10, 10, 0, 0, time.UTC)
	startTime := fixedNow.Add(-7 * time.Minute).Unix()
	deltaPath := "az://bucket/metering-data"

	checkpointData := &checkpoint.CheckpointData{
		Checkpoint:        0,
		TaskName:          "nextgen-metering",
		DeltaTablePath:    deltaPath,
		DeltaTableVersion: 0,
	}
	mockCheckpoint.EXPECT().GetCheckpointData(gomock.Any(), "nextgen-metering").Return(checkpointData, nil)
	mockCheckpoint.EXPECT().SaveCheckpointData(gomock.Any(), gomock.Any()).Times(0)

	openErr := errors.New("failed to open delta table at az://bucket/metering-data: Delta error (type 0): Not a Delta table: Generic delta kernel error: No files in log segment")
	fakeMeter := &fakeMeteringService{
		tableExistsResponses: []bool{true},
		getTableVersionResponses: []getVersionResponse{
			{err: openErr},
		},
	}
	restore := overrideFactories(t, fakeMeter, nil, &fakeDeltaHelper{}, func() time.Time { return fixedNow })
	defer restore()

	conf := &config.ServiceConfig{
		Vendor: "azure",
		Bucket: "test-bucket",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			StartTime:      startTime,
			DelayMinutes:   5,
			DeltaTablePath: deltaPath,
		},
	}

	task, err := NewNextGenTaskWithCheckpointManager(conf, zaptest.NewLogger(t), mockCheckpoint)
	require.NoError(t, err)

	err = task.Run(context.Background())
	require.Error(t, err)
	assert.Contains(t, err.Error(), "delta table exists but could not be opened during initialization")
	assert.ErrorIs(t, err, openErr)
	assert.Empty(t, fakeMeter.processMinuteCalls)
	assert.Equal(t, 0, fakeMeter.writeDeltaTableCalls)
	assert.Equal(t, 1, fakeMeter.tableExistsCalls)
	assert.True(t, fakeMeter.closed)
}

func TestNextGenTask_Run_ShouldRecoverCheckError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	resetGlobalDeltaLakeLock()
	mockCheckpoint := mockcheckpoint.NewMockNextGenCheckpoint(ctrl)
	InitGlobalDeltaLakeLock(zaptest.NewLogger(t))

	fixedNow := time.Date(2025, time.September, 25, 10, 10, 0, 0, time.UTC)
	deltaPath := "s3://delta/table"
	checkpointTimestamp := fixedNow.Add(-8 * time.Minute).Unix()

	checkpointData := &checkpoint.CheckpointData{
		Checkpoint:        checkpointTimestamp,
		TaskName:          "nextgen-metering",
		DeltaTablePath:    deltaPath,
		DeltaTableVersion: 10,
	}
	mockCheckpoint.EXPECT().GetCheckpointData(gomock.Any(), "nextgen-metering").Return(checkpointData, nil)

	startTime := checkpointTimestamp + 60
	expectedTimestamps := []int64{startTime, startTime + 60}
	finalTableVersion := int64(31) // Version after batch write
	saveIndex := 0
	// In batch processing mode, only save checkpoint once after all processing completes
	mockCheckpoint.EXPECT().SaveCheckpointData(gomock.Any(), gomock.AssignableToTypeOf(&checkpoint.CheckpointData{})).
		Times(1).
		DoAndReturn(func(_ context.Context, data *checkpoint.CheckpointData) error {
			require.Equal(t, expectedTimestamps[len(expectedTimestamps)-1], data.Checkpoint) // Last timestamp
			require.Equal(t, finalTableVersion, data.DeltaTableVersion)
			require.Equal(t, deltaPath, data.DeltaTablePath)
			saveIndex++
			return nil
		})

	shouldRecoverErr := errors.New("recover check failed")
	fakeMeter := &fakeMeteringService{
		processMinuteRecords: map[int64][]models.BillingReportV2Deltalake{
			expectedTimestamps[0]: {{SchemeVersion: "v1"}}, // Add some test data
			expectedTimestamps[1]: {{SchemeVersion: "v1"}},
		},
		getTableVersionResponses: []getVersionResponse{
			{value: finalTableVersion}, // Version after batch write
		},
	}
	fakeHelper := &fakeDeltaHelper{
		shouldRecoverResponses: []shouldRecoverResponse{{err: shouldRecoverErr}},
	}
	restore := overrideFactories(t, fakeMeter, nil, fakeHelper, func() time.Time { return fixedNow })
	defer restore()

	conf := &config.ServiceConfig{
		Vendor: "aws",
		Bucket: "test-bucket",
		Region: "us-west-2",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			StartTime:      checkpointTimestamp,
			DelayMinutes:   5,
			DeltaTablePath: deltaPath,
		},
	}

	task, err := NewNextGenTaskWithCheckpointManager(conf, zaptest.NewLogger(t), mockCheckpoint)
	require.NoError(t, err)

	err = task.Run(context.Background())
	require.NoError(t, err)
	require.Equal(t, 1, saveIndex) // Only save checkpoint once after batch write
	assert.Equal(t, expectedTimestamps, fakeMeter.processMinuteCalls)
	assert.Equal(t, 1, fakeMeter.writeDeltaTableCalls) // Only one batch write
	assert.Equal(t, 1, fakeHelper.shouldRecoverCalls)
	assert.Equal(t, 0, fakeHelper.recoverCalls)
	assert.True(t, fakeMeter.closed)
}

func TestNextGenTask_Run_GetTableVersionError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	resetGlobalDeltaLakeLock()
	mockCheckpoint := mockcheckpoint.NewMockNextGenCheckpoint(ctrl)
	InitGlobalDeltaLakeLock(zaptest.NewLogger(t))

	fixedNow := time.Date(2025, time.September, 25, 10, 10, 0, 0, time.UTC)
	deltaPath := "s3://delta/table"
	checkpointTimestamp := fixedNow.Add(-7 * time.Minute).Unix()

	checkpointData := &checkpoint.CheckpointData{
		Checkpoint:        checkpointTimestamp,
		TaskName:          "nextgen-metering",
		DeltaTablePath:    deltaPath,
		DeltaTableVersion: 12,
	}
	mockCheckpoint.EXPECT().GetCheckpointData(gomock.Any(), "nextgen-metering").Return(checkpointData, nil)

	expectedTimestamp := checkpointTimestamp + 60
	mockCheckpoint.EXPECT().SaveCheckpointData(gomock.Any(), gomock.AssignableToTypeOf(&checkpoint.CheckpointData{})).
		Times(1).
		DoAndReturn(func(_ context.Context, data *checkpoint.CheckpointData) error {
			require.Equal(t, expectedTimestamp, data.Checkpoint)
			require.Equal(t, int64(0), data.DeltaTableVersion)
			require.Equal(t, deltaPath, data.DeltaTablePath)
			return nil
		})

	versionErr := errors.New("get version failed")
	fakeMeter := &fakeMeteringService{
		processMinuteRecords: map[int64][]models.BillingReportV2Deltalake{
			expectedTimestamp: {{SchemeVersion: "v1"}}, // Add some test data
		},
		getTableVersionResponses: []getVersionResponse{{err: versionErr}},
	}
	restore := overrideFactories(t, fakeMeter, nil, &fakeDeltaHelper{}, func() time.Time { return fixedNow })
	defer restore()

	conf := &config.ServiceConfig{
		Vendor: "aws",
		Bucket: "test-bucket",
		Region: "us-west-2",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			StartTime:      checkpointTimestamp,
			DelayMinutes:   5,
			DeltaTablePath: deltaPath,
		},
	}

	task, err := NewNextGenTaskWithCheckpointManager(conf, zaptest.NewLogger(t), mockCheckpoint)
	require.NoError(t, err)

	err = task.Run(context.Background())
	require.NoError(t, err)
	assert.Equal(t, []int64{expectedTimestamp}, fakeMeter.processMinuteCalls)
	assert.Equal(t, 1, fakeMeter.writeDeltaTableCalls) // Only one batch write
	assert.True(t, fakeMeter.closed)
}

func TestNextGenTask_Run_SaveCheckpointFailure(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	resetGlobalDeltaLakeLock()
	mockCheckpoint := mockcheckpoint.NewMockNextGenCheckpoint(ctrl)
	InitGlobalDeltaLakeLock(zaptest.NewLogger(t))

	fixedNow := time.Date(2025, time.September, 25, 10, 10, 0, 0, time.UTC)
	deltaPath := "s3://delta/table"
	checkpointTimestamp := fixedNow.Add(-7 * time.Minute).Unix()

	checkpointData := &checkpoint.CheckpointData{
		Checkpoint:        checkpointTimestamp,
		TaskName:          "nextgen-metering",
		DeltaTablePath:    deltaPath,
		DeltaTableVersion: 15,
	}
	mockCheckpoint.EXPECT().GetCheckpointData(gomock.Any(), "nextgen-metering").Return(checkpointData, nil)

	expectedTimestamp := checkpointTimestamp + 60
	saveErr := errors.New("save failed")
	mockCheckpoint.EXPECT().SaveCheckpointData(gomock.Any(), gomock.AssignableToTypeOf(&checkpoint.CheckpointData{})).
		Times(1).
		Return(saveErr)

	fakeMeter := &fakeMeteringService{
		processMinuteRecords: map[int64][]models.BillingReportV2Deltalake{
			expectedTimestamp: {{SchemeVersion: "v1"}}, // Add some test data
		},
		getTableVersionResponses: []getVersionResponse{{value: 55}},
	}
	restore := overrideFactories(t, fakeMeter, nil, &fakeDeltaHelper{}, func() time.Time { return fixedNow })
	defer restore()

	conf := &config.ServiceConfig{
		Vendor: "aws",
		Bucket: "test-bucket",
		Region: "us-west-2",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			StartTime:      checkpointTimestamp,
			DelayMinutes:   5,
			DeltaTablePath: deltaPath,
		},
	}

	task, err := NewNextGenTaskWithCheckpointManager(conf, zaptest.NewLogger(t), mockCheckpoint)
	require.NoError(t, err)

	err = task.Run(context.Background())
	require.Error(t, err)
	assert.EqualError(t, err, saveErr.Error())
	assert.Equal(t, []int64{expectedTimestamp}, fakeMeter.processMinuteCalls)
	assert.Equal(t, 1, fakeMeter.writeDeltaTableCalls) // Batch write completed before checkpoint save failed
	assert.True(t, fakeMeter.closed)
}

func TestNextGenTask_Run_MeteringFactoryError(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	resetGlobalDeltaLakeLock()
	mockCheckpoint := mockcheckpoint.NewMockNextGenCheckpoint(ctrl)
	InitGlobalDeltaLakeLock(zaptest.NewLogger(t))

	factoryErr := errors.New("failed to create metering service")
	mockCheckpoint.EXPECT().GetCheckpointData(gomock.Any(), "nextgen-metering").Times(0)
	restore := overrideFactories(t, nil, factoryErr, nil, nil)
	defer restore()

	conf := &config.ServiceConfig{
		Vendor: "aws",
		Bucket: "test-bucket",
		Region: "us-west-2",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			StartTime:    time.Now().Unix(),
			DelayMinutes: 5,
		},
	}

	task, err := NewNextGenTaskWithCheckpointManager(conf, zaptest.NewLogger(t), mockCheckpoint)
	require.NoError(t, err)

	err = task.Run(context.Background())
	require.Error(t, err)
	assert.EqualError(t, err, factoryErr.Error())

}

func TestNextGenTask_GetTaskName(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockCheckpoint := mockcheckpoint.NewMockNextGenCheckpoint(ctrl)
	conf := &config.ServiceConfig{
		Vendor: "aws",
		Bucket: "test-bucket",
		Region: "us-west-2",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			DeltaTablePath: "s3://delta/table",
		},
	}

	task, err := NewNextGenTaskWithCheckpointManager(conf, zaptest.NewLogger(t), mockCheckpoint)
	require.NoError(t, err)
	assert.Equal(t, "nextgen-metering", task.GetTaskName())
}

func TestNextGenTask_NewNextGenTask_MissingConfig(t *testing.T) {
	conf := &config.ServiceConfig{}
	log := zaptest.NewLogger(t)

	task, err := NewNextGenTask(conf, log)
	assert.Error(t, err)
	assert.Nil(t, task)
	assert.Contains(t, err.Error(), "NextGenMeteringConfig is required")
}

func TestNextGenTask_NewNextGenTask_Success(t *testing.T) {
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockCheckpoint := mockcheckpoint.NewMockNextGenCheckpoint(ctrl)
	mockCheckpoint.EXPECT().GetCheckpoint(gomock.Any(), "nextgen-metering").Return(int64(0), nil)
	factory := &fakeCheckpointFactory{manager: mockCheckpoint}
	restore := overrideCheckpointFactory(t, factory)
	defer restore()

	conf := &config.ServiceConfig{
		Vendor: "aws",
		Bucket: "test-bucket",
		Region: "us-west-2",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			DeltaTablePath: "s3://delta/table",
		},
	}

	task, err := NewNextGenTask(conf, zaptest.NewLogger(t))
	require.NoError(t, err)
	assert.NotNil(t, task)
	assert.Equal(t, mockCheckpoint, task.checkpointMgr)
	assert.Equal(t, conf, task.conf)
}

func TestNextGenTask_NewNextGenTask_FactoryError(t *testing.T) {
	factoryErr := errors.New("factory failure")
	factory := &fakeCheckpointFactory{err: factoryErr}
	restore := overrideCheckpointFactory(t, factory)
	defer restore()

	conf := &config.ServiceConfig{
		Vendor: "aws",
		Bucket: "test-bucket",
		Region: "us-west-2",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			DeltaTablePath: "s3://delta/table",
		},
	}

	task, err := NewNextGenTask(conf, zaptest.NewLogger(t))
	require.Error(t, err)
	assert.Nil(t, task)
	assert.Contains(t, err.Error(), "failed to create checkpoint manager")
	assert.ErrorIs(t, err, factoryErr)
}

func TestNextGenTask_NewNextGenTask_Azure(t *testing.T) {
	t.Setenv(constants.AzureStorageAccountNameEnv, "teststorageaccount")
	t.Setenv(constants.AzureStorageAccountKeyEnv, "dGVzdC1hY2NvdW50LWtleQ==")
	t.Setenv(constants.AzureStorageSASTokenEnv, "")

	conf := &config.ServiceConfig{
		Vendor: constants.VendorAzure,
		Bucket: "test-container",
		NextGenMeteringConfig: &config.NextGenMeteringConfig{
			DeltaTablePath:    "azure://test-container/delta-table",
			CheckpointPrefix:  "checkpoints",
			MeteringPrefix:    "metering-data/",
			StartTime:         1700000000,
			DelayMinutes:      5,
			TaskInterval:      1,
			ReporterInterval:  60,
			StorageOptions:    map[string]string{},
			MaintenanceConfig: &config.NextGenMaintenanceConfig{},
		},
	}

	task, err := NewNextGenTask(conf, zaptest.NewLogger(t))

	require.NoError(t, err)
	assert.NotNil(t, task)
	assert.Equal(t, conf, task.conf)
	assert.NotNil(t, task.checkpointMgr)
}