/*
  Dynagate C++20 test suite (port of gateway_test.go)
  Created by Luxferre in 2026, released into the public domain
*/

#ifndef _DARWIN_C_SOURCE
#define _DARWIN_C_SOURCE
#endif

#include "dynagate.hpp"

#include <atomic>
#include <cerrno>
#include <chrono>
#include <csignal>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <functional>
#include <iostream>
#include <mutex>
#include <sstream>
#include <string>
#include <thread>
#include <vector>

#include <arpa/inet.h>
#include <curl/curl.h>
#include <netinet/in.h>
#include <poll.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

namespace {

int g_failures = 0;

void fail_test(const char *test_name, int line, const std::string &msg) {
  std::fprintf(stderr, "  [FAIL] %s (line %d): %s\n",
               test_name, line, msg.c_str());
  ++g_failures;
}

#define EXPECT_TRUE(cond, msg)                                                 \
  do {                                                                         \
    if (!(cond)) {                                                             \
      fail_test(__func__, __LINE__, (msg));                                    \
      return;                                                                  \
    }                                                                          \
  } while (0)

#define EXPECT_EQ(actual, expected, msg)                                       \
  do {                                                                         \
    if (!((actual) == (expected))) {                                           \
      std::ostringstream _oss;                                                 \
      _oss << (msg);                                                           \
      fail_test(__func__, __LINE__, _oss.str());                               \
      return;                                                                  \
    }                                                                          \
  } while (0)

bool write_all_fd(int fd, std::string_view data) {
  size_t written = 0;
  while (written < data.size()) {
    ssize_t n = write(fd, data.data() + written, data.size() - written);
    if (n > 0) {
      written += static_cast<size_t>(n);
    } else if (n < 0 && errno == EINTR) {
      continue;
    } else {
      return false;
    }
  }
  return true;
}

class MockHttpServer {
public:
  using HandlerFunc = std::function<void(const HttpRequest &req, int client_fd)>;

