package cd.casic.server;

import cd.casic.module.execute.docker.dataobject.model.DockerEndpoint;
import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.async.ResultCallbackTemplate;
import com.github.dockerjava.api.command.CreateContainerResponse;
import com.github.dockerjava.api.command.ExecCreateCmdResponse;
import com.github.dockerjava.api.model.Bind;
import com.github.dockerjava.api.model.Frame;
import com.github.dockerjava.api.model.HostConfig;
import com.github.dockerjava.core.DockerClientBuilder;
import com.github.dockerjava.httpclient5.ApacheDockerHttpClient;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.assertj.core.util.Lists;
import org.junit.jupiter.api.Test;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.ActiveProfiles;
import org.springframework.util.CollectionUtils;

import java.io.Closeable;
import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.TimeUnit;

import static java.lang.String.format;

@SpringBootTest(classes = {OpsServerApplication.class})
@ActiveProfiles("local")
@Slf4j
public class DockerRunTest {
    @Test
    public void test01(){
        DockerEndpoint dockerEndpoint = new DockerEndpoint();
        dockerEndpoint.setHost("175.6.27.228");
        dockerEndpoint.setPort(22375);
        dockerEndpoint.setType(DockerEndpoint.DockerEndpointTypeEnum.REMOTE);
        String command = "docker run -v /home/casic/706/yunqi/:/test  -it aflplusplus/aflplusplus bash\n" +
                "cd /test\n" +
                "afl-fuzz -i case -o ai_afl -t 3000 -Q ./CaseGenerator/testdata/libpng/libpng/pngfix @@";
        dockerRun(command,dockerEndpoint);
    }
    public void dockerRun(String command, DockerEndpoint dockerEndpoint){
        // 第一行必须是docker run 命令 option 支持 -it  -v
        String[] split = command.split("\n");
        List<String> commandLine = new ArrayList<>(Arrays.stream(split).filter(StringUtils::isNotBlank).toList());
        Optional<String> first = commandLine.stream().filter(StringUtils::isNotBlank).findFirst();
        String dockerRunCommand="";
        // 找到第一个不为空的判断是否为docker run
        if (first.isPresent()) {
            dockerRunCommand = first.get();
        }
        if (StringUtils.isEmpty(dockerRunCommand)) {
            log.error("未找到有效指令");
            return;
        }
        DockerExecHandler handler = null;
        try {
            handler = loadRunCommand(dockerRunCommand, dockerEndpoint);
            // 删除第一个 dockerRun 其他都要在容器内部执行
            commandLine.remove(0);
            if (handler == null) {
                log.error("容器创建失败");
            }
            if (!CollectionUtils.isEmpty(commandLine)) {
                StringBuffer stringBuffer = new StringBuffer();
                // docker连接
                DockerClient client = handler.getClient();
                for (String cmd : commandLine) {
                    stringBuffer.append(cmd);
                    stringBuffer.append("&&");
                }
                stringBuffer.setLength(stringBuffer.length()-2);
                // 启动的容器id
                String containerId = handler.getContainerId();
                String allCommand = stringBuffer.toString();
                String[] commandList = new String[3];
                commandList[0]="sh";
                commandList[1]="-c";
                commandList[2]=allCommand;
                log.info("容器内执行命令：{}",commandList);
                // 创建命令
                ExecCreateCmdResponse exec = client.execCreateCmd(containerId)
                        .withAttachStdin(true)
                        .withAttachStdout(true)
                        .withAttachStderr(true)
                        .withCmd(commandList).exec();

                try {
                    client.execStartCmd(exec.getId()).exec(new ResultCallbackTemplate<>() {
                        @Override
                        public void onStart(Closeable stream) {
                            super.onStart(stream);
                            log.info("命令执行开始,容器id：{},执行id：{}",containerId,exec.getId());
                        }

                        @Override
                        public void onError(Throwable throwable) {
                            super.onError(throwable);
                            log.error("命令执行出现错误,容器id：{},执行id：{},错误原因",containerId,exec.getId(),throwable);
                        }

                        @Override
                        public void onComplete() {
                            super.onComplete();
                            log.info("命令执行完毕,容器id：{},执行id：{}",containerId,exec.getId());
                        }

                        @Override
                        public void onNext(Frame frame) {
                            String output = new String(frame.getPayload(), StandardCharsets.UTF_8);
                            switch (frame.getStreamType()) {
                                case STDOUT:
                                    System.out.print(output);
                                    log.info("标准输出: {}", output.trim());
                                    break;
                                case STDERR:
                                    System.err.print(output);
                                    log.error("错误输出: {}", output.trim());
                                    break;
                                default:
                                    log.warn("未知流类型: {}", frame.getStreamType());
                            }
                        }
                    }).awaitCompletion(60,TimeUnit.SECONDS);
                } catch (InterruptedException e) {
                    log.error("执行异常",e);
                } finally {
                    log.info("停止容器{}",containerId);
                    handler.getClient().stopContainerCmd(handler.containerId).exec();
                    log.info("删除容器{}",containerId);
                    handler.getClient().removeContainerCmd(handler.containerId).exec();
                }
            }
        } finally {
            if (handler!=null) {
                try {
                    handler.getClient().close();
                } catch (IOException e) {
                    log.error("关闭容器失败",e);
                }
            }
        }

    }
    public DockerExecHandler loadRunCommand(String runCommand, DockerEndpoint dockerEndpoint){
        String[] split = runCommand.split("\\s+");
        List<String> keywords = new ArrayList<>(split.length);
        keywords.addAll(Lists.list(split));
        keywords.removeIf(next -> next.equals("docker") || next.equals("run"));
        // 容器路径映射
        String volume ="";
        // 镜像名
        String image = "";
        // 容器命令 通常为bash或者 /bin/bash
        String initCommand = "";
        for (int i = 0; i < keywords.size(); i++) {
            String keyword = keywords.get(i);
            switch (keyword){
                case "-v":{
                    volume = keywords.get(++i);
                    break;
                }
                case "-p":{
                    log.warn("忽略暂时不支持的-p {}",keywords.get(++i));
                    break;
                }
                case "--name":{
                    log.warn("忽略暂时不支持的--name {}",keywords.get(++i));
                    break;
                }
                default:{
                    if(!keyword.startsWith("-")){
                        if (StringUtils.isEmpty(image)){
                            image = keyword;
                            log.info("识别到镜像 ：{}",image);
                        } else if (StringUtils.isEmpty(initCommand)) {
                            initCommand = keyword;
                            log.info("识别到容器运行命令 ：{}",initCommand);
                        }
                    }
                    else {
                        log.warn("未识别部分：{}",keyword);
                    }
                }
            }
        }
        URI uri = null;
        try {
            if (dockerEndpoint.getType() == DockerEndpoint.DockerEndpointTypeEnum.LOCAL) {
                // 使用本地挂载
                uri = new URI("unix:///var/run/docker.sock");
            } else if (dockerEndpoint.getType() == DockerEndpoint.DockerEndpointTypeEnum.REMOTE) {
                // 远程挂载
                uri = new URI(format("tcp://%s:%s", dockerEndpoint.getHost(), dockerEndpoint.getPort()));
            } else {
                log.error("Unsupported Docker endpoint type: {}", dockerEndpoint.getType());
                return null;
            }
        } catch (URISyntaxException e) {
            log.error("docker运行时创建uri失败",e);
        }
        if (uri == null) {
            return null;
        }
        ApacheDockerHttpClient httpClient = new ApacheDockerHttpClient.Builder().dockerHost(uri).build();
        DockerClient build = DockerClientBuilder.getInstance().withDockerHttpClient(httpClient).build();
        HostConfig hostConfig = HostConfig
                .newHostConfig();
        if (StringUtils.isNotEmpty(volume)) {
            hostConfig.withBinds(Bind.parse(volume));
        }
//                        .withPortBindings(PortBinding.parse()) // 暂时先不支持端口绑定
        if (StringUtils.isEmpty(image)||StringUtils.isEmpty(initCommand)) {
            log.error("缺少参数image或者initCommand");
            return null;
        }
        // 创建容器
        CreateContainerResponse exec = build.createContainerCmd(image).withHostConfig(hostConfig)
                .withTty(true).withCmd(initCommand).exec();
        String containerId = exec.getId();
        // 启动容器
        build.startContainerCmd(containerId).exec();
        log.info("容器创建并启动完成");
        return new DockerExecHandler(build, containerId);
    }
    @Data
    @AllArgsConstructor
    private static class DockerExecHandler{
        private DockerClient client;
        private String containerId;
    }
}
