modelRT/database/async_task_extended.go

// Package database define database operation functions
package database

import (
	"context"
	"time"

	"modelRT/orm"

	"github.com/gofrs/uuid"
	"gorm.io/gorm"
)

// UpdateTaskStarted updates task start time and status to running
func UpdateTaskStarted(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, startedAt int64) error {
	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	result := tx.WithContext(cancelCtx).
		Model(&orm.AsyncTask{}).
		Where("task_id = ?", taskID).
		Updates(map[string]any{
			"status":     orm.AsyncTaskStatusRunning,
			"started_at": startedAt,
		})

	return result.Error
}

// UpdateTaskRetryInfo updates task retry information
func UpdateTaskRetryInfo(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, retryCount int, nextRetryTime int64) error {
	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	updateData := map[string]any{
		"retry_count": retryCount,
	}
	if nextRetryTime <= 0 {
		updateData["next_retry_time"] = nil
	} else {
		updateData["next_retry_time"] = nextRetryTime
	}

	result := tx.WithContext(cancelCtx).
		Model(&orm.AsyncTask{}).
		Where("task_id = ?", taskID).
		Updates(updateData)

	return result.Error
}

// UpdateTaskErrorInfo updates task error information
func UpdateTaskErrorInfo(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, errorMsg, stackTrace string) error {
	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	result := tx.WithContext(cancelCtx).
		Model(&orm.AsyncTask{}).
		Where("task_id = ?", taskID).
		Updates(map[string]any{
			"failure_reason": errorMsg,
			"stack_trace":    stackTrace,
		})

	return result.Error
}

// UpdateTaskExecutionTime updates task execution time
func UpdateTaskExecutionTime(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, executionTime int64) error {
	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	result := tx.WithContext(cancelCtx).
		Model(&orm.AsyncTask{}).
		Where("task_id = ?", taskID).
		Update("execution_time", executionTime)

	return result.Error
}

// UpdateTaskWorkerID updates the worker ID that is processing the task
func UpdateTaskWorkerID(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, workerID string) error {
	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	result := tx.WithContext(cancelCtx).
		Model(&orm.AsyncTask{}).
		Where("task_id = ?", taskID).
		Update("worker_id", workerID)

	return result.Error
}

// UpdateTaskPriority updates task priority
func UpdateTaskPriority(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, priority int) error {
	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	result := tx.WithContext(cancelCtx).
		Model(&orm.AsyncTask{}).
		Where("task_id = ?", taskID).
		Update("priority", priority)

	return result.Error
}

// UpdateTaskQueueName updates task queue name
func UpdateTaskQueueName(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, queueName string) error {
	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	result := tx.WithContext(cancelCtx).
		Model(&orm.AsyncTask{}).
		Where("task_id = ?", taskID).
		Update("queue_name", queueName)

	return result.Error
}

// UpdateTaskCreatedBy updates task creator information
func UpdateTaskCreatedBy(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, createdBy string) error {
	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	result := tx.WithContext(cancelCtx).
		Model(&orm.AsyncTask{}).
		Where("task_id = ?", taskID).
		Update("created_by", createdBy)

	return result.Error
}

// UpdateTaskResultWithMetrics updates task result with execution metrics
func UpdateTaskResultWithMetrics(ctx context.Context, tx *gorm.DB, taskID uuid.UUID, executionTime int64, memoryUsage *int64, cpuUsage *float64, retryCount int, completedAt int64) error {
	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	result := tx.WithContext(cancelCtx).
		Model(&orm.AsyncTaskResult{}).
		Where("task_id = ?", taskID).
		Updates(map[string]any{
			"execution_time": executionTime,
			"memory_usage":   memoryUsage,
			"cpu_usage":      cpuUsage,
			"retry_count":    retryCount,
			"completed_at":   completedAt,
		})

	return result.Error
}

// GetTasksForRetry retrieves tasks that are due for retry
func GetTasksForRetry(ctx context.Context, tx *gorm.DB, limit int) ([]orm.AsyncTask, error) {
	var tasks []orm.AsyncTask

	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	now := time.Now().Unix()
	result := tx.WithContext(cancelCtx).
		Where("status = ? AND next_retry_time IS NOT NULL AND next_retry_time <= ?", orm.AsyncTaskStatusFailed, now).
		Order("next_retry_time ASC").
		Limit(limit).
		Find(&tasks)

	if result.Error != nil {
		return nil, result.Error
	}

	return tasks, nil
}

// GetTasksByPriority retrieves tasks by priority order
func GetTasksByPriority(ctx context.Context, tx *gorm.DB, status orm.AsyncTaskStatus, limit int) ([]orm.AsyncTask, error) {
	var tasks []orm.AsyncTask

	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	result := tx.WithContext(cancelCtx).
		Where("status = ?", status).
		Order("priority DESC, created_at ASC").
		Limit(limit).
		Find(&tasks)

	if result.Error != nil {
		return nil, result.Error
	}

	return tasks, nil
}

// GetTasksByWorkerID retrieves tasks being processed by a specific worker
func GetTasksByWorkerID(ctx context.Context, tx *gorm.DB, workerID string) ([]orm.AsyncTask, error) {
	var tasks []orm.AsyncTask

	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
	defer cancel()

	result := tx.WithContext(cancelCtx).
		Where("worker_id = ? AND status = ?", workerID, orm.AsyncTaskStatusRunning).
		Find(&tasks)

	if result.Error != nil {
		return nil, result.Error
	}

	return tasks, nil
}

// CleanupStaleTasks marks tasks as failed if they have been running for too long
func CleanupStaleTasks(ctx context.Context, tx *gorm.DB, timeoutSeconds int64) (int64, error) {
	cancelCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()

	threshold := time.Now().Unix() - timeoutSeconds
	result := tx.WithContext(cancelCtx).
		Model(&orm.AsyncTask{}).
		Where("status = ? AND started_at IS NOT NULL AND started_at < ?", orm.AsyncTaskStatusRunning, threshold).
		Updates(map[string]any{
			"status":        orm.AsyncTaskStatusFailed,
			"failure_reason": "task timeout",
			"finished_at":   time.Now().Unix(),
		})

	return result.RowsAffected, result.Error
}