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


import cd.casic.ci.process.common.WorkAtom;
import cd.casic.ci.process.constant.CommandConstant;
import cd.casic.ci.process.dto.req.testCase.TestCaseAITaskCreateReq;
import cd.casic.ci.process.dto.resp.resource.ResourceDetailResp;
import cd.casic.ci.process.engine.runContext.TaskRunContext;
import cd.casic.ci.process.engine.worker.base.DockerWorker;
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.dataObject.testCase.TestCaseInfo;
import cd.casic.ci.process.process.service.target.TargetVersionService;
import cd.casic.ci.process.process.service.testCase.TestCaseAIGeneratorService;
import cd.casic.ci.process.process.service.testCase.TestCaseInfoService;
import cd.casic.ci.process.util.SftpUploadUtil;
import cd.casic.module.execute.docker.dataobject.dto.DockerEndpointDo;
import cd.casic.module.machine.dal.dataobject.MachineInfoDO;
import cd.casic.module.machine.utils.CryptogramUtil;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.jcraft.jsch.*;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.*;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Properties;

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

@WorkAtom(taskType = "TEST_CASE_GENERATION")
@Slf4j
public class TestCaseGenerationWorker extends DockerWorker {
    @Resource
    private TestCaseInfoService testCaseInfoService;
    @Resource
    private TargetVersionService targetVersionService;
    @Resource
    private TestCaseAIGeneratorService aiGeneratorService;
    @Override
    public void execute(TaskRunContext context) {
        int statusCode = -1;
        Map<String, Object> globalVariables = context.getGlobalVariables();
        String workDir = globalVariables.get(AFL_DOCKER_WORK_DIR_KEY) instanceof String ? ((String) globalVariables.get(AFL_DOCKER_WORK_DIR_KEY)) : null;
        if (context.getContextDef() instanceof PipTask taskDef) {
            String seedPath= workDir+File.separator+ AFL_DOCKER_SEED;
            Map<String, Object> taskProperties = taskDef.getTaskProperties();
            Object caseType = taskProperties.get(CASE_TYPE_KEY);
            String managerId = getVariableNearby(context, AFL_RESOURCE_MANAGER_ID_KEY) instanceof String ? ((String) getVariableNearby(context, AFL_RESOURCE_MANAGER_ID_KEY)) : null;
            Object itemListObject = taskProperties.get(ITEM_LIST);
            List<String> itemList = JSON.parseArray(JSON.toJSONString(itemListObject),String.class);
            ResourceDetailResp resourceById = resourceManagerService.findResourceDetailById(managerId);
            MachineInfoDO machineInfo = resourceById.getMachineInfo();
            DockerEndpointDo dockerInfo = resourceById.getDockerInfo();
            if (Objects.isNull(machineInfo)||Objects.isNull(dockerInfo)) {
                append(context,"该资源不支持docker或者ssh");
            }
            
            //如果machineId为0，则说明该节点没有配置机器，则使用开始节点的机器
            if (CASE_TYPE_AI.equals(caseType)) {
//                // 容器内部test 目录, 获取编译后的文件
//                String runScript = "docker run -v "+ AFL_VOLUME_WORK_DIR_PREFIX +":"+ AFL_DOCKER_BASE_DIR+"  -it "+imageName+" bash\n" // 目前测试用例ai生成只有这个路径能跑
//                + "cd "+ AFL_DOCKER_BASE_DIR+"\n" +
//                        "mkdir -p "+seedPath +"\n"+
//                        "cd "+seedPath +"\n"+
//                        "PYTHONPATH="+AFL_DOCKER_BASE_DIR+File.separator+"CaseGenerator/src python3 "+AFL_DOCKER_BASE_DIR+File.separator+"CaseGenerator/src/CaseGenerator/main.py --work-dir "+workDir+" --binary "+workDir+File.separator+getSourceName(fileName)+File.separator+binaryName+" --output-dir "+seedPath+" --count 100";
//                    //将节点的配置信息反编译成对象
//                    log.info("测试用例生成-AI生成，实际执行命令:{}" , runScript);
//                    append(context,"测试用例生成-AI生成，实际执行命令:"+runScript);
//                    //获取机器
//                dockerRun(runScript,dockerInfo,context);
                handleAIGenerator(context);
            } else if (CASE_TYPE_FILE.equals(caseType)){
                // 文件上传
                Object filePathObj = taskProperties.get(SEED_SOURCE);
                log.info("当前测试用例taskProperties:{}",taskProperties);
                if (Objects.isNull(filePathObj)) {
                    append(context,"缺少文件");
                    toBadEnding();
                }
                List<String> filePathList = JSON.parseArray(JSON.toJSONString(filePathObj),String.class);
                for (String filePath : filePathList) {
                    File file = new File(filePath);
                    String seedTarget = AFL_VOLUME_WORK_DIR_PREFIX+File.separator+seedPath;
                    // 将文件上传到服务器的 目录底下
                    log.info("测试用例选用上传模式,种子文件路径：{},种子上传路径:{}",filePath,seedTarget);
                    append(context,"测试用例选用上传模式,种子文件路径："+filePath+",种子上传路径:"+seedPath);
                    try {
                        SftpUploadUtil.uploadFileViaSftp(
                                machineInfo.getHostIp()
                                ,machineInfo.getSshPort(),machineInfo.getUsername(), CryptogramUtil.doDecrypt(machineInfo.getPassword()),"",filePath,seedTarget,file.getName());
                    } catch (SftpUploadUtil.SftpUploadException e) {
                        append(context,"seed文件上传失败");
                        log.error("seed文件上传失败",e);
                    }
                    append(context,file.getName()+"上传文件成功！");
                }
            }else if (CASE_TYPE_ITEM.equals(caseType)){
                List<TestCaseInfo> testCaseInfos = testCaseInfoService.listByIds(itemList);
                String seedTarget = AFL_VOLUME_WORK_DIR_PREFIX+File.separator+seedPath;
                for (TestCaseInfo testCaseInfo : testCaseInfos) {
                    String testFilePath = testCaseInfo.getTestFilePath();
                    if (StringUtils.isEmpty(testFilePath)) {
                        continue;
                    }
                    try {
                        append(context,"当前上传文件:"+testCaseInfo.getFileName());
                        SftpUploadUtil.uploadFileViaSftp(
                                machineInfo.getHostIp()
                                ,machineInfo.getSshPort()
                                ,machineInfo.getUsername()
                                ,CryptogramUtil.doDecrypt(machineInfo.getPassword())
                                ,"", testFilePath,seedTarget,testCaseInfo.getFileName());
                    } catch (SftpUploadUtil.SftpUploadException e) {
                        toBadEnding();
                    }
                }
            }
            globalVariables.put(AFL_DOCKER_SEED_PATH_KEY,seedPath);
        }
    }
    public void handleAIGenerator(TaskRunContext context){
        PipTask taskDef = context.getContextDef() instanceof PipTask ? ((PipTask) context.getContextDef()) : null;
        Map<String, Object> taskProperties = taskDef.getTaskProperties();
        String binaryName = taskProperties.get(BINARY_NAME) instanceof String ? ((String) taskProperties.get("binaryName")) : null;
        Integer count = taskProperties.get(COUNT) instanceof Integer ? ((Integer) taskProperties.get(COUNT)) : null;
        String prompt = taskProperties.get(PROMPT) instanceof String ? ((String) taskProperties.get(PROMPT)) : null;
        PipPipeline pipeline = (PipPipeline) getContextManager().getContext(taskDef.getPipelineId()).getContextDef();
        TargetVersion targetVersion = targetVersionService.getById(pipeline.getTargetVersionId());
        String managerId = getVariableNearby(context, AFL_RESOURCE_MANAGER_ID_KEY) instanceof String ? ((String) getVariableNearby(context, AFL_RESOURCE_MANAGER_ID_KEY)) : null;
        ResourceDetailResp resourceById = resourceManagerService.findResourceDetailById(managerId);
        MachineInfoDO machineInfo = resourceById.getMachineInfo();
        DockerEndpointDo dockerInfo = resourceById.getDockerInfo();
        String fileName = targetVersion.getFileName();
        if (Objects.isNull(machineInfo)||Objects.isNull(dockerInfo)) {
            append(context,"该资源不支持docker或者ssh");
        }
        // 获取编译后的文文件的inputStream
        JSch jsch = new JSch();
        Session session = null;
        ChannelExec exec = null;
        ChannelSftp sftp = null;
        try {
            session = jsch.getSession(machineInfo.getUsername(),machineInfo.getHostIp(), machineInfo.getSshPort());
            session.setPassword( CryptogramUtil.doDecrypt(machineInfo.getPassword()));
            // 跳过 host key 检查
            Properties config = new Properties();
            config.put("StrictHostKeyChecking", "no");
            session.setConfig(config);


            session.connect();
            sftp = (ChannelSftp) session.openChannel("sftp");
            sftp.connect();
            String workDir = (String)getGlobalVariable(context,AFL_DOCKER_WORK_DIR_KEY);
            append(context,"AI测试用例生成开始上传目标文件。。。");
            String filePath = AFL_VOLUME_WORK_DIR_PREFIX + File.separator + workDir+File.separator+getSourceName(fileName)+File.separator+binaryName;
            append(context,"读取到目标文件路径："+filePath);
            InputStream inputStream = sftp.get(filePath);

            String targetUploadRes = aiGeneratorService.targetUpload("./" + binaryName, inputStream);
            append(context,"上传结果:"+targetUploadRes);
            String binaryId = JSON.parseObject(targetUploadRes).getString("id");
            append(context,"开始创建测试用例生成task");
            sftp.disconnect();

            // 创建task
            TestCaseAITaskCreateReq req = new TestCaseAITaskCreateReq();
            req.setBinaryId(binaryId);
            req.setCustomPrompt(prompt);
            req.setCount(count);
            String taskCreateRes = aiGeneratorService.taskCreate(req);
            String taskId = JSON.parseObject(taskCreateRes).getString("task_id");
            append(context,"创建测试用例生成task成功，开始读取task执行日志");
            while (true){
                // 获取状态及日志
                String taskStatusRes = aiGeneratorService.getTaskStatus(taskId);
                String status = JSONObject.parseObject(taskStatusRes).getString("status");
                if ("success".equals(status)) {
                    break;
                } else if("failed".equals(status)){
                    toBadEnding();
                } else {
                    String logRes = aiGeneratorService.getTaskLog(taskId);
                    String logStr = JSON.parseObject(logRes).getString("log");
                    if (StringUtils.isNotEmpty(logStr)) {
                        append(context,logStr);
                    }
                }
//                append(context,"task当前执行状态"+taskStatusRes);
            }
            append(context,"开始下载生成的测试用例文件");
            File tempFile = File.createTempFile("seed", ".zip");
            FileOutputStream fos = new FileOutputStream(tempFile);
            // 下载生成的种子
            aiGeneratorService.downloadTestCase(taskId,fos);
            fos.close();
            append(context,"下载生成的模糊测试种子文件成功");

            FileInputStream fis = new FileInputStream(tempFile);
            // 上传种子zip到目标服务器
            String seedPath= workDir+File.separator+ AFL_DOCKER_SEED;
            String seedTarget = AFL_VOLUME_WORK_DIR_PREFIX+File.separator+seedPath+File.separator+tempFile.getName();
            String seedDir = AFL_VOLUME_WORK_DIR_PREFIX+File.separator+seedPath;
            append(context,"上传模糊测试种子压缩包到指定工作目录:"+seedDir);
            String cmd = "mkdir -p "+seedDir;
            append(context,cmd);
            // 解压种子zip
            exec = (ChannelExec)session.openChannel("exec");
            exec.setCommand(cmd);
            exec.connect();
            exec.disconnect();
            sftp = (ChannelSftp) session.openChannel("sftp");
            sftp.connect();
            sftp.put(fis,seedTarget);
            sftp.disconnect();
            fis.close();
            append(context,"解压种子到指定工作目录"+seedDir);

            cmd = "unzip -o " + seedTarget + " -d " + seedDir;
            append(context,cmd);
            // 解压种子zip
            exec = (ChannelExec)session.openChannel("exec");
            exec.setCommand(cmd);
            exec.connect();
            InputStream errStream = exec.getErrStream();
            InputStream inStream = exec.getInputStream();
            BufferedReader reader = new BufferedReader(new InputStreamReader(inStream));
            BufferedReader errReader = new BufferedReader(new InputStreamReader(errStream));
            String line;
            while ((line = reader.readLine()) != null) {
                append(context, "STDOUT: " + line);
            }
            while ((line = errReader.readLine()) != null) {
                append(context, "STDERR: " + line);
            }
            tempFile.deleteOnExit();
        } catch (Exception e) {
            log.error("ai测试用例生成报错",e);
            append(context,"测试用例生成失败");
        }finally {
            if (exec!=null) {
                if (exec.isConnected()) {
                    exec.disconnect();
                }
            }
            if (sftp!=null) {
                if (sftp.isConnected()) {
                    sftp.disconnect();
                }
            }
            if (session!=null) {
                if (session.isConnected()) {
                    session.disconnect();
                }
            }
        }

    }
    private String getSourceName(String fileName){
        if (fileName.contains(".zip")||fileName.contains(".tar")||fileName.contains(".tar.gz")) {
            return fileName
                    .replace(".zip","")
                    .replace(".tar.gz","")
                    .replace(".tar","");
        }
        return "";
    }
}
