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

type meteringService interface {
	checkpoint.MeteringServiceInterface
	Close() error
	TableExists(ctx context.Context) (bool, error)
	CreateDeltaTable() error
	ProcessMinute(ctx context.Context, timestamp int64) ([]models.BillingReportV2Deltalake, error)
	WriteToDeltaTable(records []models.BillingReportV2Deltalake, saveMode string) error
}

type deltaLakeHelper interface {
	ShouldRecoverFromTable(ctx context.Context, meteringService checkpoint.MeteringServiceInterface, checkpointData *checkpoint.CheckpointData) (bool, int64, error)
	RecoverCheckpointFromTable(ctx context.Context, meteringService checkpoint.MeteringServiceInterface, currentTableVersion int64, checkpointData *checkpoint.CheckpointData) (*checkpoint.CheckpointData, error)
}

type checkpointFactory interface {
	CreateCheckpointManager() (checkpoint.NextGenCheckpoint, error)
}

var (
	newNextGenMeteringService = func(conf *config.ServiceConfig, log *zap.Logger) (meteringService, error) {
		return nextgenmetering.NewNextGenMetering(conf, log)
	}
	newDeltaLakeCheckpointHelper = func(conf *config.ServiceConfig, log *zap.Logger) deltaLakeHelper {
		return checkpoint.NewDeltaLakeCheckpointHelper(conf, log)
	}
	newCheckpointFactory = func(conf *config.ServiceConfig, log *zap.Logger) checkpointFactory {
		return checkpoint.NewFactory(conf, log)
	}
	timeNow = time.Now
)

// NextGenTask is a task for next generation metering
// It will run the metering and write to delta lake
type NextGenTask struct {
	conf          *config.ServiceConfig
	log           *zap.Logger
	checkpointMgr checkpoint.NextGenCheckpoint
}

// NewNextGenTask creates a new NextGenTask instance
func NewNextGenTask(conf *config.ServiceConfig, log *zap.Logger) (*NextGenTask, error) {
	if conf.NextGenMeteringConfig == nil {
		return nil, fmt.Errorf("NextGenMeteringConfig is required")
	}

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

	// Load checkpoint and set metric at startup to avoid metric gaps
	initCheckpoint, err := checkpointMgr.GetCheckpoint(context.Background(), "nextgen-metering")
	if err != nil {
		log.Warn("failed to load initial checkpoint for metric", zap.Error(err))
	} else if initCheckpoint > 0 {
		metrics.MetricsCheckpointTimestamp.WithLabelValues("nextgen-metering").Set(float64(initCheckpoint))
		log.Info("loaded initial checkpoint metric", zap.Int64("checkpoint", initCheckpoint))
	}

	return &NextGenTask{
		conf:          conf,
		log:           log,
		checkpointMgr: checkpointMgr,
	}, nil
}

// NewNextGenTaskWithCheckpointManager creates a new NextGenTask instance with a custom checkpoint manager (for testing)
func NewNextGenTaskWithCheckpointManager(conf *config.ServiceConfig, log *zap.Logger, checkpointMgr checkpoint.NextGenCheckpoint) (*NextGenTask, error) {
	if conf.NextGenMeteringConfig == nil {
		return nil, fmt.Errorf("NextGenMeteringConfig is required")
	}

	return &NextGenTask{
		conf:          conf,
		log:           log,
		checkpointMgr: checkpointMgr,
	}, nil
}

