package cd.casic.ci.process.engine.worker.afl;

import cd.casic.ci.process.common.WorkAtom;
import cd.casic.ci.process.dto.req.aflManager.AflManagerReq;
import cd.casic.ci.process.dto.req.localAflCrashes.LocalAflCrashesSaveReq;
import cd.casic.ci.process.dto.resp.resource.ResourceDetailResp;
import cd.casic.ci.process.engine.constant.DIYImageExecuteCommandConstant;
import cd.casic.ci.process.engine.postHandler.ExecuteTaskPostHandler;
import cd.casic.ci.process.engine.runContext.TaskRunContext;
import cd.casic.ci.process.engine.worker.base.DockerWorker;
import cd.casic.ci.process.process.dataObject.history.PipPipelineHisInstance;
import cd.casic.ci.process.process.dataObject.pipeline.PipPipeline;
import cd.casic.ci.process.process.dataObject.target.TargetVersion;
import cd.casic.ci.process.process.dataObject.task.PipTask;
import cd.casic.ci.process.process.service.aflAnalyzer.AflCrashesAnalyzerService;
import cd.casic.ci.process.process.service.aflManager.AflInfoService;
import cd.casic.ci.process.process.service.aflManager.AflPlotInfoService;
import cd.casic.ci.process.process.service.aflManager.AflSeedInfoService;
import cd.casic.ci.process.process.service.localAflCrashes.LocalAflCrashesService;
import cd.casic.ci.process.process.service.target.TargetVersionService;
import cd.casic.ci.process.util.SftpUploadUtil;
import com.alibaba.fastjson.JSON;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.File;
import java.util.Map;
import java.util.Objects;
import java.util.UUID;

import static cd.casic.ci.process.engine.constant.AFLConstant.*;
import static cd.casic.ci.process.engine.constant.PipelineVariableConstant.*;

