package checkpoint

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

	awsstorage "github.com/tidbcloud/metering/pkg/common/utils/storage/aws"
	"go.uber.org/zap"
	"gopkg.in/yaml.v3"
)

const (
	checkpointFileName = "checkpoint.yaml"
)

// S3CheckpointManager implements NextGenCheckpoint using AWS S3 storage
type S3CheckpointManager struct {
	s3Client  awsstorage.S3Client
	bucket    string
	prefix    string
	cachePath string
	log       *zap.Logger
}

// NewS3CheckpointManager creates a new S3-based checkpoint manager
func NewS3CheckpointManager(s3Client awsstorage.S3Client, bucket, prefix, cachePath string, log *zap.Logger) NextGenCheckpoint {
	return &S3CheckpointManager{
		s3Client:  s3Client,
		bucket:    bucket,
		prefix:    prefix,
		cachePath: cachePath,
		log:       log,
	}
}

// GetCheckpoint retrieves the checkpoint for a specific task from S3
func (s *S3CheckpointManager) GetCheckpoint(ctx context.Context, taskName string) (int64, error) {
	s3FileName := fmt.Sprintf("%s/%s/%s", s.prefix, taskName, checkpointFileName)

	s.log.Debug("getting checkpoint from S3",
		zap.String("bucket", s.bucket),
		zap.String("key", s3FileName))

	// Create task-specific temporary directory to avoid conflicts
	taskTempDir := fmt.Sprintf("%s/%s_temp", s.cachePath, taskName)
	err := os.MkdirAll(taskTempDir, 0755)
	if err != nil {
		s.log.Error("failed to create task temp directory",
			zap.String("dir", taskTempDir),
			zap.Error(err))
		return 0, fmt.Errorf("failed to create task temp directory: %w", err)
	}
	defer os.RemoveAll(taskTempDir) // Clean up temp directory

	// Use the S3 client to download the checkpoint file to task-specific temp directory
	err = s.s3Client.DownloadObjects(ctx, []string{s3FileName}, taskTempDir)
	if err != nil {
		if strings.Contains(err.Error(), "NoSuchKey") || strings.Contains(err.Error(), "404") {
			s.log.Info("no checkpoint found, returning 0", zap.String("task", taskName))
			return 0, nil
		}
		s.log.Error("failed to download checkpoint file",
			zap.String("key", s3FileName),
			zap.Error(err))
		return 0, fmt.Errorf("failed to download checkpoint file: %w", err)
	}

	// Read the downloaded file - the file will be saved with its base name in temp directory
	tempFilePath := fmt.Sprintf("%s/%s", taskTempDir, checkpointFileName)
	content, err := os.ReadFile(tempFilePath)
	if err != nil {
		s.log.Error("failed to read checkpoint file",
			zap.String("file", tempFilePath),
			zap.Error(err))
		return 0, fmt.Errorf("failed to read checkpoint file: %w", err)
	}

	// Parse the checkpoint data
	data := CheckpointData{}
	err = yaml.Unmarshal(content, &data)
	if err != nil {
		s.log.Error("failed to unmarshal checkpoint data",
			zap.String("content", string(content)),
			zap.Error(err))
		return 0, fmt.Errorf("failed to unmarshal checkpoint data: %w", err)
	}

	s.log.Info("successfully retrieved checkpoint",
		zap.String("task", taskName),
		zap.Int64("checkpoint", data.Checkpoint),
		zap.Int64("updated_at", data.UpdatedAt))

	return data.Checkpoint, nil
}

// SaveCheckpoint saves the checkpoint for a specific task to S3
func (s *S3CheckpointManager) SaveCheckpoint(ctx context.Context, taskName string, checkpoint int64) error {
	data := &CheckpointData{
		Checkpoint: checkpoint,
		UpdatedAt:  time.Now().Unix(),
		TaskName:   taskName,
	}

	return s.SaveCheckpointData(ctx, data)
}

