MeBox/internal/repository/scrape_task_repository.go

package repository

import (
	"context"
	"errors"
	"strings"
	"sync"
	"time"

	"gorm.io/gorm"

	"github.com/ShukeBta/MMTL/internal/model"
)

var scrapeClaimMu sync.Mutex

// ScrapeTaskRepository persists model.ScrapeTask.
type ScrapeTaskRepository struct{ db *gorm.DB }

func (r *ScrapeTaskRepository) Create(ctx context.Context, t *model.ScrapeTask) error {
	return withSQLiteBusyRetry(ctx, func() error {
		return r.db.WithContext(ctx).Create(t).Error
	})
}

func (r *ScrapeTaskRepository) CreateBatch(ctx context.Context, tasks []model.ScrapeTask) error {
	if len(tasks) == 0 {
		return nil
	}
	return withSQLiteBusyRetry(ctx, func() error {
		return r.db.WithContext(ctx).CreateInBatches(tasks, 100).Error
	})
}

func (r *ScrapeTaskRepository) FindByID(ctx context.Context, id string) (*model.ScrapeTask, error) {
	var t model.ScrapeTask
	err := r.db.WithContext(ctx).Where("id = ?", id).First(&t).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return nil, nil
	}
	return &t, err
}

func (r *ScrapeTaskRepository) FindActiveByMediaID(ctx context.Context, mediaID string) (*model.ScrapeTask, error) {
	var t model.ScrapeTask
	err := r.db.WithContext(ctx).
		Where("media_id = ? AND status IN ?", mediaID, []string{model.ScrapeTaskPending, model.ScrapeTaskRunning}).
		First(&t).Error
	if errors.Is(err, gorm.ErrRecordNotFound) {
		return nil, nil
	}
	return &t, err
}

func (r *ScrapeTaskRepository) List(ctx context.Context, status string, page, pageSize int) ([]model.ScrapeTask, int64, error) {
	if page < 1 {
		page = 1
	}
	if pageSize < 1 || pageSize > 200 {
		pageSize = 50
	}
	q := r.db.WithContext(ctx).Model(&model.ScrapeTask{})
	if strings.TrimSpace(status) != "" && status != "all" {
		q = q.Where("status = ?", strings.TrimSpace(status))
	}
	var total int64
	if err := q.Count(&total).Error; err != nil {
		return nil, 0, err
	}
	var rows []model.ScrapeTask
	err := q.Order("created_at desc").
		Offset((page - 1) * pageSize).
		Limit(pageSize).
		Find(&rows).Error
	return rows, total, err
}

func (r *ScrapeTaskRepository) CountByStatus(ctx context.Context) (map[string]int64, error) {
	var rows []struct {
		Status string
		Count  int64
	}
	err := r.db.WithContext(ctx).Model(&model.ScrapeTask{}).
		Select("status, count(*) as count").
		Group("status").Scan(&rows).Error
	if err != nil {
		return nil, err
	}
	out := map[string]int64{}
	for _, row := range rows {
		out[row.Status] = row.Count
	}
	return out, nil
}

// ClaimPending picks pending scrape tasks and marks them running.
func (r *ScrapeTaskRepository) ClaimPending(ctx context.Context, limit int) ([]model.ScrapeTask, error) {
	scrapeClaimMu.Lock()
	defer scrapeClaimMu.Unlock()

	var rows []model.ScrapeTask
	err := withSQLiteBusyRetry(ctx, func() error {
		return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
			if err := tx.Where("status = ?", model.ScrapeTaskPending).
				Order("created_at asc").Limit(limit).Find(&rows).Error; err != nil {
				return err
			}
			if len(rows) == 0 {
				return nil
			}
			ids := make([]string, 0, len(rows))
			now := time.Now()
			for i := range rows {
				ids = append(ids, rows[i].ID)
				rows[i].Status = model.ScrapeTaskRunning
				rows[i].StartedAt = &now
			}
			return tx.Model(&model.ScrapeTask{}).Where("id IN ?", ids).
				Updates(map[string]any{"status": model.ScrapeTaskRunning, "started_at": now}).Error
		})
	})
	if err != nil {
		return nil, err
	}
	return rows, nil
}

func (r *ScrapeTaskRepository) Update(ctx context.Context, t *model.ScrapeTask) error {
	return withSQLiteBusyRetry(ctx, func() error {
		return r.db.WithContext(ctx).Model(&model.ScrapeTask{}).Where("id = ?", t.ID).Updates(map[string]any{
			"status":        t.Status,
			"error":         t.Error,
			"provider":      t.Provider,
			"matched_title": t.MatchedTitle,
			"matched_year":  t.MatchedYear,
			"poster_url":    t.PosterURL,
			"backdrop_url":  t.BackdropURL,
			"retry_count":   t.RetryCount,
			"started_at":    t.StartedAt,
			"finished_at":   t.FinishedAt,
			"updated_at":    time.Now(),
		}).Error
	})
}

