package nextgentasks

import (
	"context"
	"fmt"
	"time"

	"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/nextgenmaintenance"
	"go.uber.org/zap"
)

type maintenanceRunner interface {
	Run(ctx context.Context) error
	GetTableVersion(ctx context.Context) (int64, error)
}

var newNextGenMaintenanceRunner = func(conf *config.ServiceConfig, log *zap.Logger) (maintenanceRunner, error) {
	return nextgenmaintenance.NewNextGenMaintenance(conf, log)
}

// NextGenMaintenanceTask coordinates scheduled maintenance for the NextGen Delta Lake table.
type NextGenMaintenanceTask struct {
	conf          *config.ServiceConfig
	runner        maintenanceRunner
	log           *zap.Logger
	checkpointMgr checkpoint.NextGenCheckpoint
}

// NewNextGenMaintenanceTask creates a new maintenance task instance when maintenance is enabled in config.
func NewNextGenMaintenanceTask(conf *config.ServiceConfig, log *zap.Logger) (*NextGenMaintenanceTask, error) {
	if conf == nil || conf.NextGenMeteringConfig == nil {
		return nil, fmt.Errorf("NextGenMeteringConfig is required")
	}

	maintCfg := conf.NextGenMeteringConfig.MaintenanceConfig
	if maintCfg == nil {
		return nil, fmt.Errorf("NextGen maintenance config is required")
	}
	if !maintCfg.Enable {
		return nil, fmt.Errorf("NextGen maintenance is disabled in configuration")
	}

	runner, err := newNextGenMaintenanceRunner(conf, log)
	if err != nil {
		return nil, fmt.Errorf("failed to create NextGen maintenance runner: %w", err)
	}

	// Create checkpoint manager for updating metering checkpoint version
	factory := newCheckpointFactory(conf, log)
	checkpointMgr, err := factory.CreateCheckpointManager()
	if err != nil {
		return nil, fmt.Errorf("failed to create checkpoint manager: %w", err)
	}

	return &NextGenMaintenanceTask{
		conf:          conf,
		runner:        runner,
		log:           log,
		checkpointMgr: checkpointMgr,
	}, nil
}

// Run executes the maintenance workflow with proper Delta Lake locking.
func (t *NextGenMaintenanceTask) Run(ctx context.Context) (err error) {
	lock := GetGlobalDeltaLakeLock()
	lock.Lock(t.GetTaskName())
	defer lock.Unlock(t.GetTaskName())
	runStart := time.Now()
	defer func() {
		recordNextGenTaskDurationMetric(t.GetTaskName(), t.conf.Vendor, "maintenance", time.Since(runStart))
		if err != nil {
			recordNextGenTaskResultMetric(t.GetTaskName(), t.conf.Vendor, constants.MetricsResultFailed)
		} else {
			recordNextGenTaskResultMetric(t.GetTaskName(), t.conf.Vendor, constants.MetricsResultSuccess)
		}
	}()

	// Get nextgen-metering checkpoint data to track version changes
	meteringTaskName := "nextgen-metering"
	checkpointData, err := t.checkpointMgr.GetCheckpointData(ctx, meteringTaskName)
	if err != nil {
		t.log.Warn("failed to get metering checkpoint data, will skip version update",
			zap.Error(err))
		checkpointData = nil
	}

	// Get current table version before maintenance
	var versionBeforeMaintenance int64
	var shouldUpdateVersion bool
	if checkpointData != nil {
		versionBeforeMaintenance, err = t.runner.GetTableVersion(ctx)
		if err != nil {
			t.log.Warn("failed to get table version before maintenance",
				zap.Error(err))
		} else {
			// Only update version if checkpoint version matches current table version
			shouldUpdateVersion = checkpointData.DeltaTableVersion == versionBeforeMaintenance
			t.log.Info("checked table version before maintenance",
				zap.Int64("checkpoint_version", checkpointData.DeltaTableVersion),
				zap.Int64("current_version", versionBeforeMaintenance),
				zap.Bool("should_update_version", shouldUpdateVersion))
		}
	}

	t.log.Info("starting NextGen maintenance task")
	err = t.runner.Run(ctx)
	if err != nil {
		t.log.Error("NextGen maintenance task failed", zap.Error(err))
		return err
	}

	// After successful maintenance, check if version increased and update checkpoint if needed
	if shouldUpdateVersion && checkpointData != nil {
		versionAfterMaintenance, versionErr := t.runner.GetTableVersion(ctx)
		if versionErr != nil {
			t.log.Warn("failed to get table version after maintenance",
				zap.Error(versionErr))
		} else if versionAfterMaintenance > versionBeforeMaintenance {
			// Version increased, update checkpoint with new version
			checkpointData.DeltaTableVersion = versionAfterMaintenance
			checkpointData.UpdatedAt = time.Now().Unix()

			if saveErr := t.checkpointMgr.SaveCheckpointData(ctx, checkpointData); saveErr != nil {
				t.log.Error("failed to update checkpoint with new table version",
					zap.Int64("old_version", versionBeforeMaintenance),
					zap.Int64("new_version", versionAfterMaintenance),
					zap.Error(saveErr))
			} else {
				t.log.Info("updated metering checkpoint with new table version",
					zap.Int64("old_version", versionBeforeMaintenance),
					zap.Int64("new_version", versionAfterMaintenance))
			}
		} else {
			t.log.Debug("table version unchanged after maintenance",
				zap.Int64("version", versionAfterMaintenance))
		}
	}

	t.log.Info("NextGen maintenance task completed successfully")
	return nil
}

// GetTaskName returns the identifier used for locking and scheduling.
func (t *NextGenMaintenanceTask) GetTaskName() string {
	return "nextgen_maintenance"
}