// GetCheckpointData retrieves full checkpoint data including Delta Lake information
func (s *S3CheckpointManager) GetCheckpointData(ctx context.Context, taskName string) (*CheckpointData, error) {
	s3FileName := fmt.Sprintf("%s/%s/%s", s.prefix, taskName, checkpointFileName)

	s.log.Debug("getting checkpoint data from S3",
		zap.String("bucket", s.bucket),
		zap.String("key", s3FileName))

	// Create task-specific temporary directory to avoid conflicts
	taskTempDir := fmt.Sprintf("%s/%s_temp", s.cachePath, taskName)
	err := os.MkdirAll(taskTempDir, 0755)
	if err != nil {
		s.log.Error("failed to create task temp directory",
			zap.String("dir", taskTempDir),
			zap.Error(err))
		return nil, fmt.Errorf("failed to create task temp directory: %w", err)
	}
	defer os.RemoveAll(taskTempDir) // Clean up temp directory

	// Use the S3 client to download the checkpoint file to task-specific temp directory
	err = s.s3Client.DownloadObjects(ctx, []string{s3FileName}, taskTempDir)
	if err != nil {
		if strings.Contains(err.Error(), "NoSuchKey") || strings.Contains(err.Error(), "404") {
			s.log.Info("no checkpoint found, returning empty data", zap.String("task", taskName))
			return &CheckpointData{
				TaskName:   taskName,
				Checkpoint: 0,
				UpdatedAt:  0,
			}, nil
		}
		s.log.Error("failed to download checkpoint file",
			zap.String("key", s3FileName),
			zap.Error(err))
		return nil, fmt.Errorf("failed to download checkpoint file: %w", err)
	}

	// Read the downloaded file - the file will be saved with its base name in temp directory
	tempFilePath := fmt.Sprintf("%s/%s", taskTempDir, checkpointFileName)
	content, err := os.ReadFile(tempFilePath)
	if err != nil {
		s.log.Error("failed to read checkpoint file",
			zap.String("file", tempFilePath),
			zap.Error(err))
		return nil, fmt.Errorf("failed to read checkpoint file: %w", err)
	}

	// Parse the checkpoint data
	data := CheckpointData{}
	err = yaml.Unmarshal(content, &data)
	if err != nil {
		s.log.Error("failed to unmarshal checkpoint data",
			zap.String("content", string(content)),
			zap.Error(err))
		return nil, fmt.Errorf("failed to unmarshal checkpoint data: %w", err)
	}

	s.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 (s *S3CheckpointManager) SaveCheckpointData(ctx context.Context, data *CheckpointData) error {
	if data == nil {
		return fmt.Errorf("checkpoint data cannot be nil")
	}

	// Set update time if not already set
	if data.UpdatedAt == 0 {
		data.UpdatedAt = time.Now().Unix()
	}

	s.log.Debug("saving checkpoint data to S3",
		zap.String("task", data.TaskName),
		zap.Int64("checkpoint", data.Checkpoint),
		zap.String("delta_table_path", data.DeltaTablePath),
		zap.Int64("delta_table_version", data.DeltaTableVersion))

	// Marshal the checkpoint data to YAML
	content, err := yaml.Marshal(data)
	if err != nil {
		s.log.Error("failed to marshal checkpoint data",
			zap.String("task", data.TaskName),
			zap.Error(err))
		return fmt.Errorf("failed to marshal checkpoint data: %w", err)
	}

	// Create a temporary file with the checkpoint content
	tempFilePath := fmt.Sprintf("%s/%s_%s_%d", s.cachePath, data.TaskName, checkpointFileName, time.Now().Unix())
	err = os.WriteFile(tempFilePath, content, 0600)
	if err != nil {
		s.log.Error("failed to write temporary checkpoint file",
			zap.String("file", tempFilePath),
			zap.Error(err))
		return fmt.Errorf("failed to write temporary checkpoint file: %w", err)
	}

	// Upload the file to S3
	s3FileName := fmt.Sprintf("%s/%s/%s", s.prefix, data.TaskName, checkpointFileName)
	err = s.s3Client.UploadObject(ctx, s3FileName, tempFilePath)
	if err != nil {
		s.log.Error("failed to upload checkpoint file to S3",
			zap.String("key", s3FileName),
			zap.Error(err))
		return fmt.Errorf("failed to upload checkpoint file to S3: %w", err)
	}

	s.log.Info("successfully saved checkpoint data",
		zap.String("task", data.TaskName),
		zap.Int64("checkpoint", data.Checkpoint),
		zap.String("s3_key", s3FileName),
		zap.String("delta_table_path", data.DeltaTablePath),
		zap.Int64("delta_table_version", data.DeltaTableVersion))

	// Clean up temporary file
	_ = os.Remove(tempFilePath)

	return nil
}