// Run executes the next generation metering task for a time range from start to end
func (n *NextGenTask) Run(ctx context.Context) (err error) {
	// Acquire Delta Lake lock to prevent concurrent access with reporter task
	lock := GetGlobalDeltaLakeLock()
	lock.Lock(n.GetTaskName())
	defer lock.Unlock(n.GetTaskName())

	runStart := timeNow()
	defer func() {
		n.recordTaskDuration("metering", timeNow().Sub(runStart))
		if err != nil {
			n.recordTaskResult(constants.MetricsResultFailed)
		} else {
			n.recordTaskResult(constants.MetricsResultSuccess)
		}
	}()

	// Create Delta Lake checkpoint helper with config
	deltaHelper := newDeltaLakeCheckpointHelper(n.conf, n.log)

	// Create NextGen metering instance early for checkpoint optimization
	meteringService, err := newNextGenMeteringService(n.conf, n.log)
	if err != nil {
		n.log.Error("failed to create NextGen metering service", zap.Error(err))
		return err
	}

	// Ensure proper cleanup of metering service
	defer func() {
		if closeErr := meteringService.Close(); closeErr != nil {
			n.log.Error("failed to close metering service", zap.Error(closeErr))
		}
	}()

	// Get full checkpoint data to determine start time and check for recovery
	checkpointData, err := n.checkpointMgr.GetCheckpointData(ctx, n.GetTaskName())
	if err != nil {
		n.log.Error("failed to get checkpoint data", zap.Error(err))
		return err
	}

	// Set the Delta table path if not already set in checkpoint
	if checkpointData.DeltaTablePath == "" && n.conf.NextGenMeteringConfig != nil {
		checkpointData.DeltaTablePath = n.conf.NextGenMeteringConfig.DeltaTablePath
	}

	// Check if we need to create the delta table before checkpoint recovery
	// Only check and create table if no checkpoint exists
	if checkpointData.Checkpoint == 0 && checkpointData.DeltaTablePath != "" {
		n.log.Info("No checkpoint exists, checking if Delta table needs to be created")

		exists, err := meteringService.TableExists(ctx)
		if err != nil {
			n.log.Error("failed to check table existence", zap.Error(err))
			return err
		}

		if !exists {
			n.log.Info("Delta table does not exist, creating new table")
			if err := meteringService.CreateDeltaTable(); err != nil {
				n.log.Error("failed to create delta table", zap.Error(err))
				return err
			}
		} else {
			n.log.Info("Delta table already exists, no need to create")

			// Fail fast if the path looks populated but cannot actually be opened as a Delta table.
			if _, err := meteringService.GetTableVersion(ctx); err != nil {
				return fmt.Errorf("delta table exists but could not be opened during initialization at %s: %w", checkpointData.DeltaTablePath, err)
			}
		}
	}

	// Check if we need to recover from Delta Lake table
	if checkpointData.DeltaTablePath != "" {
		shouldRecover, currentTableVersion, err := deltaHelper.ShouldRecoverFromTable(ctx, meteringService, checkpointData)
		if err != nil {
			n.log.Warn("failed to check table version for recovery",
				zap.Error(err),
				zap.String("table_path", checkpointData.DeltaTablePath))
			// Continue with normal processing if table check fails
		} else if shouldRecover {
			n.log.Info("Delta Lake table version is higher than checkpoint, performing recovery",
				zap.Int64("checkpoint_version", checkpointData.DeltaTableVersion),
				zap.Int64("current_version", currentTableVersion))

			recoveredData, err := deltaHelper.RecoverCheckpointFromTable(ctx, meteringService, currentTableVersion, checkpointData)
			if err != nil {
				n.log.Error("failed to recover checkpoint from table", zap.Error(err))
				// Continue with original checkpoint data if recovery fails
			} else {
				n.log.Info("successfully recovered checkpoint from Delta Lake table",
					zap.Int64("original_checkpoint", checkpointData.Checkpoint),
					zap.Int64("recovered_checkpoint", recoveredData.Checkpoint),
					zap.Int64("new_table_version", recoveredData.DeltaTableVersion))
				// check if recovered checkpoint is zero
				if recoveredData.Checkpoint == 0 {
					n.log.Error("recovered checkpoint is zero, which may indicate no valid data in table",
						zap.Int64("original_checkpoint", checkpointData.Checkpoint))
					return fmt.Errorf("recovered checkpoint is zero, cannot proceed with metering")
				}
				// check if recovered checkpoint is ahead of original
				if recoveredData.Checkpoint <= checkpointData.Checkpoint {
					n.log.Warn("recovered checkpoint is not ahead of original, mean's no progress metering, we can recover from here safety",
						zap.Int64("original_checkpoint", checkpointData.Checkpoint),
						zap.Int64("recovered_checkpoint", recoveredData.Checkpoint))
				}
				// Use recovered data and save it as the new checkpoint
				checkpointData = recoveredData
				if saveErr := n.checkpointMgr.SaveCheckpointData(ctx, checkpointData); saveErr != nil {
					n.log.Error("failed to save recovered checkpoint", zap.Error(saveErr))
					// Continue anyway, recovery was successful
				}
			}
		}
	}

	// Determine start time
	var startTime int64
	if checkpointData.Checkpoint == 0 {
		// No checkpoint exists, use configured start time
		if n.conf.NextGenMeteringConfig.StartTime > 0 {
			startTime = n.conf.NextGenMeteringConfig.StartTime
		} else {
			panic("NextGenMeteringConfig.StartTime must be set if no checkpoint exists")
		}
		n.log.Info("no checkpoint found, using start time", zap.Int64("start_time", startTime))
	} else {
		// Continue from checkpoint (next minute after last processed)
		startTime = checkpointData.Checkpoint + 60
		n.log.Info("found checkpoint, continuing from",
			zap.Int64("checkpoint", checkpointData.Checkpoint),
			zap.Int64("start_time", startTime),
			zap.String("delta_table_path", checkpointData.DeltaTablePath),
			zap.Int64("delta_table_version", checkpointData.DeltaTableVersion))
	}

	// Calculate end time (current time minus delay)
	delayMinutes := n.conf.NextGenMeteringConfig.DelayMinutes
	if delayMinutes == 0 {
		delayMinutes = 5 // Default 5 minutes delay
	}

	// Round down to minute boundary and subtract delay
	now := timeNow().UTC()
	roundedTime := time.Date(now.Year(), now.Month(), now.Day(), now.Hour(), now.Minute(), 0, 0, time.UTC)
	endTime := roundedTime.Add(-time.Duration(delayMinutes) * time.Minute).Unix()

	n.log.Info("calculated time range",
		zap.Int64("start_time", startTime),
		zap.Int64("end_time", endTime),
		zap.Int("delay_minutes", delayMinutes))

	// Check if there's any work to do
	if startTime >= endTime {
		n.log.Info("no new data to process",
			zap.Int64("start_time", startTime),
			zap.Int64("end_time", endTime))
		return nil
	}

	// Collect all SKU records for the entire time range (up to 15 minutes)
	var allSkuRecords []models.BillingReportV2Deltalake
	var lastProcessedTime int64

	// Process each minute from start to end
	for currentTime := startTime; currentTime < endTime; currentTime += 60 {
		n.log.Debug("processing timestamp", zap.Int64("timestamp", currentTime))

		// Process minute and collect SKU records
		skuRecords, err := meteringService.ProcessMinute(ctx, currentTime)
		if err != nil {
			n.log.Error("failed to process minute",
				zap.Int64("timestamp", currentTime),
				zap.Error(err))
			return err
		}

		// Append to the batch
		allSkuRecords = append(allSkuRecords, skuRecords...)
		lastProcessedTime = currentTime

		n.log.Debug("collected SKU records",
			zap.Int64("timestamp", currentTime),
			zap.Int("minute_records", len(skuRecords)),
			zap.Int("total_records", len(allSkuRecords)))
	}

	// Write all collected records to Delta Lake in one batch
	if len(allSkuRecords) > 0 {
		n.log.Info("writing batch to Delta Lake",
			zap.Int("total_records", len(allSkuRecords)),
			zap.Int64("start_time", startTime),
			zap.Int64("end_time", endTime))

		err := meteringService.WriteToDeltaTable(allSkuRecords, "Append")
		if err != nil {
			n.log.Error("failed to write batch to delta table",
				zap.Int("records", len(allSkuRecords)),
				zap.Error(err))
			return err
		}

		n.log.Info("successfully wrote batch to Delta Lake",
			zap.Int("records", len(allSkuRecords)))
	} else {
		n.log.Info("no records to write for this time range",
			zap.Int64("start_time", startTime),
			zap.Int64("end_time", endTime))
	}

	// Get current table version for checkpoint after batch write
	tableVersion, versionErr := meteringService.GetTableVersion(ctx)
	if versionErr != nil {
		n.log.Warn("failed to get table version for checkpoint, using 0", zap.Error(versionErr))
		tableVersion = 0
	}

	// Save checkpoint after successful batch processing
	if lastProcessedTime > 0 {
		checkpointData := &checkpoint.CheckpointData{
			TaskName:          n.GetTaskName(),
			Checkpoint:        lastProcessedTime,
			UpdatedAt:         timeNow().Unix(),
			DeltaTablePath:    n.conf.NextGenMeteringConfig.DeltaTablePath,
			DeltaTableVersion: tableVersion,
		}

		err = n.checkpointMgr.SaveCheckpointData(ctx, checkpointData)
		if err != nil {
			n.log.Error("failed to save checkpoint data",
				zap.Int64("last_processed", lastProcessedTime),
				zap.Int64("table_version", tableVersion),
				zap.Error(err))
			return err
		}

		metrics.MetricsCheckpointTimestamp.WithLabelValues(n.GetTaskName()).Set(float64(lastProcessedTime))

		n.log.Info("next generation metering task completed successfully",
			zap.Int64("start_time", startTime),
			zap.Int64("end_time", endTime),
			zap.Int64("last_processed", lastProcessedTime),
			zap.Int("total_records", len(allSkuRecords)),
			zap.Int64("table_version", tableVersion))
	}

	return nil
}

// recordTaskResult increments the task count metric for the given result type.
func (n *NextGenTask) recordTaskResult(result string) {
	recordNextGenTaskResultMetric(n.GetTaskName(), n.conf.Vendor, result)
}

// recordTaskDuration records the duration for a specific operation within the task lifecycle.
func (n *NextGenTask) recordTaskDuration(operation string, duration time.Duration) {
	recordNextGenTaskDurationMetric(n.GetTaskName(), n.conf.Vendor, operation, duration)
}

// GetTaskName returns the task name for checkpoint management
func (n *NextGenTask) GetTaskName() string {
	return "nextgen-metering"
}