#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <vector>
#include <iostream>
#include <string>
#include <cstdlib>
#include <sstream>
#include "utils/assert.hpp"

#include "execute.hpp"
#include "utils/utils.hpp"
#include "utils/b64ed.hpp"
#include "config.hpp"
#include "utils/directories.hpp"
#include "utils/output.hpp"

namespace dropshell
{
    bool EXITSTATUSCHECK(int ret)
    {
        return (ret != -1 && WIFEXITED(ret) && (WEXITSTATUS(ret) == 0)); // ret is -1 if the command failed to execute.
    }

    // ----------------------------------------------------------------------------------------------------------
    // execute_local_command_interactive
    // ----------------------------------------------------------------------------------------------------------
    bool execute_local_command_interactive(const sCommand &command)
    {
        if (command.get_command_to_run().empty())
            return false;
        std::string full_command = command.construct_cmd(localpath::agent()+"/bb64"); // Get the command string

        pid_t pid = fork();

        if (pid == -1)
        {
            // Fork failed
            perror("fork failed");
            return false;
        }
        else if (pid == 0)
        {
            int rval = system(full_command.c_str());
            exit(rval);
        }
        else
        {
            // Parent process
            int ret;
            // Wait for the child process to complete
            waitpid(pid, &ret, 0);

            return EXITSTATUSCHECK(ret);
        }
    }

    // ----------------------------------------------------------------------------------------------------------
    // execute_local_command
    // ----------------------------------------------------------------------------------------------------------

    class fancypinter
    {
    public:
        fancypinter(sColour startColour) : startColour_(startColour), currentColour_(startColour) {}

        void print_chunk(std::string chunk)
        {
            if (chunk.empty())
                return;

            if (newline_)
            {
                // sniff the mode... if the string starts with warning or warning: then set mode to WARNING. etc.
                if (chunk.find("warning") == 0)
                    currentColour_ = sColour::WARNING;
                else if (chunk.find("error") == 0)
                    currentColour_ = sColour::ERROR;
                else if (chunk.find("debug") == 0)
                    currentColour_ = sColour::DEBUG;
                else if (chunk.find("info") == 0)
                    currentColour_ = sColour::INFO;
                else
                    currentColour_ = startColour_;
            }
            colourstream(currentColour_) << chunk;
            newline_ = (chunk[chunk.size() - 1] == '\n');
        }

        void print(const std::string &buffer)
        {
            size_t start = 0;
            while (start < buffer.size())
            {
                size_t newline_pos = buffer.find('\n', start);
                if (newline_pos == std::string::npos)
                {
                    if (start < buffer.size())
                    {
                        print_chunk(buffer.substr(start));
                    }
                    break;
                }
                print_chunk(buffer.substr(start, newline_pos - start + 1)); // include the newline
                start = newline_pos + 1;
            }
        }

    private:
        bool newline_ = true;
        sColour startColour_;
        sColour currentColour_;
    };

    bool execute_local_command(std::string directory_to_run_in, std::string command_to_run, const std::map<std::string, std::string> &env_vars, std::string *output, cMode mode)
    {
        sCommand command(directory_to_run_in, command_to_run, env_vars);

        if (hasFlag(mode, cMode::Interactive))
        {
            ASSERT(output == nullptr, "Interactive mode and an output string cannot be used together");

            return execute_local_command_interactive(command);
        }

        if (command.get_command_to_run().empty())
            return false;

        bool silent = hasFlag(mode, cMode::Silent);

        std::string full_cmd;
        if (!hasFlag(mode, cMode::NoBB64))
            full_cmd = command.construct_cmd(localpath::agent()+"/bb64");
        else
            full_cmd = command.construct_cmd("");

        if (output != nullptr)
            full_cmd += " 2>&1"; // capture both stdout and stderr

        FILE *pipe = popen(full_cmd.c_str(), "r");
        if (!pipe)
        {
            return false;
        }
        char buffer[128];
        fancypinter fancyprint(sColour::DEBUG);
        while (fgets(buffer, sizeof(buffer), pipe) != nullptr)
        {
            if (output != nullptr)
                (*output) += buffer;

            if (!silent)
                fancyprint.print(buffer);
        }
        int ret = pclose(pipe);
        return EXITSTATUSCHECK(ret);
    }

    // ----------------------------------------------------------------------------------------------------------
    // execute_ssh_command
    // ----------------------------------------------------------------------------------------------------------
    bool execute_ssh_command(const sSSHInfo &ssh_info, const sCommand &remote_command, cMode mode, std::string *output)
    {
        if (remote_command.get_command_to_run().empty())
            return false;

        std::stringstream ssh_cmd;
        ssh_cmd << "ssh -p " << ssh_info.port << " " << (hasFlag(mode, cMode::Interactive) ? "-tt " : "")
                << ssh_info.user << "@" << ssh_info.host;

        std::string remote_bb64_path;

        if (!hasFlag(mode, cMode::NoBB64))
            remote_bb64_path = remotepath::agent(ssh_info.server_ID) + "/bb64";

        bool rval = execute_local_command(
            "",                                                                    // local directory to run in
            ssh_cmd.str() + " " + remote_command.construct_cmd(remote_bb64_path),  // local command to run
            {},                                                                    // environment variables
            output,                                                                // output string
            mode                                                                   // mode
        );

        if (!rval && !hasFlag(mode, cMode::Silent))
        {
            error << "Error: Failed to execute ssh command:" << std::endl;
            debug << ssh_cmd.str() + " " + remote_command.construct_cmd(remote_bb64_path) << std::endl;
        }
        return rval;
    }

    // ----------------------------------------------------------------------------------------------------------
    // makesafecmd
    // ----------------------------------------------------------------------------------------------------------
    std::string sCommand::makesafecmd(std::string bb64path, const std::string &command) const
    {
        if (command.empty())
            return "";
        std::string encoded = base64_encode(dequote(trim(command)));
        std::string commandstr = bb64path + " " + encoded;
        return commandstr;
    }

    // ----------------------------------------------------------------------------------------------------------
    // construct_cmd
    // ----------------------------------------------------------------------------------------------------------
    std::string sCommand::construct_cmd(std::string bb64path) const
    {
        if (mCmd.empty())
            return "";

        // need to construct to change directory and set environment variables
        std::string cmdstr;

        if (!bb64path.empty())
        {
            if (!mDir.empty())
                cmdstr += "cd " + quote(mDir) + " && ";

            if (!mVars.empty())
                for (const auto &env_var : mVars)
                    cmdstr += env_var.first + "=" + quote(dequote(trim(env_var.second))) + " ";

            cmdstr += mCmd;

            cmdstr = makesafecmd(bb64path, cmdstr);
        }
        else 
        { // raw! bootstrapping only.
            ASSERT(mVars.empty(), "Bootstrapping command must not have environment variables");
            if (!mDir.empty())
                cmdstr += mDir + "/" + mCmd;
            else
                cmdstr += mCmd;
        }

        return cmdstr;
    }

} // namespace dropshell