  explicit MockHttpServer(HandlerFunc handler) : handler_(std::move(handler)) {
    fd_ = socket(AF_INET, SOCK_STREAM, 0);
    int one = 1;
    setsockopt(fd_, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
#ifdef SO_NOSIGPIPE
    setsockopt(fd_, SOL_SOCKET, SO_NOSIGPIPE, &one, sizeof(one));
#endif

    struct sockaddr_in addr;
    std::memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
    addr.sin_port = htons(0);

    if (bind(fd_, reinterpret_cast<struct sockaddr *>(&addr), sizeof(addr)) <
        0) {
      std::perror("MockHttpServer bind");
      std::exit(1);
    }
    if (listen(fd_, 64) < 0) {
      std::perror("MockHttpServer listen");
      std::exit(1);
    }

    struct sockaddr_in bound;
    socklen_t len = sizeof(bound);
    getsockname(fd_, reinterpret_cast<struct sockaddr *>(&bound), &len);
    port_ = ntohs(bound.sin_port);
    url_ = "http://127.0.0.1:" + std::to_string(port_);

    thread_ = std::thread([this]() { serve_loop(); });
  }

  ~MockHttpServer() {
    stop_.store(true);
    if (fd_ >= 0) {
      shutdown(fd_, SHUT_RDWR);
      close(fd_);
      fd_ = -1;
    }
    if (thread_.joinable()) {
      thread_.join();
    }
  }

  const std::string &url() const { return url_; }

  static void send_response(int fd,
                            int status_code,
                            std::string_view content_type,
                            std::string_view body) {
    std::string resp = "HTTP/1.1 " + std::to_string(status_code) + " Status\r\n";
    if (!content_type.empty()) {
      resp += "Content-Type: " + std::string(content_type) + "\r\n";
    }
    resp += "Content-Length: " + std::to_string(body.size()) + "\r\n";
    resp += "Connection: close\r\n\r\n";
    resp.append(body.data(), body.size());
    write_all_fd(fd, resp);
  }

private:
  void serve_loop() {
    while (!stop_.load()) {
      struct pollfd pfd;
      pfd.fd = fd_;
      pfd.events = POLLIN;
      pfd.revents = 0;
      int prc = poll(&pfd, 1, 50);
      if (prc <= 0 || !(pfd.revents & POLLIN)) {
        continue;
      }

      struct sockaddr_in caddr;
      socklen_t clen = sizeof(caddr);
      int cfd =
        accept(fd_, reinterpret_cast<struct sockaddr *>(&caddr), &clen);
      if (cfd < 0) {
        continue;
      }

      HttpRequest req;
      if (read_req(cfd, req)) {
        handler_(req, cfd);
      }
      close(cfd);
    }
  }

  static bool read_req(int fd, HttpRequest &req) {
    std::string buf;
    char tmp[4096];
    size_t hend = std::string::npos;
    for (;;) {
      hend = buf.find("\r\n\r\n");
      if (hend != std::string::npos) {
        break;
      }
      ssize_t n = read(fd, tmp, sizeof(tmp));
      if (n <= 0) {
        return false;
      }
      buf.append(tmp, static_cast<size_t>(n));
    }

    std::string_view hblock(buf.data(), hend);
    size_t fl = hblock.find("\r\n");
    std::string_view rline =
      (fl == std::string_view::npos) ? hblock : hblock.substr(0, fl);
    size_t sp1 = rline.find(' ');
    if (sp1 == std::string_view::npos) {
      return false;
    }
    req.method = std::string(rline.substr(0, sp1));
    std::string_view rest = rline.substr(sp1 + 1);
    size_t sp2 = rest.find(' ');
    req.path = std::string(
      (sp2 == std::string_view::npos) ? rest : rest.substr(0, sp2));

    size_t pos = (fl == std::string_view::npos) ? hblock.size() : fl + 2;
    long clen = 0;
    while (pos < hblock.size()) {
      size_t le = hblock.find("\r\n", pos);
      std::string_view line =
        (le == std::string_view::npos)
          ? hblock.substr(pos)
          : hblock.substr(pos, le - pos);
      pos = (le == std::string_view::npos) ? hblock.size() : le + 2;
      size_t colon = line.find(':');
      if (colon != std::string_view::npos) {
        std::string k = str_trim(line.substr(0, colon));
        std::string v = str_trim(line.substr(colon + 1));
        if (str_iequals(k, "Content-Length")) {
          clen = std::atol(v.c_str());
        }
        req.headers.add(k, v);
      }
    }

    size_t boff = hend + 4;
    size_t avail = buf.size() - boff;
    size_t needed = clen > 0 ? static_cast<size_t>(clen) : 0;
    if (avail > 0) {
      req.body.append(buf.data() + boff, std::min(avail, needed));
    }
    while (req.body.size() < needed) {
      ssize_t n = read(fd, tmp, std::min(sizeof(tmp), needed - req.body.size()));
      if (n <= 0) {
        break;
      }
      req.body.append(tmp, static_cast<size_t>(n));
    }
    return true;
  }

  int fd_ = -1;
  int port_ = 0;
  std::string url_;
  std::atomic<bool> stop_{false};
  HandlerFunc handler_;
  std::thread thread_;
};

std::filesystem::path make_temp_dir(std::string_view prefix) {
  auto base = std::filesystem::temp_directory_path();
  auto now = std::chrono::steady_clock::now().time_since_epoch().count();
  auto path =
    base / (std::string(prefix) + "-" + std::to_string(getpid()) + "-" +
            std::to_string(now));
  std::filesystem::create_directories(path);
  return path;
}

struct TempDirGuard {
  std::filesystem::path path;
  explicit TempDirGuard(std::string_view prefix)
    : path(make_temp_dir(prefix)) {}
  ~TempDirGuard() {
    std::error_code ec;
    std::filesystem::remove_all(path, ec);
  }
};

void TestURLBuilding() {
  struct TestCase {
    std::string endpoint;
    std::string req_path;
    std::string expected;
  };

  std::vector<TestCase> tests = {
    {"https://api.openai.com/v1", "/v1/chat/completions",
     "https://api.openai.com/v1/chat/completions"},
    {"https://api.openai.com/v1/", "/v1/chat/completions",
     "https://api.openai.com/v1/chat/completions"},
    {"https://api.openai.com", "/v1/chat/completions",
     "https://api.openai.com/v1/chat/completions"},
    {"http://localhost:8000/api/v1", "/v1/chat/completions",
     "http://localhost:8000/api/v1/chat/completions"},
    {"http://localhost:11434", "/v1/chat/completions",
     "http://localhost:11434/v1/chat/completions"},
    {"https://custom.endpoint/v1/custom", "/v1/chat/completions",
     "https://custom.endpoint/v1/custom/v1/chat/completions"},
  };

  for (const auto &tc : tests) {
    std::string got = build_url(tc.endpoint, tc.req_path);
    EXPECT_EQ(got, tc.expected,
              "build_url(" + tc.endpoint + ", " + tc.req_path + ") = " + got +
                "; want " + tc.expected);
  }
}

void TestTrialOrdering() {
  std::vector<ModelConfig> configs = {
    {"gpt-4", "key-a1", "ep-a", ""},
    {"gpt-4", "key-a2", "ep-a", ""},
    {"gpt-3.5", "key-b1", "ep-b", ""},
    {"claude", "key-c1", "ep-c", ""},
  };
  std::vector<std::string> unique_models = {"gpt-4", "gpt-3.5", "claude"};

  // Case 1: "gpt-4"
  {
    auto trials = get_trial_configs(configs, unique_models, "gpt-4");
    std::vector<ModelConfig> expected = {
      {"gpt-4", "key-a1", "ep-a", ""},
      {"gpt-4", "key-a2", "ep-a", ""},
      {"gpt-3.5", "key-b1", "ep-b", ""},
      {"claude", "key-c1", "ep-c", ""},
    };
    EXPECT_EQ(trials, expected, "Trial ordering for gpt-4 mismatch");
  }

  // Case 2: "gpt-3.5"
  {
    auto trials = get_trial_configs(configs, unique_models, "gpt-3.5");
    std::vector<ModelConfig> expected = {
      {"gpt-3.5", "key-b1", "ep-b", ""},
      {"claude", "key-c1", "ep-c", ""},
      {"gpt-4", "key-a1", "ep-a", ""},
      {"gpt-4", "key-a2", "ep-a", ""},
    };
    EXPECT_EQ(trials, expected, "Trial ordering for gpt-3.5 mismatch");
  }

  // Case 3: "claude"
  {
    auto trials = get_trial_configs(configs, unique_models, "claude");
    std::vector<ModelConfig> expected = {
      {"claude", "key-c1", "ep-c", ""},
      {"gpt-4", "key-a1", "ep-a", ""},
      {"gpt-4", "key-a2", "ep-a", ""},
      {"gpt-3.5", "key-b1", "ep-b", ""},
    };
    EXPECT_EQ(trials, expected, "Trial ordering for claude mismatch");
  }

  // Case 4: "non-existent"
  {
    auto trials = get_trial_configs(configs, unique_models, "non-existent");
    std::vector<ModelConfig> expected = {
      {"gpt-4", "key-a1", "ep-a", ""},
      {"gpt-4", "key-a2", "ep-a", ""},
      {"gpt-3.5", "key-b1", "ep-b", ""},
      {"claude", "key-c1", "ep-c", ""},
    };
    EXPECT_EQ(trials, expected, "Trial ordering for non-existent mismatch");
  }
}

void TestConfigManagerCSVWatcher() {
  TempDirGuard tmp("dynagate-test");
  auto csv_file = (tmp.path / "models.csv").string();

  std::string initial_content =
    "model,key,endpoint,extra\n"
    "gpt-4,key-1,http://localhost:8001,\"{\"\"User-Agent\"\":\"\"test/1.0\"\"}\"\n"
    "gpt-3.5,key-2,http://localhost:8002,\n";

  {
    std::ofstream ofs(csv_file, std::ios::binary);
    ofs << initial_content;
  }

  ConfigManager cm(csv_file);
  auto configs = cm.get_configs();
  EXPECT_EQ(configs.size(), 2u, "Expected 2 initial configs");
  EXPECT_EQ(configs[0].model, "gpt-4", "Unexpected model 0");
  EXPECT_EQ(configs[0].extra, "{\"User-Agent\":\"test/1.0\"}",
            "Unexpected quoted extra JSON");
  EXPECT_EQ(configs[1].model, "gpt-3.5", "Unexpected model 1");

  std::atomic<bool> stop_flag{false};
  std::thread watcher([&]() { cm.watch(stop_flag); });

  std::this_thread::sleep_for(std::chrono::milliseconds(1100));

  std::string updated_content =
    "model,key,endpoint\n"
    "gpt-4,key-1,http://localhost:8001\n"
    "gpt-3.5,key-2,http://localhost:8002\n"
    "claude,key-3,http://localhost:8003\n";

  {
    std::ofstream ofs(csv_file, std::ios::binary | std::ios::trunc);
    ofs << updated_content;
  }

  auto deadline =
    std::chrono::steady_clock::now() + std::chrono::seconds(3);
  std::vector<ModelConfig> updated_configs;
  while (std::chrono::steady_clock::now() < deadline) {
    updated_configs = cm.get_configs();
    if (updated_configs.size() == 3) {
      break;
    }
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
  }

  stop_flag.store(true);
  if (watcher.joinable()) {
    watcher.join();
  }

  EXPECT_EQ(updated_configs.size(), 3u, "Hot-reload failed to load 3 configs");
  EXPECT_EQ(updated_configs[2].model, "claude",
            "Hot-reload third model mismatch");
}

void TestGatewayCompletionsFailover() {
  MockHttpServer server_a([](const HttpRequest &, int fd) {
    MockHttpServer::send_response(
      fd, 500, "application/json", R"({"error": "Internal Server Error A"})");
  });

  MockHttpServer server_b([](const HttpRequest &, int fd) {
    MockHttpServer::send_response(
      fd, 401, "application/json", R"({"error": "Unauthorized B"})");
  });

  MockHttpServer server_c([](const HttpRequest &req, int fd) {
    Json body = Json::parse(req.body);
    std::string model_sent = body["model"].str();
    std::string resp_body =
      R"({"choices":[{"message":{"content":"Hello from )" + model_sent +
      R"("}}]})";
    MockHttpServer::send_response(fd, 200, "application/json", resp_body);
  });

  std::vector<ModelConfig> configs = {
    {"model-1", "key-fail-1", server_a.url(), ""},
    {"model-1", "key-fail-2", server_b.url(), ""},
    {"model-2", "key-success", server_c.url(), ""},
  };
  ConfigManager cm(configs, {"model-1", "model-2"});

  Json req_body_obj = Json::object({
    {"model", "model-1"},
    {"messages", Json::array({
                   Json::object({{"role", "user"}, {"content", "Hi"}}),
                 })},
  });

  HttpRequest req;
  req.method = "POST";
  req.path = "/v1/chat/completions";
  req.headers.set("Content-Type", "application/json");
  req.body = req_body_obj.dump();

  RecorderResponseWriter w;
  handle_chat_completions(cm, {}, std::chrono::milliseconds(10), req, w);

  EXPECT_EQ(w.status_code(), 200, "Expected status 200 after failover");

  Json resp_obj = Json::parse(w.body());
  EXPECT_TRUE(resp_obj.contains("choices") && resp_obj["choices"].is_array() &&
                !resp_obj["choices"].empty(),
              "Invalid response choices structure");
  std::string content =
    resp_obj["choices"].at(0)["message"]["content"].str();
  EXPECT_EQ(content, "Hello from model-2",
            "Expected content 'Hello from model-2', got: " + content);
}

