// Copyright 2021 The Tint Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <stdio.h>
#include <stdlib.h>
#include <fstream>
#include <sstream>
#include <string>
#include <thread>
#include <type_traits>
#include <vector>

#include "tools/src/cmd/remote-compile/compile.h"
#include "tools/src/cmd/remote-compile/socket.h"

namespace {

#if 0
#define DEBUG(msg, ...) printf(msg "\n", ##__VA_ARGS__)
#else
#define DEBUG(...)
#endif

/// Print the tool usage, and exit with 1.
void ShowUsage() {
    const char* name = "tint-remote-compile";
    printf(R"(%s is a tool for compiling a shader on a remote machine

usage as server:
  %s -s [-p port-number]

usage as client:
  %s [-p port-number] [server-address] shader-file-path

  [server-address] can be omitted if the TINT_REMOTE_COMPILE_ADDRESS environment
  variable is set.
  Alternatively, you can pass xcrun arguments so %s can be used as a
  drop-in replacement.
)",
           name, name, name, name);
    exit(1);
}

/// The protocol version code. Bump each time the protocol changes
constexpr uint32_t kProtocolVersion = 1;

/// Supported shader source languages
enum SourceLanguage {
    MSL,
};

/// Stream is a serialization wrapper around a socket
struct Stream {
    /// The underlying socket
    Socket* const socket;
    /// Error state
    std::string error;

    /// Writes a uint32_t to the socket
    Stream operator<<(uint32_t v) {
        if (error.empty()) {
            Write(&v, sizeof(v));
        }
        return *this;
    }

    /// Reads a uint32_t from the socket
    Stream operator>>(uint32_t& v) {
        if (error.empty()) {
            Read(&v, sizeof(v));
        }
        return *this;
    }

    /// Writes a std::string to the socket
    Stream operator<<(const std::string& v) {
        if (error.empty()) {
            uint32_t count = static_cast<uint32_t>(v.size());
            *this << count;
            if (count) {
                Write(v.data(), count);
            }
        }
        return *this;
    }

    /// Reads a std::string from the socket
    Stream operator>>(std::string& v) {
        uint32_t count = 0;
        *this >> count;
        if (count) {
            std::vector<char> buf(count);
            if (Read(buf.data(), count)) {
                v = std::string(buf.data(), buf.size());
            }
        } else {
            v.clear();
        }
        return *this;
    }

    /// Writes an enum value to the socket
    template <typename T>
    std::enable_if_t<std::is_enum<T>::value, Stream> operator<<(T e) {
        return *this << static_cast<uint32_t>(e);
    }

    /// Reads an enum value from the socket
    template <typename T>
    std::enable_if_t<std::is_enum<T>::value, Stream> operator>>(T& e) {
        uint32_t v;
        *this >> v;
        e = static_cast<T>(v);
        return *this;
    }

  private:
    bool Write(const void* data, size_t size) {
        if (error.empty()) {
            if (!socket->Write(data, size)) {
                error = "Socket::Write() failed";
            }
        }
        return error.empty();
    }

    bool Read(void* data, size_t size) {
        auto buf = reinterpret_cast<uint8_t*>(data);
        while (size > 0 && error.empty()) {
            if (auto n = socket->Read(buf, size)) {
                if (n > size) {
                    error = "Socket::Read() returned more bytes than requested";
                    return false;
                }
                size -= n;
                buf += n;
            }
        }
        return error.empty();
    }
};

////////////////////////////////////////////////////////////////////////////////
// Messages
////////////////////////////////////////////////////////////////////////////////

/// Base class for all messages
struct Message {
    /// The type of the message
    enum class Type {
        ConnectionRequest,
        ConnectionResponse,
        CompileRequest,
        CompileResponse,
    };

    explicit Message(Type ty) : type(ty) {}

    const Type type;
};

struct ConnectionResponse : Message {  // Server -> Client
    ConnectionResponse() : Message(Type::ConnectionResponse) {}

    template <typename T>
    void Serialize(T&& f) {
        f(error);
    }

    std::string error;
};

struct ConnectionRequest : Message {  // Client -> Server
    using Response = ConnectionResponse;

    explicit ConnectionRequest(uint32_t proto_ver = kProtocolVersion)
        : Message(Type::ConnectionRequest), protocol_version(proto_ver) {}

    template <typename T>
    void Serialize(T&& f) {
        f(protocol_version);
    }

    uint32_t protocol_version;
};

struct CompileResponse : Message {  //  Server -> Client
    CompileResponse() : Message(Type::CompileResponse) {}

    template <typename T>
    void Serialize(T&& f) {
        f(error);
    }

    std::string error;
};

struct CompileRequest : Message {  // Client -> Server
    using Response = CompileResponse;

    CompileRequest() : Message(Type::CompileRequest) {}
    CompileRequest(SourceLanguage lang, std::string src)
        : Message(Type::CompileRequest), language(lang), source(src) {}

    template <typename T>
    void Serialize(T&& f) {
        f(language);
        f(source);
    }

    SourceLanguage language;
    std::string source;
};

/// Writes the message `m` to the stream `s`
template <typename MESSAGE>
std::enable_if_t<std::is_base_of<Message, MESSAGE>::value, Stream>& operator<<(Stream& s,
                                                                               const MESSAGE& m) {
    s << m.type;
    const_cast<MESSAGE&>(m).Serialize([&s](const auto& value) { s << value; });
    return s;
}

