package checkpoint

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"time"

	sdkstorage "github.com/pingcap/metering_sdk/storage"
	"go.uber.org/zap"
	"gopkg.in/yaml.v3"
)

// ObjectStorageCheckpointManager implements NextGenCheckpoint using a generic object storage provider.
type ObjectStorageCheckpointManager struct {
	provider sdkstorage.ObjectStorageProvider
	log      *zap.Logger
}

// NewObjectStorageCheckpointManager creates a checkpoint manager backed by a generic object storage provider.
func NewObjectStorageCheckpointManager(provider sdkstorage.ObjectStorageProvider, log *zap.Logger) NextGenCheckpoint {
	return &ObjectStorageCheckpointManager{
		provider: provider,
		log:      log,
	}
}

func checkpointObjectPath(taskName string) string {
	return fmt.Sprintf("%s/%s", taskName, checkpointFileName)
}

// GetCheckpoint retrieves the checkpoint for a specific task.
func (m *ObjectStorageCheckpointManager) GetCheckpoint(ctx context.Context, taskName string) (int64, error) {
	data, err := m.GetCheckpointData(ctx, taskName)
	if err != nil {
		return 0, err
	}
	return data.Checkpoint, nil
}

// SaveCheckpoint saves the checkpoint for a specific task.
func (m *ObjectStorageCheckpointManager) SaveCheckpoint(ctx context.Context, taskName string, checkpoint int64) error {
	return m.SaveCheckpointData(ctx, &CheckpointData{
		Checkpoint: checkpoint,
		UpdatedAt:  time.Now().Unix(),
		TaskName:   taskName,
	})
}

// GetCheckpointData retrieves full checkpoint data including Delta Lake information.
func (m *ObjectStorageCheckpointManager) GetCheckpointData(ctx context.Context, taskName string) (*CheckpointData, error) {
	objectPath := checkpointObjectPath(taskName)

	exists, err := m.provider.Exists(ctx, objectPath)
	if err != nil {
		return nil, fmt.Errorf("failed to check checkpoint existence: %w", err)
	}
	if !exists {
		m.log.Info("no checkpoint found, returning empty data", zap.String("task", taskName))
		return &CheckpointData{
			TaskName:   taskName,
			Checkpoint: 0,
			UpdatedAt:  0,
		}, nil
	}

	reader, err := m.provider.Download(ctx, objectPath)
	if err != nil {
		return nil, fmt.Errorf("failed to download checkpoint file: %w", err)
	}
	defer reader.Close()

	content, err := io.ReadAll(reader)
	if err != nil {
		return nil, fmt.Errorf("failed to read checkpoint file: %w", err)
	}

	data := CheckpointData{}
	if err := yaml.Unmarshal(content, &data); err != nil {
		return nil, fmt.Errorf("failed to unmarshal checkpoint data: %w", err)
	}

	m.log.Info("successfully retrieved checkpoint data",
		zap.String("task", taskName),
		zap.Int64("checkpoint", data.Checkpoint),
		zap.Int64("updated_at", data.UpdatedAt),
		zap.String("delta_table_path", data.DeltaTablePath),
		zap.Int64("delta_table_version", data.DeltaTableVersion))

	return &data, nil
}

// SaveCheckpointData saves full checkpoint data including Delta Lake information.
func (m *ObjectStorageCheckpointManager) SaveCheckpointData(ctx context.Context, data *CheckpointData) error {
	if data == nil {
		return fmt.Errorf("checkpoint data cannot be nil")
	}

	if data.UpdatedAt == 0 {
		data.UpdatedAt = time.Now().Unix()
	}

	content, err := yaml.Marshal(data)
	if err != nil {
		return fmt.Errorf("failed to marshal checkpoint data: %w", err)
	}

	objectPath := checkpointObjectPath(data.TaskName)
	if err := m.provider.Upload(ctx, objectPath, bytes.NewReader(content)); err != nil {
		return fmt.Errorf("failed to upload checkpoint file: %w", err)
	}

	m.log.Info("successfully saved checkpoint data",
		zap.String("task", data.TaskName),
		zap.Int64("checkpoint", data.Checkpoint),
		zap.String("object_path", objectPath),
		zap.String("delta_table_path", data.DeltaTablePath),
		zap.Int64("delta_table_version", data.DeltaTableVersion))

	return nil
}