// Copyright (C) 2025-2026 Joel Rosdahl and other contributors
//
// See doc/authors.adoc for a complete list of contributors.
//
// This program is free software; you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the Free
// Software Foundation; either version 3 of the License, or (at your option)
// any later version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
// more details.
//
// You should have received a copy of the GNU General Public License along with
// this program; if not, write to the Free Software Foundation, Inc., 51
// Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

#include "client.hpp"

#include <ccache/util/expected.hpp>
#include <ccache/util/format.hpp>

#include <algorithm>
#include <cstring>

namespace storage::remote {

namespace {

constexpr uint8_t k_request_get = 0x00;
constexpr uint8_t k_request_put = 0x01;
constexpr uint8_t k_request_remove = 0x02;
constexpr uint8_t k_request_stop = 0x03;
constexpr uint8_t k_request_info = 0x04;
constexpr uint8_t k_request_exists = 0x05;

} // namespace

static Client::Error
make_error(const util::IpcError& ipc_error)
{
  switch (ipc_error.failure) {
  case util::IpcError::Failure::error:
    return Client::Error(Client::Failure::error, ipc_error.message);
  case util::IpcError::Failure::permission_denied:
    return Client::Error(Client::Failure::permission_denied, ipc_error.message);
  case util::IpcError::Failure::timeout:
    return Client::Error(Client::Failure::timeout, ipc_error.message);
  default:
    return Client::Error(Client::Failure::error, "internal error");
  }
}

Client::Client(std::chrono::milliseconds data_timeout,
               std::chrono::milliseconds request_timeout)
  : m_data_timeout(data_timeout),
    m_request_timeout(request_timeout)
{
}

Client::~Client()
{
  close();
}

std::chrono::milliseconds
Client::calculate_timeout() const
{
  // Calculate remaining time for current request
  auto now = std::chrono::steady_clock::now();
  auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
    now - m_request_start_time);
  auto remaining_request_timeout = m_request_timeout - elapsed;

  if (remaining_request_timeout <= std::chrono::milliseconds(0)) {
    // Already expired.
    return std::chrono::milliseconds(0);
  }

  return std::min(m_data_timeout, remaining_request_timeout);
}

tl::expected<void, Client::Error>
Client::connect(const std::string& path)
{
  if (m_connected) {
    return tl::unexpected(Error(Failure::error, "Already connected"));
  }

  m_request_start_time = std::chrono::steady_clock::now();

  auto timeout = calculate_timeout();
  auto result = m_channel.connect(path, timeout);
  if (!result) {
    return tl::unexpected(make_error(result.error()));
  }

  if (auto greeting_result = read_greeting(); !greeting_result) {
    // Don't leave the channel connected, otherwise a later connect attempt
    // would fail since the transport is still in use.
    close();
    return tl::unexpected(greeting_result.error());
  }

  m_connected = true;
  return {};
}

uint8_t
Client::protocol_version() const
{
  return m_protocol_version;
}

const std::vector<Client::Capability>&
Client::capabilities() const
{
  return m_capabilities;
}

bool
Client::has_capability(Capability cap) const
{
  return std::find(m_capabilities.begin(), m_capabilities.end(), cap)
         != m_capabilities.end();
}

tl::expected<void, Client::Error>
Client::verify_connected() const
{
  if (!m_connected) {
    return tl::unexpected(Error(Failure::error, "Not connected"));
  }
  return {};
}

tl::expected<void, Client::Error>
Client::verify_key_size(std::span<const uint8_t> key)
{
  if (key.size() > 255) {
    return tl::unexpected(
      Error(Failure::error, "Key too long (max 255 bytes)"));
  }
  return {};
}

tl::expected<bool, Client::Error>
Client::exists(std::span<const uint8_t> key)
{
  TRY(verify_connected());
  TRY(verify_key_size(key));

  m_request_start_time = std::chrono::steady_clock::now();

  util::Bytes msg;
  msg.reserve(2 + key.size());
  msg.push_back(k_request_exists);
  msg.push_back(static_cast<uint8_t>(key.size()));
  msg.insert(msg.end(), key.data(), key.size());
  TRY(send_bytes(msg));
  return receive_response_exists();
}

