package me.zacharias.chat;

import com.github.dockerjava.api.DockerClient;
import com.github.dockerjava.api.async.ResultCallback;
import com.github.dockerjava.api.async.ResultCallbackTemplate;
import com.github.dockerjava.api.command.AttachContainerCmd;
import com.github.dockerjava.api.command.LogContainerCmd;
import com.github.dockerjava.api.command.StartContainerCmd;
import com.github.dockerjava.api.model.Frame;
import com.github.dockerjava.core.DefaultDockerClientConfig;
import com.github.dockerjava.core.DockerClientBuilder;
import com.github.dockerjava.core.DockerClientConfig;
import com.github.dockerjava.core.command.LogContainerResultCallback;
import me.zacharias.chat.ollama.*;
import me.zacharias.chat.ollama.exceptions.OllamaToolErrorException;
import org.json.JSONObject;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.logging.Logger;

import static me.zacharias.chat.Main.writeLog;

public class PythonRunner extends OllamaFuntionTool {
    DockerClient dockerClient;
    public PythonRunner() {

        DefaultDockerClientConfig.Builder config
                = DefaultDockerClientConfig.createDefaultConfigBuilder()
                .withDockerHost("tcp://localhost:2375")
                .withDockerTlsVerify(false);
        dockerClient = DockerClientBuilder
                .getInstance(config)
                .build();
    }

    @Override
    public String name() {
        return "python_runner";
    }

    @Override
    public String description() {
        return "Runs python code";
    }

    @Override
    public OllamaPerameter parameters() {
        return OllamaPerameter.builder()
                .addProperty("code", OllamaPerameter.OllamaPerameterBuilder.Type.STRING, "The code to be executed", true)
                .addProperty("name", OllamaPerameter.OllamaPerameterBuilder.Type.STRING, "The name of the python code")
                .build();
    }

    @Override
    public OllamaToolRespnce function(OllamaFunctionArgument... args) {
        if(args.length == 0)
        {
            throw new OllamaToolErrorException(name(), "Missing code argument");
        }

        String name = null;
        String code = null;

        for(OllamaFunctionArgument arg : args)
        {
            if(arg.getArgument().equals("name"))
            {
                name = (String) arg.getValue();
                if(!name.endsWith(".py"))
                {
                    name += ".py";
                }
            } else if (arg.getArgument().equals("code")) {
                code = (String) arg.getValue();
            }
        }

        if(name == null)
        {
            try {
                MessageDigest digest = MessageDigest.getInstance("SHA-256");
                byte[] encodedhash = digest.digest(String.valueOf(args[0].getValue()).getBytes(StandardCharsets.UTF_8));
                StringBuffer hexString = new StringBuffer();
                for(byte b : encodedhash)
                {
                    hexString.append(String.format("%02x", b));
                }
                name = hexString.toString()+".py";
            }catch (Exception e) {}
        }

        name = name.replace(" ", "_");

        writeLog("Running python code `" + name + "`");

        File pythonFile = new File("./pythonFiles", name);

        if(!pythonFile.exists())
        {
            try {
                BufferedWriter writer = new BufferedWriter(new FileWriter(pythonFile));
                writer.write(code);
                writer.close();
            }catch(IOException e) {}
        }

        String containerId = dockerClient.createContainerCmd("python").withCmd("python", name).exec().getId();
        dockerClient.copyArchiveToContainerCmd(containerId)
                .withHostResource(pythonFile.getPath())
                //.withRemotePath("~/")
                .exec();

        dockerClient.startContainerCmd(containerId).exec();

        GetContainerLog log = new GetContainerLog(dockerClient, containerId);

        List<String> logs = new ArrayList<>();

        do {
            try {
                Thread.sleep(2000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            logs.addAll(log.getDockerLogs());
        }
        while (logs.isEmpty());

        StringBuilder output = new StringBuilder();

        for(String s : logs)
        {
            output.append(s).append("\n");
        }

        //writeLog("Result from python: " + output.toString());

        return new OllamaToolRespnce(name(), output.toString());
    }

    public class GetContainerLog {
        private DockerClient dockerClient;
        private String containerId;
        private int lastLogTime;

        private static String nameOfLogger = "dockertest.PrintContainerLog";
        private static Logger myLogger = Logger.getLogger(nameOfLogger);

        public GetContainerLog(DockerClient dockerClient, String containerId) {
            this.dockerClient = dockerClient;
            this.containerId = containerId;
            this.lastLogTime = (int) (System.currentTimeMillis() / 1000);
        }

        public List<String> getDockerLogs() {

            final List<String> logs = new ArrayList<>();

            LogContainerCmd logContainerCmd = dockerClient.logContainerCmd(containerId);
            logContainerCmd.withStdOut(true).withStdErr(true);
            logContainerCmd.withSince(lastLogTime);  // UNIX timestamp (integer) to filter logs. Specifying a timestamp will only output log-entries since that timestamp.
            // logContainerCmd.withTail(4);  // get only the last 4 log entries

            logContainerCmd.withTimestamps(true);

            try {
                logContainerCmd.exec(new ResultCallback.Adapter<Frame>() {
                    @Override
                    public void onNext(Frame item) {
                        logs.add(item.toString());
                    }
                }).awaitCompletion();
            } catch (InterruptedException e) {
                myLogger.severe("Interrupted Exception!" + e.getMessage());
            }

            lastLogTime = (int) (System.currentTimeMillis() / 1000) + 5;  // assumes at least a 5 second wait between calls to getDockerLogs

            return logs;
        }
    }
}