void TestGatewayCompletionsStreamingFailover() {
  MockHttpServer server_a([](const HttpRequest &, int fd) {
    MockHttpServer::send_response(fd, 500, "text/plain", "Error");
  });

  MockHttpServer server_b([](const HttpRequest &, int fd) {
    std::string headers =
      "HTTP/1.1 200 OK\r\n"
      "Content-Type: text/event-stream\r\n"
      "Cache-Control: no-cache\r\n"
      "Connection: close\r\n\r\n";
    write_all_fd(fd, headers);
    write_all_fd(fd, "data: chunk1\n\n");
    write_all_fd(fd, "data: chunk2\n\n");
  });

  std::vector<ModelConfig> configs = {
    {"model-1", "key-fail", server_a.url(), ""},
    {"model-2", "key-success", server_b.url(), ""},
  };
  ConfigManager cm(configs, {"model-1", "model-2"});

  Json req_body_obj = Json::object({
    {"model", "model-1"},
    {"stream", true},
    {"messages", Json::array({
                   Json::object({{"role", "user"}, {"content", "Hi"}}),
                 })},
  });

  HttpRequest req;
  req.method = "POST";
  req.path = "/v1/chat/completions";
  req.headers.set("Content-Type", "application/json");
  req.body = req_body_obj.dump();

  RecorderResponseWriter w;
  handle_chat_completions(cm, {}, std::chrono::milliseconds(10), req, w);

  EXPECT_EQ(w.status_code(), 200, "Expected status 200 on stream failover");
  EXPECT_TRUE(w.body().find("data: chunk1") != std::string::npos &&
                w.body().find("data: chunk2") != std::string::npos,
              "Expected streaming chunks to be forwarded, got: " + w.body());
}

void TestModelsEndpoint() {
  std::vector<ModelConfig> configs = {
    {"model-1", "key-1", "ep-1", ""},
    {"model-1", "key-2", "ep-1", ""},
    {"model-2", "key-3", "ep-2", ""},
  };
  ConfigManager cm(configs, {"model-1", "model-2"});

  HttpRequest req;
  req.method = "GET";
  req.path = "/v1/models";

  RecorderResponseWriter w;
  handle_models(cm, {}, req, w);

  EXPECT_EQ(w.status_code(), 200, "Expected status 200 on /v1/models");

  Json resp_obj = Json::parse(w.body());
  EXPECT_TRUE(resp_obj.contains("data") && resp_obj["data"].is_array() &&
                resp_obj["data"].size() == 2,
              "Expected 2 models in response data");
  EXPECT_EQ(resp_obj["data"].at(0)["id"].str(), "model-1",
            "Expected first model to be model-1");
  EXPECT_EQ(resp_obj["data"].at(1)["id"].str(), "model-2",
            "Expected second model to be model-2");
}