tl::expected<std::optional<util::Bytes>, Client::Error>
Client::get(std::span<const uint8_t> key)
{
  TRY(verify_connected());
  TRY(verify_key_size(key));

  m_request_start_time = std::chrono::steady_clock::now();

  util::Bytes msg;
  msg.reserve(2 + key.size());
  msg.push_back(k_request_get);
  msg.push_back(static_cast<uint8_t>(key.size()));
  msg.insert(msg.end(), key.data(), key.size());
  TRY(send_bytes(msg));

  return receive_response_get();
}

tl::expected<Client::InfoResponse, Client::Error>
Client::info()
{
  TRY(verify_connected());

  m_request_start_time = std::chrono::steady_clock::now();

  util::Bytes msg({k_request_info});
  TRY(send_bytes(msg));

  Client::InfoResponse response;

  TRY_ASSIGN(uint8_t identity_len, receive_u8());
  TRY_ASSIGN(auto identity_bytes, receive_bytes(identity_len));
  response.server_identity.assign(
    reinterpret_cast<const char*>(identity_bytes.data()), identity_len);

  TRY_ASSIGN(uint8_t diag_num, receive_u8());
  response.diagnostics.reserve(diag_num);
  for (uint8_t i = 0; i < diag_num; ++i) {
    TRY_ASSIGN(uint8_t msg_len, receive_u8());
    TRY_ASSIGN(auto msg_bytes, receive_bytes(msg_len));
    response.diagnostics.emplace_back(
      reinterpret_cast<const char*>(msg_bytes.data()), msg_len);
  }

  return response;
}

tl::expected<bool, Client::Error>
Client::put(std::span<const uint8_t> key,
            std::span<const uint8_t> value,
            PutFlags flags)
{
  TRY(verify_connected());
  TRY(verify_key_size(key));

  m_request_start_time = std::chrono::steady_clock::now();

  uint8_t flag_byte = flags.overwrite ? 0x01 : 0x00;
  uint64_t value_len = value.size();
  uint8_t len_bytes[sizeof(uint64_t)];
  std::memcpy(len_bytes, &value_len, sizeof(uint64_t)); // host byte order
  util::Bytes header;
  header.reserve(1 + 1 + key.size() + 1 + sizeof(uint64_t));
  header.push_back(k_request_put);
  header.push_back(static_cast<uint8_t>(key.size()));
  header.insert(header.end(), key.data(), key.size());
  header.push_back(flag_byte);
  header.insert(header.end(), len_bytes, sizeof(uint64_t));
  TRY(send_bytes(header));
  TRY(send_bytes(value));
  return receive_response_ok_noop_error();
}

tl::expected<bool, Client::Error>
Client::remove(std::span<const uint8_t> key)
{
  TRY(verify_connected());
  TRY(verify_key_size(key));

  m_request_start_time = std::chrono::steady_clock::now();

  util::Bytes msg;
  msg.reserve(2 + key.size());
  msg.push_back(k_request_remove);
  msg.push_back(static_cast<uint8_t>(key.size()));
  msg.insert(msg.end(), key.data(), key.size());
  TRY(send_bytes(msg));
  return receive_response_ok_noop_error();
}

tl::expected<void, Client::Error>
Client::stop()
{
  TRY(verify_connected());

  m_request_start_time = std::chrono::steady_clock::now();

  const uint8_t msg[] = {k_request_stop};
  TRY(send_bytes(msg));
  return receive_response_void();
}

void
Client::close()
{
  m_channel.close();
  m_connected = false;
  m_protocol_version = 0;
  m_capabilities.clear();
}

tl::expected<void, Client::Error>
Client::read_greeting()
{
  TRY_ASSIGN(m_protocol_version, receive_u8());
  if (m_protocol_version != k_protocol_version) {
    return tl::unexpected(
      Error(Failure::error,
            FMT("Unsupported protocol version: {}", m_protocol_version)));
  }

  TRY_ASSIGN(uint8_t cap_len, receive_u8());
  m_capabilities.clear();
  m_capabilities.reserve(cap_len);
  for (uint8_t i = 0; i < cap_len; ++i) {
    TRY_ASSIGN(uint8_t cap_byte, receive_u8());
    m_capabilities.push_back(static_cast<Capability>(cap_byte));
  }

  return {};
}

