package nextgentasks

import (
	"context"
	"errors"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/metering/pkg/common/checkpoint"
	"github.com/tidbcloud/metering/pkg/common/config"
	"go.uber.org/zap"
	"go.uber.org/zap/zaptest"
)

type fakeMaintenanceRunner struct {
	calls               int
	ctx                 context.Context
	err                 error
	tableVersion        int64
	tableVersionErr     error
	tableVersionChanged bool // If true, version increases after Run
}

func (f *fakeMaintenanceRunner) Run(ctx context.Context) error {
	f.calls++
	f.ctx = ctx
	if f.tableVersionChanged {
		f.tableVersion++ // Simulate version increase after maintenance
	}
	if f.err != nil {
		return f.err
	}
	return nil
}

func (f *fakeMaintenanceRunner) GetTableVersion(ctx context.Context) (int64, error) {
	if f.tableVersionErr != nil {
		return 0, f.tableVersionErr
	}
	return f.tableVersion, nil
}

type fakeCheckpointManager struct {
	checkpointData    *checkpoint.CheckpointData
	getErr            error
	saveErr           error
	savedData         *checkpoint.CheckpointData
	getCheckpointData func(ctx context.Context, taskName string) (*checkpoint.CheckpointData, error)
}

func (f *fakeCheckpointManager) GetCheckpoint(ctx context.Context, taskName string) (int64, error) {
	if f.checkpointData != nil {
		return f.checkpointData.Checkpoint, nil
	}
	return 0, nil
}

func (f *fakeCheckpointManager) SaveCheckpoint(ctx context.Context, taskName string, cp int64) error {
	return f.saveErr
}

func (f *fakeCheckpointManager) GetCheckpointData(ctx context.Context, taskName string) (*checkpoint.CheckpointData, error) {
	if f.getCheckpointData != nil {
		return f.getCheckpointData(ctx, taskName)
	}
	if f.getErr != nil {
		return nil, f.getErr
	}
	if f.checkpointData != nil {
		return f.checkpointData, nil
	}
	return &checkpoint.CheckpointData{TaskName: taskName}, nil
}

func (f *fakeCheckpointManager) SaveCheckpointData(ctx context.Context, data *checkpoint.CheckpointData) error {
	f.savedData = data
	return f.saveErr
}

func overrideMaintenanceRunnerFactory(t *testing.T, runner maintenanceRunner, err error) func() {
	t.Helper()
	original := newNextGenMaintenanceRunner
	newNextGenMaintenanceRunner = func(*config.ServiceConfig, *zap.Logger) (maintenanceRunner, error) {
		if err != nil {
			return nil, err
		}
		return runner, nil
	}
	return func() {
		newNextGenMaintenanceRunner = original
	}
}

func overrideCheckpointFactoryForMaintenance(t *testing.T, checkpointMgr checkpoint.NextGenCheckpoint) func() {
	t.Helper()
	original := newCheckpointFactory
	newCheckpointFactory = func(*config.ServiceConfig, *zap.Logger) checkpointFactory {
		return &fakeCheckpointFactoryForMaintenance{checkpointMgr: checkpointMgr}
	}
	return func() {
		newCheckpointFactory = original
	}
}

type fakeCheckpointFactoryForMaintenance struct {
	checkpointMgr checkpoint.NextGenCheckpoint
}

func (f *fakeCheckpointFactoryForMaintenance) CreateCheckpointManager() (checkpoint.NextGenCheckpoint, error) {
	return f.checkpointMgr, nil
}

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

	task, err := NewNextGenMaintenanceTask(nil, logger)
	assert.Error(t, err)
	assert.Nil(t, task)

	task, err = NewNextGenMaintenanceTask(&config.ServiceConfig{}, logger)
	assert.Error(t, err)
	assert.Nil(t, task)

	task, err = NewNextGenMaintenanceTask(&config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{}}, logger)
	assert.Error(t, err)
	assert.Nil(t, task)

	task, err = NewNextGenMaintenanceTask(&config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{MaintenanceConfig: &config.NextGenMaintenanceConfig{}}}, logger)
	assert.Error(t, err)
	assert.Nil(t, task)
}

func TestNewNextGenMaintenanceTask_FactoryError(t *testing.T) {
	logger := zaptest.NewLogger(t)
	factoryErr := errors.New("factory failed")
	restore := overrideMaintenanceRunnerFactory(t, nil, factoryErr)
	defer restore()

	conf := &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{
		MaintenanceConfig: &config.NextGenMaintenanceConfig{Enable: true},
	}}

	task, err := NewNextGenMaintenanceTask(conf, logger)
	require.Error(t, err)
	assert.Nil(t, task)
	assert.ErrorIs(t, err, factoryErr)
}

func TestNewNextGenMaintenanceTask_Success(t *testing.T) {
	logger := zaptest.NewLogger(t)
	fake := &fakeMaintenanceRunner{}
	restore := overrideMaintenanceRunnerFactory(t, fake, nil)
	defer restore()
	restoreCheckpoint := overrideCheckpointFactoryForMaintenance(t, &fakeCheckpointManager{})
	defer restoreCheckpoint()

	conf := &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{
		MaintenanceConfig: &config.NextGenMaintenanceConfig{Enable: true},
	}}

	task, err := NewNextGenMaintenanceTask(conf, logger)
	require.NoError(t, err)
	assert.NotNil(t, task)
	assert.Equal(t, conf, task.conf)
}