void TestGatewayAuthentication() {
  std::vector<ModelConfig> configs = {
    {"model-1", "key-1", "http://127.0.0.1:1", ""},
  };
  ConfigManager cm(configs, {"model-1"});
  std::vector<std::string> expected_tokens = {
    "secret-token-1",
    "secret-token-2",
  };

  // Case 1: No Authorization header
  {
    HttpRequest req;
    req.method = "GET";
    req.path = "/v1/models";
    RecorderResponseWriter w;
    handle_models(cm, expected_tokens, req, w);
    EXPECT_EQ(w.status_code(), 401,
              "Expected 401 Unauthorized for missing token");
  }

  // Case 2: Incorrect Authorization header
  {
    HttpRequest req;
    req.method = "GET";
    req.path = "/v1/models";
    req.headers.set("Authorization", "Bearer wrong-token");
    RecorderResponseWriter w;
    handle_models(cm, expected_tokens, req, w);
    EXPECT_EQ(w.status_code(), 401,
              "Expected 401 Unauthorized for wrong token");
  }

  // Case 3: Correct Authorization header (first token)
  {
    HttpRequest req;
    req.method = "GET";
    req.path = "/v1/models";
    req.headers.set("Authorization", "Bearer secret-token-1");
    RecorderResponseWriter w;
    handle_models(cm, expected_tokens, req, w);
    EXPECT_EQ(w.status_code(), 200, "Expected 200 OK for correct token 1");
  }

  // Case 3b: Correct Authorization header (second token)
  {
    HttpRequest req;
    req.method = "GET";
    req.path = "/v1/models";
    req.headers.set("Authorization", "Bearer secret-token-2");
    RecorderResponseWriter w;
    handle_models(cm, expected_tokens, req, w);
    EXPECT_EQ(w.status_code(), 200, "Expected 200 OK for correct token 2");
  }

  // Case 4: Completions endpoint with correct token
  {
    HttpRequest req;
    req.method = "POST";
    req.path = "/v1/chat/completions";
    req.headers.set("Authorization", "Bearer secret-token-2");
    req.headers.set("Content-Type", "application/json");
    req.body = R"({"model":"model-1"})";
    RecorderResponseWriter w;
    handle_chat_completions(cm, expected_tokens, std::chrono::milliseconds(0),
                            req, w);
    EXPECT_TRUE(w.status_code() != 401,
                "Expected completions request to pass authentication");
  }
}

void TestCSVUpdater() {
  TempDirGuard tmp("dynagate-updater-test");
  auto csv_file = (tmp.path / "models.csv").string();

  {
    std::ofstream ofs(csv_file, std::ios::binary);
    ofs << "model,key,endpoint\ngpt-4,key-1,http://localhost:8001\n";
  }

  std::string command =
    "echo 'gpt-3.5,key-2,http://localhost:8002' >> " + csv_file;

  std::atomic<bool> stop_flag{false};
  std::thread updater([&]() {
    start_csv_updater(stop_flag, command, std::chrono::milliseconds(100));
  });

  auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2);
  bool contains_new_model = false;
  while (std::chrono::steady_clock::now() < deadline) {
    std::ifstream ifs(csv_file);
    std::ostringstream oss;
    oss << ifs.rdbuf();
    if (oss.str().find("gpt-3.5") != std::string::npos) {
      contains_new_model = true;
      break;
    }
    std::this_thread::sleep_for(std::chrono::milliseconds(50));
  }

  stop_flag.store(true);
  if (updater.joinable()) {
    updater.join();
  }

  EXPECT_TRUE(contains_new_model,
              "Expected CSV file to contain updater-appended model");
}

void TestGatewayImageGenerations() {
  MockHttpServer server([](const HttpRequest &req, int fd) {
    if (req.path != "/v1/images/generations") {
      MockHttpServer::send_response(fd, 404, "text/plain", "Not Found");
      return;
    }
    Json body = Json::parse(req.body);
    std::string model_sent = body["model"].str();
    std::string resp =
      R"({"created":1589478378,"data":[{"url":"http://image-url/generated-by-)" +
      model_sent + R"("}]})";
    MockHttpServer::send_response(fd, 200, "application/json", resp);
  });

  std::vector<ModelConfig> configs = {
    {"dall-e-3", "dalle-key", server.url(), ""},
  };
  ConfigManager cm(configs, {"dall-e-3"});

  Json req_body = Json::object({
    {"prompt", "a beautiful kitten"},
    {"model", "dall-e-3"},
  });

  HttpRequest req;
  req.method = "POST";
  req.path = "/v1/images/generations";
  req.headers.set("Content-Type", "application/json");
  req.body = req_body.dump();

  RecorderResponseWriter w;
  handle_image_generations(cm, {}, std::chrono::milliseconds(0), req, w);

  EXPECT_EQ(w.status_code(), 200, "Expected 200 OK for image generations");
  Json resp_obj = Json::parse(w.body());
  std::string url_val = resp_obj["data"].at(0)["url"].str();
  EXPECT_EQ(url_val, "http://image-url/generated-by-dall-e-3",
            "Unexpected image URL: " + url_val);
}

void TestGatewayNoAuthHeaderOnEmptyKey() {
  std::atomic<bool> auth_was_empty{true};

  MockHttpServer server([&](const HttpRequest &req, int fd) {
    if (!req.headers.get("Authorization").empty()) {
      auth_was_empty.store(false);
      MockHttpServer::send_response(fd, 400, "text/plain", "Unexpected Auth");
      return;
    }
    MockHttpServer::send_response(
      fd, 200, "application/json",
      R"({"choices":[{"message":{"content":"Hello"}}]})");
  });

  std::vector<ModelConfig> configs = {
    {"model-no-key", "", server.url(), ""},
    {"model-no-key-2", "-", server.url(), ""},
  };
  ConfigManager cm(configs, {"model-no-key", "model-no-key-2"});

  // Case 1: Empty key
  {
    Json req_body = Json::object({
      {"model", "model-no-key"},
      {"messages", Json::array({
                     Json::object({{"role", "user"}, {"content", "Hi"}}),
                   })},
    });
    HttpRequest req;
    req.method = "POST";
    req.path = "/v1/chat/completions";
    req.headers.set("Content-Type", "application/json");
    req.headers.set("Authorization", "Bearer some-client-token");
    req.body = req_body.dump();

    RecorderResponseWriter w;
    handle_chat_completions(cm, {}, std::chrono::milliseconds(0), req, w);
    EXPECT_EQ(w.status_code(), 200, "Expected status 200 for empty key");
    EXPECT_TRUE(auth_was_empty.load(),
                "Expected no Authorization header on upstream request");
  }

  // Case 2: Dash key
  {
    Json req_body = Json::object({
      {"model", "model-no-key-2"},
      {"messages", Json::array({
                     Json::object({{"role", "user"}, {"content", "Hi"}}),
                   })},
    });
    HttpRequest req;
    req.method = "POST";
    req.path = "/v1/chat/completions";
    req.headers.set("Content-Type", "application/json");
    req.headers.set("Authorization", "Bearer some-client-token");
    req.body = req_body.dump();

    RecorderResponseWriter w;
    handle_chat_completions(cm, {}, std::chrono::milliseconds(0), req, w);
    EXPECT_EQ(w.status_code(), 200, "Expected status 200 for dash key");
    EXPECT_TRUE(auth_was_empty.load(),
                "Expected no Authorization header on upstream request");
  }
}

