EasyVQD/internal/core/vqdtask/core.go

package vqdtask

import (
	"bufio"
	"context"
	"easyvqd/internal/conf"
	"easyvqd/internal/core/host"
	"easyvqd/internal/core/vqd"
	"easyvqd/pkg/vqdcms"
	"fmt"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

type Core struct {
	HostCore    *host.Core
	VqdTaskCore *vqd.Core
	Cfg         *conf.Bootstrap
	ResultCb    vqdcms.VQDResultCB
}

var (
	VqdTaskMap = vqdcms.VqdTaskMap{M: make(map[string]*vqdcms.VQDHandle)}
)

func NewCore(HostCore *host.Core, VqdTaskCore *vqd.Core, Cfg *conf.Bootstrap) *Core {
	core := &Core{
		HostCore:    HostCore,
		VqdTaskCore: VqdTaskCore,
		Cfg:         Cfg,
	}
	core.ResultCb = func(v vqdcms.AbnormalModel) {

		in := &vqd.AddVqdAlarmInput{
			ChannelName:      v.ChannelName,
			ChannelID:        v.ChannelID,
			TaskTemplateName: v.TemplateName,
			TaskName:         v.TaskName,
			PlanName:         v.PlanName,
			TaskID:           int64(v.ID),
			TaskTemplateID:   int64(v.TemplateID),
			PlanID:           int64(v.PlanID),
			IsDeep:           v.IsDeep,
			FilePath:         v.FilePath,
		}
		var Abnormals vqd.Abnormals
		var AbnormalStr []string
		if len(v.Abnormals) > 0 {
			for _, abnormal := range v.Abnormals {
				item := vqd.Abnormal{
					Value: abnormal.Value,
					Name:  abnormal.Name,
					Mode:  abnormal.Mode,
				}
				AbnormalStr = append(AbnormalStr, abnormal.Mode)
				Abnormals = append(Abnormals, item)
			}
		}
		in.AbnormalStr = strings.Join(AbnormalStr, ",")
		in.Abnormals = Abnormals

		var DefaultValues vqd.DefaultValues
		if len(v.DefaultValues) > 0 {
			for _, defaultValue := range v.DefaultValues {
				item := vqd.DefaultValue{
					Thr1:  defaultValue.Thr1,
					Name1: defaultValue.Name1,
					Thr2:  defaultValue.Thr2,
					Name2: defaultValue.Name2,
					Ratio: defaultValue.Ratio,
				}
				DefaultValues = append(DefaultValues, item)
			}
		}
		in.DefaultValues = DefaultValues
		_, err := core.VqdTaskCore.AddVqdAlarm(context.TODO(), in)
		if err != nil {
			slog.Error("add alarm", "err", err.Error())
		}
	}
	core.HostCore.CbIFrame = func(s string, data []byte, codes int) {
		if codes == VIDEO_CODEC_H264 {
			{
				v, ok := VqdTaskMap.LoadTaskMap(s)
				if ok {
					v.SendData(data, VIDEO_CODEC_H264)
				}
			}
			{
				v, ok := VqdPollingTaskMap.LoadTaskMap(fmt.Sprintf("%s_polling", s))
				if ok {
					v.SendData(data, VIDEO_CODEC_H264)
				}
			}
		} else {
			{
				v, ok := VqdTaskMap.LoadTaskMap(s)
				if ok {
					v.SendData(data, VIDEO_CODEC_H265)
				}
			}
			{
				v, ok := VqdPollingTaskMap.LoadTaskMap(fmt.Sprintf("%s_polling", s))
				if ok {
					v.SendData(data, VIDEO_CODEC_H265)
				}
			}
		}

		//slog.Debug("cb IFrame", "name", s, "codes", codes)

	}
	time.AfterFunc(time.Duration(5)*time.Second, func() {
		// 启用诊断分析
		core.InitVqdTask()
		// 启用轮询诊断任务
		core.StartPolling()
	})
	// 启用定时清理任务
	go core.scheduleCleanTask()

	// 启用任务管理器
	return core
}

func (c *Core) InitVqdTask() {
	err := vqdcms.VQDInit()
	if err != nil {
		slog.Error("vqd cms open", "err", err.Error())
		return
	}

	all, _, err := c.VqdTaskCore.FindVqdTaskAll()
	if err == nil {
		for _, vqdTask := range all {
			if vqdTask.Enable {
				errs := c.AddTaskVqd(vqdTask.ID)
				if errs != nil {
					slog.Error("vqd init add task", "err", errs.Error())
				}
				time.Sleep(200 * time.Millisecond)
			}
		}
	}

	return
}
func (c *Core) UnVqdTask() {
	VqdTaskMap.DeleteTaskMapAll()
	vqdcms.VQDUnInit()
	// 停止轮询诊断任务
	c.StopPolling()
	return
}
func (c *Core) AddTaskVqd(taskId int) error {

	task, err := c.VqdTaskCore.GetVqdTask(context.TODO(), taskId)
	if err != nil {
		slog.Error("vqd add task find", "err", err.Error())
		return err
	}

	taskTemplate, err := c.VqdTaskCore.GetIDVqdTaskTemplate(context.TODO(), task.TaskTemplateID)
	if err != nil {
		slog.Error("vqd add task find template", "err", err.Error())
		return err
	}
	taskPlan, err := c.VqdTaskCore.GetVqdTimeTemplate(context.TODO(), int(task.TimeTemplateID))
	if err != nil {
		slog.Error("vqd add task find plan", "err", err.Error())
		return err
	}
	chnId := task.ChannelID

	para := vqdcms.NewVQDPara(taskTemplate)
	info := vqdcms.VQDHandleInfo{
		ChannelID:    chnId,
		ChannelName:  task.ChannelName,
		TaskID:       task.ID,
		TaskName:     task.Name,
		TemplateID:   taskTemplate.ID,
		TemplateName: taskTemplate.Name,
		PlanID:       taskPlan.ID,
		PlanName:     taskPlan.Name,
		Plans:        taskPlan.Plans,
	}
	v := vqdcms.NewVQDHandle(c.ResultCb, c.HostCore, info).Create(para, taskPlan.Plans)

	VqdTaskMap.StoreChildMap(fmt.Sprintf("%s", chnId), v)
	return nil
}
func (c *Core) UpdateTaskVqd(taskId int) error {
	task, err := c.VqdTaskCore.GetVqdTask(context.TODO(), taskId)
	if err != nil {
		slog.Error("vqd update task find", "err", err.Error())
		return err
	}
	v, ok := VqdTaskMap.LoadTaskMap(task.ChannelID)
	if ok {
		taskTemplate, err := c.VqdTaskCore.GetIDVqdTaskTemplate(context.TODO(), task.TaskTemplateID)
		if err != nil {
			slog.Error("vqd update task find template", "err", err.Error())
			return err
		}
		taskPlan, err := c.VqdTaskCore.GetVqdTimeTemplate(context.TODO(), int(task.TimeTemplateID))
		if err != nil {
			slog.Error("vqd add task find plan", "err", err.Error())
			return err
		}
		para := vqdcms.NewVQDPara(taskTemplate)
		info := vqdcms.VQDHandleInfo{
			ChannelID:    task.ChannelID,
			ChannelName:  task.ChannelName,
			TaskID:       task.ID,
			TaskName:     task.Name,
			TemplateID:   taskTemplate.ID,
			TemplateName: taskTemplate.Name,
			PlanID:       taskPlan.ID,
			PlanName:     taskPlan.Name,
			Plans:        taskPlan.Plans,
		}
		errs := v.SetVQDConfig(para, info)
		if errs != nil {
			slog.Error("vqd update set config err", "err", errs.Error())
			return err
		}
		VqdTaskMap.StoreChildMap(fmt.Sprintf("%s", task.ChannelID), v)
	}
	return nil

}
func (c *Core) DelTaskVqd(taskId int, chnId string) {
	v, ok := VqdTaskMap.LoadTaskMap(chnId)
	if ok {
		v.Destroy()
	}
	VqdTaskMap.DeleteTaskMap(chnId)
}

// 测试i帧数据是否可以转为图片文件
func CheckIFramesToJpg() {
	dir, _ := os.Getwd()

	h264Path := filepath.Join(dir, "check_tip.h264") // 你的H.264裸流文件路径
	//1. 检查FFmpeg
	ok, err := CheckFFmpeg()
	if !ok {
		fmt.Println("错误:", err)
		os.Exit(1)
	}

	// 2. 配置参数（根据你的H.264文件调整）

	imageFormat := "jpg" // 输出图片格式
	videoWidth := 0      // 分辨率自动探测（如需指定则改为实际值，如1920）
	videoHeight := 0     // 分辨率自动探测（如需指定则改为实际值，如1080）

	// 3. 提取H.264关键帧并转图片
	if err := ExtractH264KeyFrames(h264Path, dir, imageFormat, videoWidth, videoHeight); err != nil {
		fmt.Println("提取H.264关键帧失败:", err)
		os.Exit(1)
	}
}

// CheckFFmpeg 检查系统是否安装了FFmpeg
func CheckFFmpeg() (bool, error) {
	cmd := exec.Command("ffmpeg", "-version")
	err := cmd.Run()
	if err != nil {
		return false, fmt.Errorf("FFmpeg 未安装或未添加到系统PATH: %v", err)
	}
	return true, nil
}

// ExtractH264KeyFrames 将H.264裸流的关键帧转为图片
// h264Path: H.264裸流文件路径（.264/.h264）
// outputDir: 图片输出目录
// format: 输出图片格式 (jpg/png)
// width/height: 视频分辨率（若不指定，FFmpeg会自动探测）
func ExtractH264KeyFrames(h264Path, outputDir, format string, width, height int) error {
	// 验证输入文件
	if _, err := os.Stat(h264Path); os.IsNotExist(err) {
		return fmt.Errorf("H.264文件不存在: %s", h264Path)
	}

	// 创建输出目录
	if err := os.MkdirAll(outputDir, 0755); err != nil {
		return fmt.Errorf("创建输出目录失败: %v", err)
	}

	// 构造FFmpeg命令（适配H.264裸流）
	outputPattern := filepath.Join(outputDir, "h264_keyframe_%04d."+format)
	args := []string{
		"-f", "h264", // 明确输入格式为H.264裸流
		"-i", h264Path, // 输入H.264裸流文件
		"-skip_frame", "nokey", // 只处理关键帧（I帧）
		"-vsync", "0", // 禁用帧同步，保证每个I帧都输出
		"-q:v", "2", // 图片质量（1-31，值越小质量越高）
	}

	// 可选：指定分辨率（如果FFmpeg自动探测失败时使用）
	if width > 0 && height > 0 {
		args = append(args, "-s", fmt.Sprintf("%dx%d", width, height))
	}

	// 补全输出参数
	args = append(args,
		"-f", "image2", // 输出图片序列格式
		outputPattern, // 输出文件模板
		"-y",          // 覆盖已存在的文件
	)

	// 创建命令并捕获输出
	cmd := exec.Command("ffmpeg", args...)
	_, err := cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("创建标准输出管道失败: %v", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("创建标准错误管道失败: %v", err)
	}

	// 启动命令
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("启动FFmpeg失败: %v", err)
	}

	// 实时打印FFmpeg日志（便于调试H.264解析问题）
	scanner := bufio.NewScanner(stderr)
	for scanner.Scan() {
		line := scanner.Text()
		fmt.Println("FFmpeg日志:", line)
		// 捕获关键错误信息
		if strings.Contains(line, "error") && strings.Contains(line, "H.264") {
			fmt.Println("⚠️ H.264解析警告:", line)
		}
	}

	// 等待命令执行完成
	if err := cmd.Wait(); err != nil {
		return fmt.Errorf("FFmpeg执行失败: %v", err)
	}

	// 统计输出的关键帧图片数量
	files, err := filepath.Glob(filepath.Join(outputDir, "h264_keyframe_*."+format))
	if err != nil {
		return fmt.Errorf("统计输出图片失败: %v", err)
	}

	fmt.Printf("成功从H.264裸流提取 %d 个关键帧，保存至: %s\n", len(files), outputDir)
	return nil
}