hireme/internal/db/tasks.go

package db

import (
	"database/sql"
	"encoding/json"
	"fmt"

	"git.pengzhan.dev/haopengzhan/hireme/internal/logger"
)

// QueueStats holds task queue counters
type QueueStats struct {
	Pending   int `json:"pending"`
	Running   int `json:"running"`
	Failed    int `json:"failed"`
	Completed int `json:"completed"`
}

// GetQueueStats returns counts for each task status
func (db *DB) GetQueueStats() (*QueueStats, error) {
	query := `
	SELECT
		COALESCE(SUM(CASE WHEN status = 'PENDING' THEN 1 ELSE 0 END), 0),
		COALESCE(SUM(CASE WHEN status = 'RUNNING' THEN 1 ELSE 0 END), 0),
		COALESCE(SUM(CASE WHEN status = 'FAILED' THEN 1 ELSE 0 END), 0),
		COALESCE(SUM(CASE WHEN status = 'COMPLETED' THEN 1 ELSE 0 END), 0)
	FROM tasks
	`
	var stats QueueStats
	err := db.QueryRow(query).Scan(&stats.Pending, &stats.Running, &stats.Failed, &stats.Completed)
	if err != nil {
		return nil, err
	}
	return &stats, nil
}

// EnqueueTask adds a new task to the queue
func (db *DB) EnqueueTask(taskType string, payload interface{}) error {
	payloadBytes, err := json.Marshal(payload)
	if err != nil {
		return fmt.Errorf("failed to marshal task payload: %w", err)
	}

	query := `
	INSERT INTO tasks (task_type, payload, status)
	VALUES (?, ?, 'PENDING')
	`
	_, err = db.Exec(query, taskType, string(payloadBytes))
	return err
}

// ClaimNextTask atomically finds and claims the next pending task
func (db *DB) ClaimNextTask() (*Task, error) {
	query := `
	UPDATE tasks 
	SET status = 'RUNNING', updated_at = CURRENT_TIMESTAMP 
	WHERE id = (
		SELECT id FROM tasks 
		WHERE status = 'PENDING' AND next_retry_at <= CURRENT_TIMESTAMP 
		ORDER BY created_at ASC 
		LIMIT 1
	)
	RETURNING id, task_type, payload, status, error_message, retry_count, next_retry_at, created_at, updated_at
	`
	
	var task Task
	var errMsg sql.NullString
	err := db.QueryRow(query).Scan(
		&task.ID, &task.TaskType, &task.Payload, &task.Status, &errMsg,
		&task.RetryCount, &task.NextRetryAt, &task.CreatedAt, &task.UpdatedAt,
	)
	
	if err == sql.ErrNoRows {
		return nil, nil // No tasks available
	}
	if err != nil {
		return nil, err
	}

	if errMsg.Valid {
		task.ErrorMessage = errMsg.String
	}

	return &task, nil
}

// CompleteTask marks a task as COMPLETED
func (db *DB) CompleteTask(id int64) error {
	query := `UPDATE tasks SET status = 'COMPLETED', updated_at = CURRENT_TIMESTAMP WHERE id = ?`
	_, err := db.Exec(query, id)
	return err
}

// FailTask marks a task as FAILED or PENDING (for retry) using exponential backoff
func (db *DB) FailTask(id int64, errStr string) error {
	// First get current retry count
	var retryCount int
	err := db.QueryRow(`SELECT retry_count FROM tasks WHERE id = ?`, id).Scan(&retryCount)
	if err != nil {
		logger.Log.Errorw("FailTask: failed to get retry_count", "component", "DB", "task_id", id, "error", err)
		return err
	}

	retryCount++
	
	// Max retries = 5 (1s, 2s, 4s, 8s, 16s)
	status := "PENDING"
	if retryCount > 5 {
		status = "FAILED"
	}

	// Exponential backoff
	backoffSeconds := 1 << retryCount // 2^retryCount
	
	query := `
	UPDATE tasks 
	SET status = ?, 
	    error_message = ?, 
	    retry_count = ?,
	    next_retry_at = datetime('now', '+' || ? || ' seconds'),
	    updated_at = CURRENT_TIMESTAMP
	WHERE id = ?
	`
	_, err = db.Exec(query, status, errStr, retryCount, fmt.Sprintf("%d", backoffSeconds), id)
	return err
}

// HasPendingFetchTask checks if there is already a PENDING or RUNNING fetch task for a given canonical ID
func (db *DB) HasPendingFetchTask(canonicalID string) (bool, error) {
	query := `
	SELECT COUNT(*) FROM tasks 
	WHERE task_type = 'FETCH' 
	  AND status IN ('PENDING', 'RUNNING') 
	  AND payload LIKE ?
	`
	likePattern := fmt.Sprintf(`%%"canonical_id":"%s"%%`, canonicalID)
	
	var count int
	err := db.QueryRow(query, likePattern).Scan(&count)
	if err != nil {
		return false, err
	}
	return count > 0, nil
}

// ClearTaskQueue deletes all tasks or only pending/failed ones
func (db *DB) ClearTaskQueue() error {
	query := `DELETE FROM tasks WHERE status IN ('PENDING', 'FAILED', 'COMPLETED')`
	_, err := db.Exec(query)
	return err
}

// GetQueueStatus returns counts of tasks by status
func (db *DB) GetQueueStatus() (map[string]int, error) {
	query := `SELECT status, COUNT(*) FROM tasks GROUP BY status`
	rows, err := db.Query(query)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rows.Close() }()

	counts := make(map[string]int)
	for rows.Next() {
		var status string
		var count int
		if err := rows.Scan(&status, &count); err != nil {
			return nil, err
		}
		counts[status] = count
	}
	return counts, nil
}

// ReprocessAllJobs wipes all job analysis results and re-queues PROCESS tasks for all existing raw jobs.
func (db *DB) ReprocessAllJobs() error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()

	// 1. Delete all rows from job_analyses
	if _, err := tx.Exec(`DELETE FROM job_analyses`); err != nil {
		return fmt.Errorf("failed to clear job_analyses: %w", err)
	}

	// 2. Delete all existing PROCESS and ANALYZE tasks
	if _, err := tx.Exec(`DELETE FROM tasks WHERE task_type IN ('PROCESS', 'ANALYZE')`); err != nil {
		return fmt.Errorf("failed to clear old tasks: %w", err)
	}

	// 3. Queue up new PROCESS tasks for all raw jobs
	rows, err := tx.Query(`SELECT id FROM raw_jobs`)
	if err != nil {
		return fmt.Errorf("failed to get raw jobs: %w", err)
	}
	defer func() { _ = rows.Close() }()

	var rawJobIDs []int64
	for rows.Next() {
		var rawJobID int64
		if err := rows.Scan(&rawJobID); err != nil {
			return err
		}
		rawJobIDs = append(rawJobIDs, rawJobID)
	}
	rows.Close() // Explicitly close to release any read locks

	// 4. Enqueue PROCESS task for each job
	for _, rawJobID := range rawJobIDs {
		payloadJSON := fmt.Sprintf(`{"raw_job_id":%d}`, rawJobID)
		if _, err := tx.Exec(`INSERT INTO tasks (task_type, payload, status) VALUES ('PROCESS', ?, 'PENDING')`, payloadJSON); err != nil {
			return fmt.Errorf("failed to enqueue process task for job %d: %w", rawJobID, err)
		}
	}

	return tx.Commit()
}