void TestGatewayBlankKey() {
  std::string received_auth;
  std::mutex mu;

  MockHttpServer server([&](const HttpRequest &req, int fd) {
    {
      std::lock_guard<std::mutex> lock(mu);
      received_auth = req.headers.get("Authorization");
    }
    if (received_auth != "Bearer") {
      MockHttpServer::send_response(fd, 400, "text/plain", "Bad Auth");
      return;
    }
    MockHttpServer::send_response(
      fd, 200, "application/json",
      R"({"choices":[{"message":{"content":"Hello"}}]})");
  });

  std::vector<ModelConfig> configs = {
    {"model-blank-key", "-blank-", server.url(), ""},
  };
  ConfigManager cm(configs, {"model-blank-key"});

  Json req_body = Json::object({
    {"model", "model-blank-key"},
    {"messages", Json::array({
                   Json::object({{"role", "user"}, {"content", "Hi"}}),
                 })},
  });
  HttpRequest req;
  req.method = "POST";
  req.path = "/v1/chat/completions";
  req.headers.set("Content-Type", "application/json");
  req.headers.set("Authorization", "Bearer some-client-token");
  req.body = req_body.dump();

  RecorderResponseWriter w;
  handle_chat_completions(cm, {}, std::chrono::milliseconds(0), req, w);
  EXPECT_EQ(w.status_code(), 200, "Expected status 200 for -blank- key");
  EXPECT_EQ(received_auth, "Bearer",
            "Expected Authorization header 'Bearer', got: " + received_auth);
}

void TestGatewayExtraHeaders() {
  MockHttpServer server([](const HttpRequest &req, int fd) {
    if (req.headers.get("X-My-Custom-Header") != "HeaderValue" ||
        req.headers.get("X-Another-Header") != "123") {
      MockHttpServer::send_response(fd, 400, "text/plain", "Missing headers");
      return;
    }
    MockHttpServer::send_response(
      fd, 200, "application/json",
      R"({"choices":[{"message":{"content":"Hello"}}]})");
  });

  std::vector<ModelConfig> configs = {
    {"model-extra-headers", "my-key", server.url(),
     R"({"X-My-Custom-Header":"HeaderValue","X-Another-Header":123})"},
  };
  ConfigManager cm(configs, {"model-extra-headers"});

  Json req_body = Json::object({
    {"model", "model-extra-headers"},
    {"messages", Json::array({
                   Json::object({{"role", "user"}, {"content", "Hi"}}),
                 })},
  });
  HttpRequest req;
  req.method = "POST";
  req.path = "/v1/chat/completions";
  req.headers.set("Content-Type", "application/json");
  req.body = req_body.dump();

  RecorderResponseWriter w;
  handle_chat_completions(cm, {}, std::chrono::milliseconds(0), req, w);
  EXPECT_EQ(w.status_code(), 200, "Expected status 200 for extra headers");
}

void TestMultipleTokensKeyFile() {
  std::string key_file_content = "\n  token-a  \n\ntoken-b\n   \ntoken-c\n";
  auto auth_tokens = parse_auth_tokens(key_file_content);
  std::vector<std::string> expected = {"token-a", "token-b", "token-c"};
  EXPECT_EQ(auth_tokens, expected, "Parsed auth tokens mismatch");

  HttpRequest req_ok;
  req_ok.method = "GET";
  req_ok.path = "/v1/models";
  req_ok.headers.set("Authorization", "Bearer token-b");
  EXPECT_TRUE(check_auth(auth_tokens, req_ok),
              "Expected token-b to authenticate");

  HttpRequest req_wrong;
  req_wrong.method = "GET";
  req_wrong.path = "/v1/models";
  req_wrong.headers.set("Authorization", "Bearer token-wrong");
  EXPECT_TRUE(!check_auth(auth_tokens, req_wrong),
              "Expected token-wrong to fail authentication");

  auto empty_tokens = parse_auth_tokens("\n   \n\n  \n");
  EXPECT_TRUE(empty_tokens.empty(), "Expected empty tokens vector");
  HttpRequest req_no_auth;
  req_no_auth.method = "GET";
  req_no_auth.path = "/v1/models";
  EXPECT_TRUE(check_auth(empty_tokens, req_no_auth),
              "Expected check_auth to return true for empty expected_tokens");
}

