// 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;
}