@WorkAtom(taskType = "AFL")
@Slf4j
public class AFLWorker extends DockerWorker {
    @Resource
    private TargetVersionService targetVersionService;
    @Resource
    private AflInfoService aflInfoService;
    @Resource
    private AflSeedInfoService aflSeedInfoService;
    @Resource
    private AflPlotInfoService aflPlotInfoService;
    @Resource
    private LocalAflCrashesService localAflCrashesService;
    @Resource
    private AflCrashesAnalyzerService crashesAnalyzerService;
    @Override
    public void execute(TaskRunContext context) {
        int statusCode = -1;
        Map<String, Object> localVariables = context.getLocalVariables();
        if (context.getContextDef() instanceof PipTask taskDef) {
            log.info(taskDef.getTaskName());
            Map<String, Object> taskProperties = taskDef.getTaskProperties();
            //从上下文中获取资源id
            String managerId = getVariableNearby(context, AFL_RESOURCE_MANAGER_ID_KEY) instanceof String ? ((String) getVariableNearby(context, AFL_RESOURCE_MANAGER_ID_KEY)) : null;
            // ssh 上传目标文件
            ResourceDetailResp resourceById = resourceManagerService.findResourceDetailById(managerId);
            if (resourceById==null) {
                append(context,"缺少资源信息，请添加编译节点");
            }
            String commandEnd = taskProperties.get(COMMAND_END) instanceof String ? ((String) taskProperties.get(COMMAND_END)) : null;
            String imageName = getVariableNearby(context,IMAGE_NAME) instanceof String ? ((String) getVariableNearby(context,IMAGE_NAME)) : null;
            String executableName = taskProperties.get(EXECUTABLE_NAME) instanceof String ? ((String) taskProperties.get(EXECUTABLE_NAME)) : null;
            putGlobalVariable(context, EXECUTABLE_NAME_KEY,imageName);
            // 代码编译目录
            String workDir = (String)getGlobalVariable(context,AFL_DOCKER_WORK_DIR_KEY);
            // 入参目录
            String seedPath = (String)getGlobalVariable(context,AFL_DOCKER_SEED_PATH_KEY);
            if (StringUtils.isEmpty(seedPath) ||
                    StringUtils.isEmpty(workDir) ||
                    Objects.isNull(resourceById)||
                    StringUtils.isEmpty(imageName) ||
                    StringUtils.isEmpty(executableName) ||
                    StringUtils.isEmpty(commandEnd) ||
                    Objects.isNull(resourceById.getDockerInfo())) {
//                缺少参数
                toBadEnding();
            }
            String output = workDir +File.separator+ AFL_DOCKER_OUTPUT;
            String volumeWorkDirPath = AFL_VOLUME_WORK_DIR_PREFIX;
            PipPipeline pipeline = (PipPipeline) getContextManager().getContext(taskDef.getPipelineId()).getContextDef();
            TargetVersion targetVersion = targetVersionService.getById(pipeline.getTargetVersionId());
            String fileName = targetVersion.getFileName();
            append(context,"AFL 输出路径："+output);
            Object timeObj = taskProperties.get(RUNNING_TIME);
            Double timeHours = Double.valueOf(String.valueOf(timeObj));
            String outputPath = AFL_DOCKER_BASE_DIR + File.separator + output;
            String allDir = AFL_DOCKER_BASE_DIR + File.separator + workDir;
            String initScript = "docker run -v " + volumeWorkDirPath + ":" + AFL_DOCKER_BASE_DIR + "  -it " + imageName + " bash\n" +
                    "cd " + allDir + "\n";
            String commandScript = initScript +
                    cdSourceName(fileName) +
                    "rm -rf " + outputPath + "\n" +
                    "mkdir -p "+ outputPath +"\n" +
                    "chmod o+rwx -R "+outputPath+"\n"+
                    "afl-fuzz -i "+ AFL_DOCKER_BASE_DIR+File.separator+seedPath+" -o "+ outputPath+"  ./"+ executableName + " " +commandEnd;
            String aflCrashScript= initScript+" python "+allDir+"/afl-crash-analyzer-m/AflCrashAnalyzer.py -i "+outputPath+"/default/crashes/  -o "+allDir+"/aflCrash/ -b "+allDir+getSourcePath(fileName)+File.separator+executableName;
            try {
                //将节点的配置信息反编译成对象
                log.info("AFL模糊测试执行脚本：{}",commandScript);
                append(context,"AFL模糊测试执行脚本："+commandScript);
                Long runningTime = getRunningTime(timeHours);
                append(context,"运行时间"+runningTime);
                // 获取docker 暂时先写固定值
                dockerRun(commandScript,resourceById.getDockerInfo(),context, runningTime);
                log.info("AFLCrash文件分析命令：{}",commandScript);
                append(context,"AFLCrash文件分析脚本："+commandScript);
                dockerRun(aflCrashScript,resourceById.getDockerInfo(),context);
            } catch (Exception e) {
                String errorMessage = "执行afl失败"+e.getMessage() + "\r\n";
                log.error("执行afl失败", e);
                append(context, errorMessage);
                toBadEnding();
            }
            localVariables.put(DIYImageExecuteCommandConstant.STATUS_CODE, statusCode);
            afterTaskExecute(context);
            afterPipelineExecute(context,pipeline,resourceById,AFL_VOLUME_WORK_DIR_PREFIX+File.separator + workDir+"/aflCrash");
        }
    }
    private String cdSourceName(String fileName){
        if (fileName.contains(".zip")||fileName.contains(".tar")||fileName.contains(".tar.gz")) {
            String cdPath = fileName
                    .replace(".zip","")
                    .replace(".tar.gz","")
                    .replace(".tar","");
            return "cd "+cdPath+"\n";
        }
        return "";
    }
    private String getSourcePath(String fileName){
        if (fileName.contains(".zip")||fileName.contains(".tar")||fileName.contains(".tar.gz")) {
            return fileName
                    .replace(".zip","")
                    .replace(".tar.gz","")
                    .replace(".tar","");
        }
        return "";
    }
    private Long getRunningTime(Double runningTimeHours){
        if (runningTimeHours!=null && runningTimeHours >0) {
            return (long)(runningTimeHours*60*60*1000);
        }
        return null;
    }