void TestRequestMessagesAutodetectionAndContentArrays() {
  Json received_body;
  std::mutex mu;

  MockHttpServer server([&](const HttpRequest &req, int fd) {
    {
      std::lock_guard<std::mutex> lock(mu);
      received_body = Json::parse(req.body);
    }
    MockHttpServer::send_response(
      fd, 200, "application/json",
      R"({"choices":[{"message":{"content":"OK"}}]})");
  });

  std::vector<ModelConfig> configs = {
    {"test-model", "test-key", server.url(), ""},
  };
  ConfigManager cm(configs, {"test-model"});

  // 1. Pi agent format - array content
  {
    Json req_obj = Json::object({
      {"model", "test-model"},
      {"messages", Json::array({
                     Json::object({
                       {"role", "user"},
                       {"content", Json::array({
                                     Json::object({
                                       {"type", "text"},
                                       {"text", "Hello from Pi agent"},
                                     }),
                                   })},
                     }),
                   })},
    });
    HttpRequest req;
    req.method = "POST";
    req.path = "/v1/chat/completions";
    req.headers.set("Content-Type", "application/json");
    req.body = req_obj.dump();

    RecorderResponseWriter w;
    handle_chat_completions(cm, {}, std::chrono::milliseconds(0), req, w);
    EXPECT_EQ(w.status_code(), 200, "Expected 200 OK for Pi agent format");
    EXPECT_EQ(received_body["messages"].size(), 1u, "Expected 1 message");
    EXPECT_EQ(received_body["messages"].at(0)["content"].str(),
              "Hello from Pi agent", "Unexpected Pi agent normalized content");
  }

  // 2. Top-level system prompt
  {
    Json req_obj = Json::object({
      {"model", "test-model"},
      {"system", "You are a helpful coding assistant"},
      {"messages", Json::array({
                     Json::object({
                       {"role", "user"},
                       {"content", "Write a test"},
                     }),
                   })},
    });
    HttpRequest req;
    req.method = "POST";
    req.path = "/v1/chat/completions";
    req.headers.set("Content-Type", "application/json");
    req.body = req_obj.dump();

    RecorderResponseWriter w;
    handle_chat_completions(cm, {}, std::chrono::milliseconds(0), req, w);
    EXPECT_EQ(received_body["messages"].size(), 2u, "Expected 2 messages");
    EXPECT_EQ(received_body["messages"].at(0)["role"].str(),
              "system", "Expected system role first");
    EXPECT_EQ(received_body["messages"].at(0)["content"].str(),
              "You are a helpful coding assistant",
              "Unexpected system message content");
    EXPECT_EQ(received_body["messages"].at(1)["content"].str(),
              "Write a test", "Unexpected user message content");
    EXPECT_TRUE(!received_body.contains("system"),
                "Expected top-level system field to be removed");
  }

  // 3. Prompt field autodetection
  {
    Json req_obj = Json::object({
      {"model", "test-model"},
      {"prompt", "Explain recursion"},
    });
    HttpRequest req;
    req.method = "POST";
    req.path = "/v1/chat/completions";
    req.headers.set("Content-Type", "application/json");
    req.body = req_obj.dump();

    RecorderResponseWriter w;
    handle_chat_completions(cm, {}, std::chrono::milliseconds(0), req, w);
    EXPECT_EQ(received_body["messages"].size(), 1u, "Expected 1 message");
    EXPECT_EQ(received_body["messages"].at(0)["role"].str(), "user",
              "Expected user role");
    EXPECT_EQ(received_body["messages"].at(0)["content"].str(),
              "Explain recursion", "Unexpected prompt autodetection content");
  }

  // 4. Multimodal array content preserved
  {
    Json req_obj = Json::object({
      {"model", "test-model"},
      {"messages", Json::array({
                     Json::object({
                       {"role", "user"},
                       {"content",
                        Json::array({
                          Json::object({
                            {"type", "text"},
                            {"text", "Describe this image"},
                          }),
                          Json::object({
                            {"type", "image_url"},
                            {"image_url",
                             Json::object({
                               {"url", "http://example.com/img.png"},
                             })},
                          }),
                        })},
                     }),
                   })},
    });
    HttpRequest req;
    req.method = "POST";
    req.path = "/v1/chat/completions";
    req.headers.set("Content-Type", "application/json");
    req.body = req_obj.dump();

    RecorderResponseWriter w;
    handle_chat_completions(cm, {}, std::chrono::milliseconds(0), req, w);
    EXPECT_TRUE(received_body["messages"].at(0)["content"].is_array() &&
                  received_body["messages"].at(0)["content"].size() == 2,
                "Expected multimodal array content to be preserved");
  }
}

void TestFibonacciBackoff() {
  std::vector<std::pair<int, int64_t>> expected_seq = {
    {0, 0},  {1, 1},  {2, 1},  {3, 2},  {4, 3},  {5, 5},
    {6, 8},  {7, 13}, {8, 21}, {9, 34}, {10, 55},
  };
  for (const auto &p : expected_seq) {
    EXPECT_EQ(fibonacci(p.first), p.second, "fibonacci mismatch");
  }
}

void TestAnthropicMessagesEndpoint() {
  Json received_openai_body;
  std::mutex mu;

  MockHttpServer server([&](const HttpRequest &req, int fd) {
    if (req.path != "/v1/chat/completions") {
      MockHttpServer::send_response(fd, 404, "text/plain", "Wrong path");
      return;
    }
    Json body = Json::parse(req.body);
    {
      std::lock_guard<std::mutex> lock(mu);
      received_openai_body = body;
    }
    bool is_stream = body["stream"].as_bool(false);
    if (is_stream) {
      std::string headers =
        "HTTP/1.1 200 OK\r\n"
        "Content-Type: text/event-stream\r\n"
        "Connection: close\r\n\r\n";
      write_all_fd(fd, headers);
      write_all_fd(
        fd, "data: {\"choices\":[{\"delta\":{\"content\":\"Hello \"}}]}\n\n");
      write_all_fd(
        fd,
        "data: {\"choices\":[{\"delta\":{\"content\":\"Anthropic!\"},\"finish_reason\":\"stop\"}]}\n\n");
      write_all_fd(fd, "data: [DONE]\n\n");
      return;
    }

    Json resp = Json::object({
      {"id", "chatcmpl-test"},
      {"model", "claude-model"},
      {"choices", Json::array({
                    Json::object({
                      {"index", 0},
                      {"message", Json::object({
                                    {"role", "assistant"},
                                    {"content", "Hello from Anthropic proxy"},
                                  })},
                      {"finish_reason", "stop"},
                    }),
                  })},
      {"usage", Json::object({
                  {"prompt_tokens", 10},
                  {"completion_tokens", 15},
                  {"total_tokens", 25},
                })},
    });
    MockHttpServer::send_response(fd, 200, "application/json", resp.dump());
  });

  std::vector<ModelConfig> configs = {
    {"claude-model", "test-key", server.url(), ""},
  };
  ConfigManager cm(configs, {"claude-model"});

  // 1. Non-streaming test
  {
    Json req_obj = Json::object({
      {"model", "claude-model"},
      {"system", "You are a helpful bot"},
      {"messages", Json::array({
                     Json::object({
                       {"role", "user"},
                       {"content", "Hello"},
                     }),
                   })},
      {"max_tokens", 100},
    });
    HttpRequest req;
    req.method = "POST";
    req.path = "/v1/messages";
    req.headers.set("Content-Type", "application/json");
    req.body = req_obj.dump();

    RecorderResponseWriter w;
    handle_messages(cm, {}, std::chrono::milliseconds(0), req, w);
    EXPECT_EQ(w.status_code(), 200, "Expected 200 OK for Anthropic messages");

    Json anthropic_resp = Json::parse(w.body());
    EXPECT_EQ(anthropic_resp["type"].str(), "message",
              "Expected type 'message'");
    EXPECT_EQ(anthropic_resp["role"].str(), "assistant",
              "Expected role 'assistant'");
    EXPECT_EQ(anthropic_resp["content"].at(0)["text"].str(),
              "Hello from Anthropic proxy", "Unexpected Anthropic text");
    EXPECT_EQ(received_openai_body["messages"].size(), 2u,
              "Expected 2 messages in OpenAI body");
    EXPECT_EQ(received_openai_body["messages"].at(0)["role"].str(),
              "system", "Expected first message to be system");
  }

  // 2. Streaming test
  {
    Json req_obj = Json::object({
      {"model", "claude-model"},
      {"stream", true},
      {"messages", Json::array({
                     Json::object({
                       {"role", "user"},
                       {"content", "Stream me"},
                     }),
                   })},
    });
    HttpRequest req;
    req.method = "POST";
    req.path = "/v1/messages";
    req.headers.set("Content-Type", "application/json");
    req.body = req_obj.dump();

    RecorderResponseWriter w;
    handle_messages(cm, {}, std::chrono::milliseconds(0), req, w);
    EXPECT_EQ(w.status_code(), 200,
              "Expected 200 OK for streaming Anthropic messages");
    EXPECT_TRUE(w.body().find("event: message_start") != std::string::npos &&
                  w.body().find("event: content_block_delta") !=
                    std::string::npos &&
                  w.body().find("event: message_stop") != std::string::npos,
                "Expected Anthropic SSE stream events in response");
  }
}

