task.go

// task.go
package main

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

var db *sql.DB

// Task 定义了任务的数据结构
type Task struct {
	SourceID       string       `json:"source_id"`
	URL            string       `json:"url"`
	ModelList      string       `json:"model_list"`
	ReconnectCount int          `json:"reconnect_count"`
	OperationType  int          `json:"operation_type"`
	Second         float64      `json:"second"`
	IsCascade      int          `json:"is_cascade"`
	CascadeLabel   string       `json:"cascade_label"`
	Status         string       `json:"status"`
	CreatedAt      time.Time    `json:"created_at"`
	UpdatedAt      sql.NullTime `json:"updated_at"` // Use sql.NullTime for nullable fields
}

// InitDB 初始化数据库连接并创建表
func InitDB(filepath string) error {
	var err error
	db, err = sql.Open("sqlite3", filepath)
	if err != nil {
		return err
	}

	taskTableSQL := `
    CREATE TABLE IF NOT EXISTS tasks (
        source_id TEXT PRIMARY KEY,
        url TEXT NOT NULL,
        model_list TEXT,
        reconnect_count INTEGER,
        operation_type INTEGER,
        second REAL,
        is_cascade INTEGER,
        cascade_label TEXT,
        status TEXT,
        created_at DATETIME,
        updated_at DATETIME
    );`

	logTableSQL := `
    CREATE TABLE IF NOT EXISTS task_logs (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        log TEXT,
        source_id TEXT,
        created_at DATETIME
    );`

	if _, err := db.Exec(taskTableSQL); err != nil {
		return err
	}
	if _, err := db.Exec(logTableSQL); err != nil {
		return err
	}
	log.Println("Database initialized successfully.")
	return nil
}

// InsertLog 记录任务日志
func InsertLog(sourceID, message string) {
	_, err := db.Exec("INSERT INTO task_logs (source_id, log, created_at) VALUES (?, ?, ?)",
		sourceID, message, time.Now())
	if err != nil {
		log.Printf("Error inserting log for %s: %v\n", sourceID, err)
	}
}

// CreateTask 在数据库中创建一个新任务
func CreateTask(task *Task) error {
	_, err := db.Exec(`INSERT INTO tasks (source_id, url, model_list, reconnect_count, operation_type, second, is_cascade, status, created_at)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
		task.SourceID, task.URL, task.ModelList, task.ReconnectCount, task.OperationType, task.Second, task.IsCascade, task.Status, task.CreatedAt)

	if err != nil {
		return err
	}
	InsertLog(task.SourceID, "CREATED")
	return nil
}

// GetTaskByID 从数据库中获取一个任务
func GetTaskByID(sourceID string) (*Task, error) {
	task := &Task{}
	row := db.QueryRow("SELECT * FROM tasks WHERE source_id = ?", sourceID)
	err := row.Scan(&task.SourceID, &task.URL, &task.ModelList, &task.ReconnectCount, &task.OperationType,
		&task.Second, &task.IsCascade, &task.CascadeLabel, &task.Status, &task.CreatedAt, &task.UpdatedAt)
	if err != nil {
		return nil, err
	}
	return task, nil
}

// GetAllTasks 获取所有任务
func GetAllTasks() ([]Task, error) {
	rows, err := db.Query("SELECT * FROM tasks")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var tasks []Task
	for rows.Next() {
		var t Task
		if err := rows.Scan(&t.SourceID, &t.URL, &t.ModelList, &t.ReconnectCount, &t.OperationType,
			&t.Second, &t.IsCascade, &t.CascadeLabel, &t.Status, &t.CreatedAt, &t.UpdatedAt); err != nil {
			log.Printf("Error scanning task row: %v\n", err)
			continue
		}
		tasks = append(tasks, t)
	}
	return tasks, nil
}

// GetRunnableTasks 获取所有可运行的任务
func GetRunnableTasks() ([]Task, error) {
	rows, err := db.Query("SELECT * FROM tasks WHERE status = 'FINISHED'")
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var tasks []Task
	for rows.Next() {
		var t Task
		if err := rows.Scan(&t.SourceID, &t.URL, &t.ModelList, &t.ReconnectCount, &t.OperationType,
			&t.Second, &t.IsCascade, &t.CascadeLabel, &t.Status, &t.CreatedAt, &t.UpdatedAt); err != nil {
			log.Printf("Error scanning task row: %v\n", err)
			continue
		}
		// 检查执行间隔
		if t.UpdatedAt.Valid && time.Since(t.UpdatedAt.Time) < (time.Duration(t.Second)*time.Second) {
			continue
		}
		tasks = append(tasks, t)
	}
	return tasks, nil
}

// UpdateTaskStatus 更新任务状态
func UpdateTaskStatus(sourceID, status string) error {
	_, err := db.Exec("UPDATE tasks SET status = ? WHERE source_id = ?", status, sourceID)
	if err == nil {
		InsertLog(sourceID, status)
	}
	return err
}

// UpdateTaskTimestamp 更新任务的 `updated_at` 时间戳
func UpdateTaskTimestamp(sourceID string) error {
	_, err := db.Exec("UPDATE tasks SET updated_at = ? WHERE source_id = ?", time.Now(), sourceID)
	return err
}

// DeleteTask 从数据库中删除任务
func DeleteTask(sourceID string) error {
	res, err := db.Exec("DELETE FROM tasks WHERE source_id = ?", sourceID)
	if err != nil {
		return err
	}
	rowsAffected, _ := res.RowsAffected()
	if rowsAffected == 0 {
		return fmt.Errorf("task with source_id '%s' not found", sourceID)
	}
	InsertLog(sourceID, "DELETED")
	return nil
}