func (r *ScrapeTaskRepository) Delete(ctx context.Context, id string) error {
	return withSQLiteBusyRetry(ctx, func() error {
		return r.db.WithContext(ctx).Unscoped().Where("id = ?", id).Delete(&model.ScrapeTask{}).Error
	})
}

func (r *ScrapeTaskRepository) DeleteBatch(ctx context.Context, ids []string) (int64, error) {
	if len(ids) == 0 {
		return 0, nil
	}
	var count int64
	err := withSQLiteBusyRetry(ctx, func() error {
		res := r.db.WithContext(ctx).Unscoped().Where("id IN ?", ids).Delete(&model.ScrapeTask{})
		count = res.RowsAffected
		return res.Error
	})
	return count, err
}

func (r *ScrapeTaskRepository) RetryBatch(ctx context.Context, ids []string) (int64, error) {
	if len(ids) == 0 {
		return 0, nil
	}
	var count int64
	err := withSQLiteBusyRetry(ctx, func() error {
		res := r.db.WithContext(ctx).Model(&model.ScrapeTask{}).
			Where("id IN ? AND status IN ?", ids, []string{model.ScrapeTaskFailed, model.ScrapeTaskCanceled}).
			Updates(map[string]any{
				"status":      model.ScrapeTaskPending,
				"error":       "",
				"retry_count": 0,
				"started_at":  nil,
				"finished_at": nil,
				"updated_at":  time.Now(),
			})
		count = res.RowsAffected
		return res.Error
	})
	return count, err
}

func (r *ScrapeTaskRepository) CancelBatch(ctx context.Context, ids []string) (int64, error) {
	if len(ids) == 0 {
		return 0, nil
	}
	now := time.Now()
	var count int64
	err := withSQLiteBusyRetry(ctx, func() error {
		res := r.db.WithContext(ctx).Model(&model.ScrapeTask{}).
			Where("id IN ? AND status IN ?", ids, []string{model.ScrapeTaskPending, model.ScrapeTaskRunning}).
			Updates(map[string]any{
				"status":      model.ScrapeTaskCanceled,
				"error":       "已批量取消",
				"finished_at": now,
				"updated_at":  now,
			})
		count = res.RowsAffected
		return res.Error
	})
	return count, err
}

func (r *ScrapeTaskRepository) ClearDone(ctx context.Context) (int64, error) {
	var count int64
	err := withSQLiteBusyRetry(ctx, func() error {
		res := r.db.WithContext(ctx).Unscoped().Where("status = ?", model.ScrapeTaskDone).Delete(&model.ScrapeTask{})
		count = res.RowsAffected
		return res.Error
	})
	return count, err
}

func (r *ScrapeTaskRepository) ClearFinished(ctx context.Context) (int64, error) {
	var count int64
	err := withSQLiteBusyRetry(ctx, func() error {
		res := r.db.WithContext(ctx).Unscoped().Where("status IN ?", []string{model.ScrapeTaskDone, model.ScrapeTaskFailed, model.ScrapeTaskCanceled}).
			Delete(&model.ScrapeTask{})
		count = res.RowsAffected
		return res.Error
	})
	return count, err
}

func (r *ScrapeTaskRepository) ClearCanceled(ctx context.Context) (int64, error) {
	var count int64
	err := withSQLiteBusyRetry(ctx, func() error {
		res := r.db.WithContext(ctx).Unscoped().Where("status = ?", model.ScrapeTaskCanceled).Delete(&model.ScrapeTask{})
		count = res.RowsAffected
		return res.Error
	})
	return count, err
}

func (r *ScrapeTaskRepository) RetryAllFailed(ctx context.Context) (int64, error) {
	var count int64
	err := withSQLiteBusyRetry(ctx, func() error {
		res := r.db.WithContext(ctx).Model(&model.ScrapeTask{}).
			Where("status = ?", model.ScrapeTaskFailed).
			Updates(map[string]any{
				"status":      model.ScrapeTaskPending,
				"error":       "",
				"retry_count": 0,
				"started_at":  nil,
				"finished_at": nil,
				"updated_at":  time.Now(),
			})
		count = res.RowsAffected
		return res.Error
	})
	return count, err
}

func (r *ScrapeTaskRepository) CancelPending(ctx context.Context) (int64, error) {
	now := time.Now()
	var count int64
	err := withSQLiteBusyRetry(ctx, func() error {
		res := r.db.WithContext(ctx).Model(&model.ScrapeTask{}).
			Where("status IN ?", []string{model.ScrapeTaskPending, model.ScrapeTaskRunning}).
			Updates(map[string]any{
				"status":      model.ScrapeTaskCanceled,
				"error":       "已批量取消",
				"finished_at": now,
				"updated_at":  now,
			})
		count = res.RowsAffected
		return res.Error
	})
	return count, err
}