void TestOpenAIResponsesEndpoint() {
  Json received_openai_body;
  std::mutex mu;

  MockHttpServer server([&](const HttpRequest &req, int fd) {
    if (req.path != "/v1/chat/completions") {
      MockHttpServer::send_response(fd, 404, "text/plain", "Wrong path");
      return;
    }
    Json body = Json::parse(req.body);
    {
      std::lock_guard<std::mutex> lock(mu);
      received_openai_body = body;
    }
    bool is_stream = body["stream"].as_bool(false);
    if (is_stream) {
      std::string headers =
        "HTTP/1.1 200 OK\r\n"
        "Content-Type: text/event-stream\r\n"
        "Connection: close\r\n\r\n";
      write_all_fd(fd, headers);
      write_all_fd(
        fd,
        "data: {\"choices\":[{\"delta\":{\"content\":\"Response \"}}]}\n\n");
      write_all_fd(
        fd,
        "data: {\"choices\":[{\"delta\":{\"content\":\"Chunk\"},\"finish_reason\":\"stop\"}]}\n\n");
      write_all_fd(fd, "data: [DONE]\n\n");
      return;
    }

    Json resp = Json::object({
      {"id", "chatcmpl-resp-test"},
      {"model", "gpt-5-model"},
      {"choices", Json::array({
                    Json::object({
                      {"index", 0},
                      {"message", Json::object({
                                    {"role", "assistant"},
                                    {"content", "Hello from Responses proxy"},
                                  })},
                      {"finish_reason", "stop"},
                    }),
                  })},
      {"usage", Json::object({
                  {"prompt_tokens", 12},
                  {"completion_tokens", 18},
                  {"total_tokens", 30},
                })},
    });
    MockHttpServer::send_response(fd, 200, "application/json", resp.dump());
  });

  std::vector<ModelConfig> configs = {
    {"gpt-5-model", "test-key", server.url(), ""},
  };
  ConfigManager cm(configs, {"gpt-5-model"});

  // 1. Non-streaming test
  {
    Json req_obj = Json::object({
      {"model", "gpt-5-model"},
      {"instructions", "Be accurate"},
      {"input", "What is 2+2?"},
    });
    HttpRequest req;
    req.method = "POST";
    req.path = "/v1/responses";
    req.headers.set("Content-Type", "application/json");
    req.body = req_obj.dump();

    RecorderResponseWriter w;
    handle_responses(cm, {}, std::chrono::milliseconds(0), req, w);
    EXPECT_EQ(w.status_code(), 200, "Expected 200 OK for OpenAI Responses");

    Json resp_obj = Json::parse(w.body());
    EXPECT_EQ(resp_obj["object"].str(), "response",
              "Expected object 'response'");
    EXPECT_EQ(resp_obj["status"].str(), "completed",
              "Expected status 'completed'");
    EXPECT_EQ(
      resp_obj["output"].at(0)["content"].at(0)["text"].str(),
      "Hello from Responses proxy", "Unexpected Responses output text");
    EXPECT_EQ(received_openai_body["messages"].size(), 2u,
              "Expected 2 messages in OpenAI body");
    EXPECT_EQ(received_openai_body["messages"].at(0)["role"].str(),
              "system", "Expected system role for instructions");
  }

  // 2. Streaming test
  {
    Json req_obj = Json::object({
      {"model", "gpt-5-model"},
      {"stream", true},
      {"input", "Stream test"},
    });
    HttpRequest req;
    req.method = "POST";
    req.path = "/v1/responses";
    req.headers.set("Content-Type", "application/json");
    req.body = req_obj.dump();

    RecorderResponseWriter w;
    handle_responses(cm, {}, std::chrono::milliseconds(0), req, w);
    EXPECT_EQ(w.status_code(), 200,
              "Expected 200 OK for streaming OpenAI Responses");
    EXPECT_TRUE(
      w.body().find("event: response.created") != std::string::npos &&
        w.body().find("event: response.output_text.delta") !=
          std::string::npos &&
        w.body().find("event: response.completed") != std::string::npos,
      "Expected OpenAI Responses SSE events in response stream");
  }
}