/// Reads the message `m` from the stream `s`
template <typename MESSAGE>
std::enable_if_t<std::is_base_of<Message, MESSAGE>::value, Stream>& operator>>(Stream& s,
                                                                               MESSAGE& m) {
    Message::Type ty;
    s >> ty;
    if (ty == m.type) {
        m.Serialize([&s](auto& value) { s >> value; });
    } else {
        std::stringstream ss;
        ss << "expected message type " << static_cast<int>(m.type) << ", got "
           << static_cast<int>(ty);
        s.error = ss.str();
    }
    return s;
}

/// Writes the request message `req` to the stream `s`, then reads and returns
/// the response message from the same stream.
template <typename REQUEST, typename RESPONSE = typename REQUEST::Response>
RESPONSE Send(Stream& s, const REQUEST& req) {
    s << req;
    if (s.error.empty()) {
        RESPONSE resp;
        s >> resp;
        if (s.error.empty()) {
            return resp;
        }
    }
    return {};
}

}  // namespace

bool RunServer(std::string port);
bool RunClient(std::string address, std::string port, std::string file);

int main(int argc, char* argv[]) {
    bool run_server = false;
    std::string port = "19000";

    std::vector<std::string> args;
    for (int i = 1; i < argc; i++) {
        std::string arg = argv[i];
        if (arg == "-s" || arg == "--server") {
            run_server = true;
            continue;
        }
        if (arg == "-p" || arg == "--port") {
            if (i < argc - 1) {
                i++;
                port = argv[i];
            } else {
                printf("expected port number");
                exit(1);
            }
            continue;
        }

        // xcrun flags are ignored so this executable can be used as a replacement
        // for xcrun.
        if ((arg == "-x" || arg == "-sdk") && (i < argc - 1)) {
            i++;
            continue;
        }
        if (arg == "metal") {
            for (; i < argc; i++) {
                if (std::string(argv[i]) == "-c") {
                    break;
                }
            }
            continue;
        }

        args.emplace_back(arg);
    }

    bool success = false;

    if (run_server) {
        success = RunServer(port);
    } else {
        std::string address;
        std::string file;
        switch (args.size()) {
            case 1:
                if (auto* addr = getenv("TINT_REMOTE_COMPILE_ADDRESS")) {
                    address = addr;
                }
                file = args[0];
                break;
            case 2:
                address = args[0];
                file = args[1];
                break;
        }
        if (address.empty() || file.empty()) {
            ShowUsage();
        }
        success = RunClient(address, port, file);
    }

    if (!success) {
        exit(1);
    }

    return 0;
}

bool RunServer(std::string port) {
    auto server_socket = Socket::Listen("", port.c_str());
    if (!server_socket) {
        printf("Failed to listen on port %s\n", port.c_str());
        return false;
    }
    printf("Listening on port %s...\n", port.c_str());
    while (auto conn = server_socket->Accept()) {
        std::thread([=] {
            DEBUG("Client connected...");
            Stream stream{conn.get()};

            {
                ConnectionRequest req;
                stream >> req;
                if (!stream.error.empty()) {
                    printf("%s\n", stream.error.c_str());
                    return;
                }
                ConnectionResponse resp;
                if (req.protocol_version != kProtocolVersion) {
                    DEBUG("Protocol version mismatch");
                    resp.error = "Protocol version mismatch";
                    stream << resp;
                    return;
                }
                stream << resp;
            }
            DEBUG("Connection established");
            {
                CompileRequest req;
                stream >> req;
                if (!stream.error.empty()) {
                    printf("%s\n", stream.error.c_str());
                    return;
                }
#ifdef TINT_ENABLE_MSL_COMPILATION_USING_METAL_API
                if (req.language == SourceLanguage::MSL) {
                    auto result = CompileMslUsingMetalAPI(req.source);
                    CompileResponse resp;
                    if (!result.success) {
                        resp.error = result.output;
                    }
                    stream << resp;
                    return;
                }
#endif
                CompileResponse resp;
                resp.error = "server cannot compile this type of shader";
                stream << resp;
            }
        }).detach();
    }
    return true;
}

bool RunClient(std::string address, std::string port, std::string file) {
    // Read the file
    std::ifstream input(file, std::ios::binary);
    if (!input) {
        printf("Couldn't open '%s'\n", file.c_str());
        return false;
    }
    std::string source((std::istreambuf_iterator<char>(input)), std::istreambuf_iterator<char>());

    constexpr const int timeout_ms = 10000;
    DEBUG("Connecting to %s:%s...", address.c_str(), port.c_str());
    auto conn = Socket::Connect(address.c_str(), port.c_str(), timeout_ms);
    if (!conn) {
        printf("Connection failed\n");
        return false;
    }

    Stream stream{conn.get()};

    DEBUG("Sending connection request...");
    auto conn_resp = Send(stream, ConnectionRequest{kProtocolVersion});
    if (!stream.error.empty()) {
        printf("%s\n", stream.error.c_str());
        return false;
    }
    if (!conn_resp.error.empty()) {
        printf("%s\n", conn_resp.error.c_str());
        return false;
    }
    DEBUG("Connection established. Requesting compile...");
    auto comp_resp = Send(stream, CompileRequest{SourceLanguage::MSL, source});
    if (!stream.error.empty()) {
        printf("%s\n", stream.error.c_str());
        return false;
    }
    if (!comp_resp.error.empty()) {
        printf("%s\n", comp_resp.error.c_str());
        return false;
    }
    DEBUG("Compilation successful");
    return true;
}