tl::expected<void, Client::Error>
Client::send_bytes(std::span<const uint8_t> data)
{
  auto timeout = calculate_timeout();
  auto result = m_channel.send(data, timeout);
  if (!result) {
    return tl::unexpected(make_error(result.error()));
  }
  return {};
}

tl::expected<util::Bytes, Client::Error>
Client::receive_bytes(size_t count)
{
  util::Bytes result(count);
  size_t total_received = 0;

  while (total_received < count) {
    std::span<uint8_t> buffer(result.data() + total_received,
                              count - total_received);
    auto timeout = calculate_timeout();
    auto recv_result = m_channel.receive(buffer, timeout);
    if (!recv_result) {
      return tl::unexpected(make_error(recv_result.error()));
    }

    if (*recv_result == 0) {
      return tl::unexpected(
        Error(Failure::error, "Connection closed by server"));
    }

    total_received += *recv_result;
  }

  return result;
}

tl::expected<uint8_t, Client::Error>
Client::receive_u8()
{
  TRY_ASSIGN(auto data, receive_bytes(sizeof(uint8_t)));
  return data[0];
}

tl::expected<uint64_t, Client::Error>
Client::receive_u64()
{
  TRY_ASSIGN(auto data, receive_bytes(sizeof(uint64_t)));
  uint64_t value;
  std::memcpy(&value, data.data(), sizeof(uint64_t)); // host byte order
  return value;
}

tl::expected<std::string, Client::Error>
Client::receive_error_string()
{
  TRY_ASSIGN(uint8_t msg_len, receive_u8());
  TRY_ASSIGN(auto msg_bytes, receive_bytes(msg_len));
  return std::string(msg_bytes.begin(), msg_bytes.end());
}

tl::expected<bool, Client::Error>
Client::receive_response_exists()
{
  TRY_ASSIGN(uint8_t status_byte, receive_u8());
  auto status = static_cast<Status>(status_byte);

  switch (status) {
  case Status::ok: {
    TRY_ASSIGN(uint8_t exists, receive_u8());
    return exists;
  }

  case Status::error: {
    TRY_ASSIGN(auto err_msg, receive_error_string());
    return tl::unexpected(Error(Failure::error, err_msg));
  }

  default:
    return tl::unexpected(
      Error(Failure::error, FMT("Invalid status code: {}", status_byte)));
  }
}

tl::expected<std::optional<util::Bytes>, Client::Error>
Client::receive_response_get()
{
  TRY_ASSIGN(uint8_t status_byte, receive_u8());
  auto status = static_cast<Status>(status_byte);

  switch (status) {
  case Status::ok: {
    TRY_ASSIGN(uint64_t value_len, receive_u64());
    TRY_ASSIGN(auto value, receive_bytes(value_len));
    return value;
  }

  case Status::noop: // key not found
    return std::nullopt;

  case Status::error: {
    TRY_ASSIGN(auto err_msg, receive_error_string());
    return tl::unexpected(Error(Failure::error, err_msg));
  }

  default:
    return tl::unexpected(
      Error(Failure::error, FMT("Invalid status code: {}", status_byte)));
  }
}

tl::expected<bool, Client::Error>
Client::receive_response_ok_noop_error()
{
  TRY_ASSIGN(uint8_t status_byte, receive_u8());
  auto status = static_cast<Status>(status_byte);

  switch (status) {
  case Status::ok:
    return true;

  case Status::noop:
    return false;

  case Status::error: {
    TRY_ASSIGN(auto err_msg, receive_error_string());
    return tl::unexpected(Error(Failure::error, err_msg));
  }

  default:
    return tl::unexpected(
      Error(Failure::error, FMT("Invalid status code: {}", status_byte)));
  }
}

tl::expected<void, Client::Error>
Client::receive_response_void()
{
  TRY_ASSIGN(uint8_t status_byte, receive_u8());
  auto status = static_cast<Status>(status_byte);

  switch (status) {
  case Status::ok:
    return {};

  case Status::noop:
    // This shouldn't happen for stop, but treat it as success.
    return {};

  case Status::error: {
    TRY_ASSIGN(auto err_msg, receive_error_string());
    return tl::unexpected(Error(Failure::error, err_msg));
  }

  default:
    return tl::unexpected(
      Error(Failure::error, FMT("Invalid status code: {}", status_byte)));
  }
}

} // namespace storage::remote