void TestFullGatewayServerEndToEnd() {
  MockHttpServer upstream([](const HttpRequest &req, int fd) {
    Json body = Json::parse(req.body);
    std::string model_sent = body["model"].str();
    Json resp = Json::object({
      {"id", "chatcmpl-e2e"},
      {"model", model_sent},
      {"choices", Json::array({
                    Json::object({
                      {"index", 0},
                      {"message", Json::object({
                                    {"role", "assistant"},
                                    {"content", "E2E OK from " + model_sent},
                                  })},
                      {"finish_reason", "stop"},
                    }),
                  })},
    });
    MockHttpServer::send_response(fd, 200, "application/json", resp.dump());
  });

  std::vector<ModelConfig> configs = {
    {"e2e-model", "e2e-key", upstream.url(), ""},
  };
  ConfigManager cm(configs, {"e2e-model"});

  GatewayServer gw(0, cm, {"gw-token"}, std::chrono::milliseconds(0));
  std::string err;
  EXPECT_TRUE(gw.start(&err), "GatewayServer start failed: " + err);
  std::thread gw_thread([&]() { gw.serve(); });

  std::string base_url = "http://127.0.0.1:" + std::to_string(gw.bound_port());

  // Perform a real HTTP GET /v1/models via libcurl
  {
    CURL *c = curl_easy_init();
    std::string resp_body;
    struct curl_slist *hdrs = nullptr;
    hdrs = curl_slist_append(hdrs, "Authorization: Bearer gw-token");
    std::string url = base_url + "/v1/models";
    curl_easy_setopt(c, CURLOPT_URL, url.c_str());
    curl_easy_setopt(c, CURLOPT_HTTPHEADER, hdrs);
    curl_easy_setopt(
      c, CURLOPT_WRITEFUNCTION,
      +[](char *p, size_t s, size_t n, void *ud) -> size_t {
        static_cast<std::string *>(ud)->append(p, s * n);
        return s * n;
      });
    curl_easy_setopt(c, CURLOPT_WRITEDATA, &resp_body);
    CURLcode rc = curl_easy_perform(c);
    long code = 0;
    curl_easy_getinfo(c, CURLINFO_RESPONSE_CODE, &code);
    curl_slist_free_all(hdrs);
    curl_easy_cleanup(c);

    EXPECT_EQ(rc, CURLE_OK, "curl /v1/models failed");
    EXPECT_EQ(code, 200L, "Expected 200 from GatewayServer /v1/models");
    Json j = Json::parse(resp_body);
    EXPECT_EQ(j["data"].at(0)["id"].str(), "e2e-model",
              "Unexpected model from GatewayServer");
  }

  // Perform a real HTTP POST /v1/chat/completions via libcurl
  {
    CURL *c = curl_easy_init();
    std::string resp_body;
    struct curl_slist *hdrs = nullptr;
    hdrs = curl_slist_append(hdrs, "Authorization: Bearer gw-token");
    hdrs = curl_slist_append(hdrs, "Content-Type: application/json");
    std::string payload =
      R"({"model":"e2e-model","messages":[{"role":"user","content":"ping"}]})";
    std::string url = base_url + "/v1/chat/completions";
    curl_easy_setopt(c, CURLOPT_URL, url.c_str());
    curl_easy_setopt(c, CURLOPT_POST, 1L);
    curl_easy_setopt(c, CURLOPT_POSTFIELDS, payload.c_str());
    curl_easy_setopt(c, CURLOPT_HTTPHEADER, hdrs);
    curl_easy_setopt(
      c, CURLOPT_WRITEFUNCTION,
      +[](char *p, size_t s, size_t n, void *ud) -> size_t {
        static_cast<std::string *>(ud)->append(p, s * n);
        return s * n;
      });
    curl_easy_setopt(c, CURLOPT_WRITEDATA, &resp_body);
    CURLcode rc = curl_easy_perform(c);
    long code = 0;
    curl_easy_getinfo(c, CURLINFO_RESPONSE_CODE, &code);
    curl_slist_free_all(hdrs);
    curl_easy_cleanup(c);

    EXPECT_EQ(rc, CURLE_OK, "curl /v1/chat/completions failed");
    EXPECT_EQ(code, 200L,
              "Expected 200 from GatewayServer /v1/chat/completions");
    Json j = Json::parse(resp_body);
    EXPECT_EQ(j["choices"].at(0)["message"]["content"].str(),
              "E2E OK from e2e-model", "Unexpected E2E chat completion");
  }

  gw.stop();
  if (gw_thread.joinable()) {
    gw_thread.join();
  }
}

} // namespace

int main() {
  std::signal(SIGPIPE, SIG_IGN);
  ensure_curl_initialized();

  struct TestEntry {
    const char *name;
    void (*fn)();
  };

  std::vector<TestEntry> tests = {
    {"TestURLBuilding", TestURLBuilding},
    {"TestTrialOrdering", TestTrialOrdering},
    {"TestConfigManagerCSVWatcher", TestConfigManagerCSVWatcher},
    {"TestGatewayCompletionsFailover", TestGatewayCompletionsFailover},
    {"TestGatewayCompletionsStreamingFailover",
     TestGatewayCompletionsStreamingFailover},
    {"TestModelsEndpoint", TestModelsEndpoint},
    {"TestGatewayAuthentication", TestGatewayAuthentication},
    {"TestCSVUpdater", TestCSVUpdater},
    {"TestGatewayImageGenerations", TestGatewayImageGenerations},
    {"TestGatewayNoAuthHeaderOnEmptyKey", TestGatewayNoAuthHeaderOnEmptyKey},
    {"TestGatewayBlankKey", TestGatewayBlankKey},
    {"TestGatewayExtraHeaders", TestGatewayExtraHeaders},
    {"TestMultipleTokensKeyFile", TestMultipleTokensKeyFile},
    {"TestRequestMessagesAutodetectionAndContentArrays",
     TestRequestMessagesAutodetectionAndContentArrays},
    {"TestFibonacciBackoff", TestFibonacciBackoff},
    {"TestAnthropicMessagesEndpoint", TestAnthropicMessagesEndpoint},
    {"TestOpenAIResponsesEndpoint", TestOpenAIResponsesEndpoint},
    {"TestFullGatewayServerEndToEnd", TestFullGatewayServerEndToEnd},
  };

  for (const auto &t : tests) {
    int before = g_failures;
    std::printf("=== RUN   %s\n", t.name);
    t.fn();
    if (g_failures == before) {
      std::printf("--- PASS: %s\n", t.name);
    } else {
      std::printf("--- FAIL: %s\n", t.name);
    }
  }

  if (g_failures == 0) {
    std::printf("PASS\n");
    return 0;
  }
  std::printf("FAIL (%d failures)\n", g_failures);
  return 1;
}