    public void afterTaskExecute(TaskRunContext context){
        append(context,"开始拉取afl输出信息");
        PipTask contextDef = (PipTask)context.getContextDef();
        Map<String, Object> localVariables = context.getLocalVariables();
        String pipelineId =contextDef.getPipelineId();
        String taskId = contextDef.getId();
        String groupId = UUID.randomUUID().toString().replace("-","");
        localVariables.put(GROUP_ID,groupId);
        AflManagerReq req = new AflManagerReq();
        req.setGroupIdentifier(groupId);
        req.setTaskId(taskId);
        req.setPipelineId(pipelineId);
        try {
            aflInfoService.saveAflInfo(req);
        } catch (SftpUploadUtil.SftpUploadException e) {
            log.error("保存afl信息失败,入参{}", JSON.toJSONString(req),e);
        }
//        aflSeedInfoService.saveAflSeedInfo(req);
        try {
            aflPlotInfoService.saveAflPlotInfo(req);
        } catch (SftpUploadUtil.SftpUploadException e) {
            log.error("保存afl信息失败,入参{}", JSON.toJSONString(req),e);
        }
    }
    private void afterPipelineExecute(TaskRunContext context,PipPipeline pipeline,ResourceDetailResp resource,String crashesOutputPath){
        append(context,"开始拉取afl输出信息");
        PipTask contextDef = (PipTask)context.getContextDef();
        String pipelineId =contextDef.getPipelineId();
        String taskId = contextDef.getId();
        AflManagerReq req = new AflManagerReq();
        req.setTaskId(taskId);
        req.setPipelineId(pipelineId);
        TargetVersion targetVersion = targetVersionService.getById(pipeline.getTargetVersionId());
        postHandlerManager.registerPostHandler(new ExecuteTaskPostHandler(taskId,pipelineId) {
            @Override
            public void executeAfterDone(PipPipelineHisInstance pipPipelineHisInstance) {
                Map<String, Object> localVariables = context.getLocalVariables();
                String groupId = localVariables.get(GROUP_ID) instanceof String ? ((String) localVariables.get(GROUP_ID)) : null;
                req.setGroupIdentifier(groupId);
                req.setPipelineHistoryId(pipPipelineHisInstance.getId());
                aflInfoService.updateHistoryPipelineIdByAflInfo(req);
                aflSeedInfoService.updateHistoryPipelineIdByAflSeedInfo(req);
                aflPlotInfoService.updateHistoryPipelineIdByAflPlotInfo(req);
                LocalAflCrashesSaveReq crashesSaveReq = new LocalAflCrashesSaveReq();
                crashesSaveReq.setPipelineId(pipelineId);
                crashesSaveReq.setInstanceId(pipPipelineHisInstance.getId());
                crashesSaveReq.setTaskId(taskId);
                crashesSaveReq.setTaskType("AFL");
                crashesSaveReq.setCity("长沙市");// 暂时空置
                crashesSaveReq.setTargetType(pipeline.getTargetType());
                crashesSaveReq.setTargetName(targetVersion.getFileName());
                localAflCrashesService.saveLocalAflCrashesInfo(crashesSaveReq);
            }
        });
        postHandlerManager.registerPostHandler(new ExecuteTaskPostHandler(taskId,pipelineId) {
            @Override
            public void executeAfterDone(PipPipelineHisInstance instance) {
                crashesAnalyzerService.saveCrashesByOutputPath(crashesOutputPath,resource,instance.getPipelineId(),taskId,instance.getId());
            }
        });
    }
}
