package postgres

import (
	"context"
	"time"

	"github.com/redis/go-redis/v9"
	"gorm.io/gorm"

	"github.com/break/junhong_cmp_fiber/internal/model"
	"github.com/break/junhong_cmp_fiber/pkg/constants"
)

// ExportShardTaskStore 导出分片任务数据访问层。
type ExportShardTaskStore struct {
	db    *gorm.DB
	redis *redis.Client
}

// NewExportShardTaskStore 创建导出分片任务 Store。
func NewExportShardTaskStore(db *gorm.DB, redis *redis.Client) *ExportShardTaskStore {
	return &ExportShardTaskStore{db: db, redis: redis}
}

// CreateBatch 批量创建分片任务。
func (s *ExportShardTaskStore) CreateBatch(ctx context.Context, shards []*model.ExportShardTask) error {
	if len(shards) == 0 {
		return nil
	}
	return s.db.WithContext(ctx).Create(&shards).Error
}

// GetByIDForWorker 按 ID 查询分片任务（Worker 使用）。
func (s *ExportShardTaskStore) GetByIDForWorker(ctx context.Context, id uint) (*model.ExportShardTask, error) {
	var shard model.ExportShardTask
	if err := s.db.WithContext(ctx).Where("id = ?", id).First(&shard).Error; err != nil {
		return nil, err
	}
	return &shard, nil
}

// ListByTaskID 按任务 ID 查询分片任务列表。
func (s *ExportShardTaskStore) ListByTaskID(ctx context.Context, taskID uint) ([]*model.ExportShardTask, error) {
	var list []*model.ExportShardTask
	if err := s.db.WithContext(ctx).
		Where("task_id = ?", taskID).
		Order("shard_no ASC").
		Find(&list).Error; err != nil {
		return nil, err
	}
	return list, nil
}

// TryMarkProcessingIfPending 尝试把分片从待处理更新为处理中。
func (s *ExportShardTaskStore) TryMarkProcessingIfPending(ctx context.Context, shardID uint, updater uint) (bool, error) {
	now := time.Now()
	result := s.db.WithContext(ctx).
		Model(&model.ExportShardTask{}).
		Where("id = ? AND status = ?", shardID, constants.ExportShardStatusPending).
		Updates(map[string]any{
			"status":     constants.ExportShardStatusProcessing,
			"started_at": now,
			"updated_at": now,
			"updater":    updater,
		})
	return result.RowsAffected > 0, result.Error
}

// MarkSuccess 标记分片为成功。
func (s *ExportShardTaskStore) MarkSuccess(ctx context.Context, shardID uint, updater uint, rowCount int, fileKey string, fileSize int64) (bool, error) {
	now := time.Now()
	result := s.db.WithContext(ctx).
		Model(&model.ExportShardTask{}).
		Where("id = ? AND status IN ?", shardID, []int{constants.ExportShardStatusPending, constants.ExportShardStatusProcessing}).
		Updates(map[string]any{
			"status":         constants.ExportShardStatusSuccess,
			"row_count":      rowCount,
			"processed_rows": rowCount,
			"file_key":       fileKey,
			"file_size":      fileSize,
			"completed_at":   now,
			"updated_at":     now,
			"updater":        updater,
			"error_message":  "",
		})
	return result.RowsAffected > 0, result.Error
}

// MarkFailed 标记分片为失败。
func (s *ExportShardTaskStore) MarkFailed(ctx context.Context, shardID uint, updater uint, errorMessage string) (bool, error) {
	now := time.Now()
	result := s.db.WithContext(ctx).
		Model(&model.ExportShardTask{}).
		Where("id = ? AND status IN ?", shardID, []int{constants.ExportShardStatusPending, constants.ExportShardStatusProcessing}).
		Updates(map[string]any{
			"status":        constants.ExportShardStatusFailed,
			"error_message": errorMessage,
			"completed_at":  now,
			"updated_at":    now,
			"updater":       updater,
		})
	return result.RowsAffected > 0, result.Error
}

// MarkCancelled 标记分片为取消。
func (s *ExportShardTaskStore) MarkCancelled(ctx context.Context, shardID uint, updater uint, errorMessage string) (bool, error) {
	now := time.Now()
	result := s.db.WithContext(ctx).
		Model(&model.ExportShardTask{}).
		Where("id = ? AND status IN ?", shardID, []int{constants.ExportShardStatusPending, constants.ExportShardStatusProcessing}).
		Updates(map[string]any{
			"status":        constants.ExportShardStatusCancelled,
			"error_message": errorMessage,
			"completed_at":  now,
			"updated_at":    now,
			"updater":       updater,
		})
	return result.RowsAffected > 0, result.Error
}

// CancelByTaskID 批量取消指定主任务下未完成的分片。
func (s *ExportShardTaskStore) CancelByTaskID(ctx context.Context, taskID uint, updater uint, errorMessage string) error {
	now := time.Now()
	return s.db.WithContext(ctx).
		Model(&model.ExportShardTask{}).
		Where("task_id = ? AND status IN ?", taskID, []int{constants.ExportShardStatusPending, constants.ExportShardStatusProcessing}).
		Updates(map[string]any{
			"status":        constants.ExportShardStatusCancelled,
			"error_message": errorMessage,
			"completed_at":  now,
			"updated_at":    now,
			"updater":       updater,
		}).Error
}

// CountUnfinishedByTaskID 统计未完成分片数量。
func (s *ExportShardTaskStore) CountUnfinishedByTaskID(ctx context.Context, taskID uint) (int64, error) {
	var count int64
	err := s.db.WithContext(ctx).
		Model(&model.ExportShardTask{}).
		Where("task_id = ? AND status IN ?", taskID, []int{constants.ExportShardStatusPending, constants.ExportShardStatusProcessing}).
		Count(&count).Error
	return count, err
}

// CountByTaskIDAndStatus 统计指定状态分片数量。
func (s *ExportShardTaskStore) CountByTaskIDAndStatus(ctx context.Context, taskID uint, status int) (int64, error) {
	var count int64
	err := s.db.WithContext(ctx).
		Model(&model.ExportShardTask{}).
		Where("task_id = ? AND status = ?", taskID, status).
		Count(&count).Error
	return count, err
}
