package checkpoint

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/tidbcloud/metering/pkg/common/config"
	"github.com/tidbcloud/metering/pkg/common/constants"
	"github.com/tidbcloud/metering/pkg/common/utils/duckdb"
	"go.uber.org/zap"
)

// MeteringServiceInterface defines the interface for NextGen metering operations
type MeteringServiceInterface interface {
	GetTableVersion(ctx context.Context) (int64, error)
	QueryMaxTimestampWithPartitionFilter(ctx context.Context, partitionTime int64) (int64, error)
}

// DuckDBClientInterface defines the interface for DuckDB operations
type DuckDBClientInterface interface {
	QueryMaxTimestampWithPartitionFilter(tablePath string, partitionTime int64) (int64, error)
	Close() error
}

// DuckDBClientFactory defines how to create DuckDB clients
type DuckDBClientFactory interface {
	NewDuckDBClientWithConfig(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error)
}

// DefaultDuckDBClientFactory implements DuckDBClientFactory using the real duckdb package
type DefaultDuckDBClientFactory struct{}

func (f *DefaultDuckDBClientFactory) NewDuckDBClientWithConfig(conf *config.ServiceConfig, log *zap.Logger) (DuckDBClientInterface, error) {
	return duckdb.NewDuckDBClientWithConfig(conf, log)
}

// DeltaLakeCheckpointHelper provides utility functions for Delta Lake checkpoint operations
type DeltaLakeCheckpointHelper struct {
	log           *zap.Logger
	conf          *config.ServiceConfig
	duckDBFactory DuckDBClientFactory
}

// NewDeltaLakeCheckpointHelper creates a new helper for Delta Lake checkpoint operations
func NewDeltaLakeCheckpointHelper(conf *config.ServiceConfig, log *zap.Logger) *DeltaLakeCheckpointHelper {
	return &DeltaLakeCheckpointHelper{
		log:           log,
		conf:          conf,
		duckDBFactory: &DefaultDuckDBClientFactory{},
	}
}

// NewDeltaLakeCheckpointHelperWithFactory creates a new helper with custom DuckDB factory for testing
func NewDeltaLakeCheckpointHelperWithFactory(conf *config.ServiceConfig, log *zap.Logger, factory DuckDBClientFactory) *DeltaLakeCheckpointHelper {
	return &DeltaLakeCheckpointHelper{
		log:           log,
		conf:          conf,
		duckDBFactory: factory,
	}
}

// GetTableVersion retrieves the current version of a Delta Lake table using provided metering service
func (h *DeltaLakeCheckpointHelper) GetTableVersion(ctx context.Context, meteringService MeteringServiceInterface) (int64, error) {
	h.log.Debug("getting Delta Lake table version using provided metering service")

	// Get the current version using the provided service (will reuse opened table if available)
	version, err := meteringService.GetTableVersion(ctx)
	if err != nil {
		h.log.Error("failed to get table version", zap.Error(err))
		return 0, fmt.Errorf("failed to get table version: %w", err)
	}

	h.log.Debug("successfully retrieved table version", zap.Int64("version", version))
	return version, nil
}

