package com.backendsys.modules.sdk.comfyui.utils;

import com.backendsys.modules.common.config.security.utils.SecurityUtil;
import com.backendsys.modules.sdk.comfyui.dao.ComfyuiTaskDao;
import com.backendsys.modules.sdk.comfyui.dao.ComfyuiTaskExecuteDao;
import com.backendsys.modules.sdk.comfyui.entity.ComfyuiTask;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

@Component
public class ComfyuiUtil {

    @Autowired
    private ComfyuiTaskDao comfyuiTaskDao;
    
    @Autowired
    private ComfyuiTaskExecuteDao comfyuiTaskExecuteDao;

    /**
     * 初始化任务
     */
    public Long initComfyuiTask(String client_id, String prompt, String type) {
        Long user_id = SecurityUtil.getUserId();
        try {
            ComfyuiTask comfyuiTask = new ComfyuiTask();
            comfyuiTask.setUser_id(user_id);
            comfyuiTask.setClient_id(client_id);
            comfyuiTask.setTask_type(type);
            comfyuiTask.setGenerate_request(prompt);
            comfyuiTaskDao.insert(comfyuiTask);
            return comfyuiTask.getId();
        } catch (Exception e) {
            System.out.println(e.getMessage());
        }
        return -1L;
    }
    /**
     * [更新] 执行任务
     * - status: 任务状态 (-1:未开始, 1:进行中, 2:成功, 3:失败)
     */
    public void updateComfyuiTask(String prompt_id, String response, Integer status) {
        try {
            LambdaQueryWrapper<ComfyuiTask> wrapper = new LambdaQueryWrapper<>();
            wrapper.eq(ComfyuiTask::getExecute_prompt_id, prompt_id);

            ComfyuiTask comfyuiTask = new ComfyuiTask();
            comfyuiTask.setExecute_status(status);
            comfyuiTask.setExecute_prompt_id(prompt_id);
            comfyuiTask.setGenerate_response(response);

            comfyuiTaskExecuteDao.update(comfyuiTask, wrapper);
        } catch (Exception e) {
            System.out.println(e.getMessage());
        }
    }

}
