BiliDown/cmd/bilidown/main.go

package main

import (
	"context"
	"database/sql"
	"embed"
	"fmt"
	"io/fs"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"bilidown/internal/logger"
	"bilidown/internal/router"
	"bilidown/internal/task"
	"bilidown/internal/util"

	_ "modernc.org/sqlite"
)

//go:embed all:static
var staticFiles embed.FS

const (
	HTTP_PORT = 8098      // HTTP 服务器端口
	HTTP_HOST = ""        // HTTP 服务器主机
	VERSION   = "v2.1.2" // 软件版本号
)

var urlLocal = fmt.Sprintf("http://127.0.0.1:%d", HTTP_PORT)
var ffmpegAvailable bool
var server *http.Server

func main() {
	ffmpegAvailable = checkFFmpeg()
	logger.FFmpegStatus(ffmpegAvailable)
	mustInitTables()
	mustRunServer()
	logger.ServerStarted(HTTP_PORT, VERSION)

	// 等待中断信号，实现优雅退出
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit

	logger.Info("正在关闭服务...")
	gracefulShutdown()
}

// checkFFmpeg 检测 ffmpeg 的安装情况，返回是否可用
func checkFFmpeg() bool {
	_, err := util.GetFFmpegPath()
	return err == nil
}

// mustRunServer 配置和启动 HTTP 服务器
func mustRunServer() {
	// 从嵌入的文件系统中获取静态文件
	staticFS, err := fs.Sub(staticFiles, "static")
	if err != nil {
		logger.Fatal("加载静态文件失败: " + err.Error())
	}
	// 前端静态文件
	http.Handle("/", http.FileServer(http.FS(staticFS)))
	// 后端 API 接口
	apiRouter := router.API()
	// 添加 FFmpeg 状态检查接口
	http.HandleFunc("/api/checkFFmpeg", func(w http.ResponseWriter, r *http.Request) {
		util.Res{Success: true, Data: map[string]bool{"available": ffmpegAvailable}}.Write(w)
	})
	http.Handle("/api/", http.StripPrefix("/api", apiRouter))
	// 启动 HTTP 服务器
	server = &http.Server{
		Addr:    fmt.Sprintf("%s:%d", HTTP_HOST, HTTP_PORT),
		Handler: nil,
	}
	go func() {
		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			logger.Fatal("http.ListenAndServe: " + err.Error())
		}
	}()
}

// gracefulShutdown 优雅退出：保存任务状态并关闭服务
func gracefulShutdown() {
	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
	defer cancel()

	// 标记所有运行中的任务为暂停状态（立即中断，不等待）
	task.MarkAllTasksPaused()

	// 关闭 HTTP 服务器
	if err := server.Shutdown(ctx); err != nil {
		logger.Error("服务器关闭失败: " + err.Error())
	}

	logger.Info("服务已关闭")
}

// mustInitTables 初始化数据表
func mustInitTables() {
	db := util.MustGetDB()
	defer db.Close()

	if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS "field" (
		"name" TEXT PRIMARY KEY NOT NULL,
		"value" TEXT
	)`); err != nil {
		logger.Fatal("创建 field 表失败: " + err.Error())
	}

	if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS "log" (
		"id" integer NOT NULL PRIMARY KEY AUTOINCREMENT,
		"content" TEXT NOT NULL,
		"create_at" text NOT NULL DEFAULT CURRENT_TIMESTAMP
	)`); err != nil {
		logger.Fatal("创建 log 表失败: " + err.Error())
	}

	if _, err := db.Exec(`CREATE TABLE IF NOT EXISTS "task" (
		"id" integer NOT NULL PRIMARY KEY AUTOINCREMENT,
		"bvid" text NOT NULL,
		"cid" integer NOT NULL,
		"format" integer NOT NULL,
		"title" text NOT NULL,
		"owner" text NOT NULL,
		"cover" text NOT NULL,
		"status" text NOT NULL,
		"folder" text NOT NULL,
		"duration" integer NOT NULL,
		"download_type" text NOT NULL DEFAULT 'merge',
		"create_at" text NOT NULL DEFAULT CURRENT_TIMESTAMP
	)`); err != nil {
		logger.Fatal("创建 task 表失败: " + err.Error())
	}

	if _, err := util.GetCurrentFolder(db); err != nil {
		logger.Fatal("获取下载目录失败: " + err.Error())
	}

	if err := initHistoryTask(db); err != nil {
		logger.Fatal("初始化历史任务失败: " + err.Error())
	}

	// 添加可能缺失的列（用于数据库迁移）
	if err := addMissingColumns(db); err != nil {
		logger.Fatal("添加缺失列失败: " + err.Error())
	}

	logger.Info("数据库初始化完成")
}

// addMissingColumns 添加可能缺失的列（用于数据库迁移）
func addMissingColumns(db *sql.DB) error {
	util.SqliteLock.Lock()
	_, _ = db.Exec(`ALTER TABLE "task" ADD COLUMN "download_type" TEXT DEFAULT 'merge'`)
	_, _ = db.Exec(`UPDATE "task" SET "download_type" = 'merge' WHERE "download_type" IS NULL`)
	util.SqliteLock.Unlock()
	return nil
}

// initHistoryTask 将上一次程序运行时未完成的任务状态变为 paused，支持恢复
func initHistoryTask(db *sql.DB) error {
	util.SqliteLock.Lock()
	// 将 waiting 和 running 状态改为 paused，而非 error
	result, err := db.Exec(`UPDATE "task" SET "status" = 'paused' WHERE "status" IN ('waiting', 'running')`)
	util.SqliteLock.Unlock()
	if err != nil {
		return err
	}
	rowsAffected, _ := result.RowsAffected()
	if rowsAffected > 0 {
		logger.Infof("发现 %d 个未完成任务，状态已设为 paused", rowsAffected)
	}
	return nil
}