// GetLatestCostTimestamp retrieves the maximum cost_time_stamp from a Delta Lake table using DuckDB
func (h *DeltaLakeCheckpointHelper) GetLatestCostTimestamp(ctx context.Context, meteringService MeteringServiceInterface, tablePath string, partitiontime int64) (int64, error) {
	if h.conf != nil && strings.EqualFold(h.conf.Vendor, constants.VendorAlicloud) {
		if meteringService == nil {
			h.log.Error("metering service is required for Alicloud max timestamp query")
			return 0, fmt.Errorf("metering service is required for Alicloud vendor")
		}

		h.log.Debug("getting latest cost_time_stamp from Delta Lake table using delta-rs",
			zap.String("table_path", tablePath))

		maxTimestamp, err := meteringService.QueryMaxTimestampWithPartitionFilter(ctx, partitiontime)
		if err != nil {
			h.log.Error("failed to query max timestamp using delta-rs",
				zap.String("table_path", tablePath),
				zap.Error(err))
			return 0, fmt.Errorf("failed to query max timestamp: %w", err)
		}

		h.log.Info("successfully retrieved latest cost_time_stamp using delta-rs",
			zap.String("table_path", tablePath),
			zap.Int64("max_timestamp", maxTimestamp))

		return maxTimestamp, nil
	}

	h.log.Debug("getting latest cost_time_stamp from Delta Lake table using DuckDB", zap.String("table_path", tablePath))

	// Create DuckDB client with configuration
	duckClient, err := h.duckDBFactory.NewDuckDBClientWithConfig(h.conf, h.log)
	if err != nil {
		h.log.Error("failed to create DuckDB client", zap.Error(err))
		return 0, fmt.Errorf("failed to create DuckDB client: %w", err)
	}
	defer duckClient.Close()

	// Use DuckDB to efficiently query the maximum cost_time_stamp
	maxTimestamp, err := duckClient.QueryMaxTimestampWithPartitionFilter(tablePath, partitiontime)
	if err != nil {
		h.log.Error("failed to query max timestamp using DuckDB",
			zap.String("table_path", tablePath),
			zap.Error(err))
		return 0, fmt.Errorf("failed to query max timestamp: %w", err)
	}

	h.log.Info("successfully retrieved latest cost_time_stamp using DuckDB",
		zap.String("table_path", tablePath),
		zap.Int64("max_timestamp", maxTimestamp))

	return maxTimestamp, nil
}

// ShouldRecoverFromTable determines if recovery should be performed from the Delta Lake table
// Returns true if the table version is higher than the checkpoint version
func (h *DeltaLakeCheckpointHelper) ShouldRecoverFromTable(ctx context.Context, meteringService MeteringServiceInterface, checkpointData *CheckpointData) (bool, int64, error) {
	if checkpointData.DeltaTablePath == "" {
		h.log.Debug("no Delta table path in checkpoint, no recovery needed")
		return false, 0, nil
	}

	currentVersion, err := h.GetTableVersion(ctx, meteringService)
	if err != nil {
		return false, 0, fmt.Errorf("failed to get current table version: %w", err)
	}

	shouldRecover := currentVersion > checkpointData.DeltaTableVersion

	h.log.Info("checking if recovery is needed",
		zap.String("table_path", checkpointData.DeltaTablePath),
		zap.Int64("checkpoint_version", checkpointData.DeltaTableVersion),
		zap.Int64("current_version", currentVersion),
		zap.Bool("should_recover", shouldRecover))

	return shouldRecover, currentVersion, nil
}

// RecoverCheckpointFromTable performs intelligent checkpoint recovery from Delta Lake table
func (h *DeltaLakeCheckpointHelper) RecoverCheckpointFromTable(ctx context.Context, meteringService MeteringServiceInterface, currentTableVersion int64, checkpointData *CheckpointData) (*CheckpointData, error) {
	h.log.Info("starting checkpoint recovery from Delta Lake table",
		zap.String("table_path", checkpointData.DeltaTablePath),
		zap.Int64("checkpoint_version", checkpointData.DeltaTableVersion))

	// Get the latest cost_time_stamp from the table
	latestTimestamp, err := h.GetLatestCostTimestamp(ctx, meteringService, checkpointData.DeltaTablePath, checkpointData.Checkpoint)
	if err != nil {
		return nil, fmt.Errorf("failed to get latest timestamp from table: %w", err)
	}
	// Create recovered checkpoint data
	recoveredData := &CheckpointData{
		TaskName:          checkpointData.TaskName,
		Checkpoint:        latestTimestamp,
		UpdatedAt:         time.Now().Unix(),
		DeltaTablePath:    checkpointData.DeltaTablePath,
		DeltaTableVersion: currentTableVersion,
	}

	h.log.Info("successfully recovered checkpoint from table",
		zap.String("table_path", checkpointData.DeltaTablePath),
		zap.Int64("recovered_timestamp", latestTimestamp),
		zap.Int64("new_table_version", currentTableVersion),
		zap.Int64("original_checkpoint", checkpointData.Checkpoint))

	return recoveredData, nil
}