func TestNextGenMaintenanceTask_RunSuccess(t *testing.T) {
	logger := zaptest.NewLogger(t)
	fake := &fakeMaintenanceRunner{tableVersion: 5}
	fakeCheckpoint := &fakeCheckpointManager{
		checkpointData: &checkpoint.CheckpointData{
			TaskName:          "nextgen-metering",
			Checkpoint:        1000,
			DeltaTableVersion: 5,
		},
	}
	restore := overrideMaintenanceRunnerFactory(t, fake, nil)
	defer restore()
	restoreCheckpoint := overrideCheckpointFactoryForMaintenance(t, fakeCheckpoint)
	defer restoreCheckpoint()

	conf := &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{
		MaintenanceConfig: &config.NextGenMaintenanceConfig{Enable: true},
	}}

	task, err := NewNextGenMaintenanceTask(conf, logger)
	require.NoError(t, err)

	resetGlobalDeltaLakeLock()
	InitGlobalDeltaLakeLock(logger)

	ctx := context.Background()
	err = task.Run(ctx)
	require.NoError(t, err)
	assert.Equal(t, 1, fake.calls)
	assert.Equal(t, ctx, fake.ctx)
}

func TestNextGenMaintenanceTask_RunError(t *testing.T) {
	logger := zaptest.NewLogger(t)
	runner := &fakeMaintenanceRunner{err: errors.New("boom")}
	fakeCheckpoint := &fakeCheckpointManager{}
	restore := overrideMaintenanceRunnerFactory(t, runner, nil)
	defer restore()
	restoreCheckpoint := overrideCheckpointFactoryForMaintenance(t, fakeCheckpoint)
	defer restoreCheckpoint()

	conf := &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{
		MaintenanceConfig: &config.NextGenMaintenanceConfig{Enable: true},
	}}

	task, err := NewNextGenMaintenanceTask(conf, logger)
	require.NoError(t, err)

	resetGlobalDeltaLakeLock()
	InitGlobalDeltaLakeLock(logger)

	err = task.Run(context.Background())
	require.Error(t, err)
	assert.Equal(t, 1, runner.calls)
}

func TestNextGenMaintenanceTask_RunUpdatesVersionOnChange(t *testing.T) {
	logger := zaptest.NewLogger(t)
	// Version starts at 5, will increment to 6 after Run
	fake := &fakeMaintenanceRunner{tableVersion: 5, tableVersionChanged: true}
	fakeCheckpoint := &fakeCheckpointManager{
		checkpointData: &checkpoint.CheckpointData{
			TaskName:          "nextgen-metering",
			Checkpoint:        1000,
			DeltaTableVersion: 5, // Matches current version
		},
	}
	restore := overrideMaintenanceRunnerFactory(t, fake, nil)
	defer restore()
	restoreCheckpoint := overrideCheckpointFactoryForMaintenance(t, fakeCheckpoint)
	defer restoreCheckpoint()

	conf := &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{
		MaintenanceConfig: &config.NextGenMaintenanceConfig{Enable: true},
	}}

	task, err := NewNextGenMaintenanceTask(conf, logger)
	require.NoError(t, err)

	resetGlobalDeltaLakeLock()
	InitGlobalDeltaLakeLock(logger)

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

	// Verify checkpoint was saved with new version
	require.NotNil(t, fakeCheckpoint.savedData)
	assert.Equal(t, int64(6), fakeCheckpoint.savedData.DeltaTableVersion)
	assert.Equal(t, int64(1000), fakeCheckpoint.savedData.Checkpoint) // Checkpoint timestamp unchanged
}

func TestNextGenMaintenanceTask_RunNoUpdateWhenVersionMismatch(t *testing.T) {
	logger := zaptest.NewLogger(t)
	fake := &fakeMaintenanceRunner{tableVersion: 10, tableVersionChanged: true}
	fakeCheckpoint := &fakeCheckpointManager{
		checkpointData: &checkpoint.CheckpointData{
			TaskName:          "nextgen-metering",
			Checkpoint:        1000,
			DeltaTableVersion: 5, // Does NOT match current version (10)
		},
	}
	restore := overrideMaintenanceRunnerFactory(t, fake, nil)
	defer restore()
	restoreCheckpoint := overrideCheckpointFactoryForMaintenance(t, fakeCheckpoint)
	defer restoreCheckpoint()

	conf := &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{
		MaintenanceConfig: &config.NextGenMaintenanceConfig{Enable: true},
	}}

	task, err := NewNextGenMaintenanceTask(conf, logger)
	require.NoError(t, err)

	resetGlobalDeltaLakeLock()
	InitGlobalDeltaLakeLock(logger)

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

	// Verify checkpoint was NOT saved because versions didn't match
	assert.Nil(t, fakeCheckpoint.savedData)
}

func TestNextGenMaintenanceTask_RunNoUpdateWhenVersionUnchanged(t *testing.T) {
	logger := zaptest.NewLogger(t)
	fake := &fakeMaintenanceRunner{tableVersion: 5, tableVersionChanged: false} // Version stays at 5
	fakeCheckpoint := &fakeCheckpointManager{
		checkpointData: &checkpoint.CheckpointData{
			TaskName:          "nextgen-metering",
			Checkpoint:        1000,
			DeltaTableVersion: 5,
		},
	}
	restore := overrideMaintenanceRunnerFactory(t, fake, nil)
	defer restore()
	restoreCheckpoint := overrideCheckpointFactoryForMaintenance(t, fakeCheckpoint)
	defer restoreCheckpoint()

	conf := &config.ServiceConfig{NextGenMeteringConfig: &config.NextGenMeteringConfig{
		MaintenanceConfig: &config.NextGenMaintenanceConfig{Enable: true},
	}}

	task, err := NewNextGenMaintenanceTask(conf, logger)
	require.NoError(t, err)

	resetGlobalDeltaLakeLock()
	InitGlobalDeltaLakeLock(logger)

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

	// Verify checkpoint was NOT saved because version didn't change
	assert.Nil(t, fakeCheckpoint.savedData)
}