/*
  Dynagate HTTP/1.1 server implementation (C++20 port)
  Created by Luxferre in 2026, released into the public domain
*/

#ifndef _DARWIN_C_SOURCE
#define _DARWIN_C_SOURCE
#endif

#include "dynagate.hpp"

#include <algorithm>
#include <cctype>
#include <cerrno>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <thread>

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

namespace {

constexpr size_t kMaxHeaderBytes = 65536;
constexpr size_t kMaxBodyBytes = 64 * 1024 * 1024; // 64 MB

bool write_all(int fd, const char *data, size_t len) {
  size_t written = 0;
  while (written < len) {
    ssize_t n = write(fd, data + written, len - written);
    if (n > 0) {
      written += static_cast<size_t>(n);
    } else if (n < 0 && errno == EINTR) {
      continue;
    } else {
      return false;
    }
  }
  return true;
}

const char *status_reason_phrase(int code) {
  switch (code) {
  case 200:
    return "OK";
  case 201:
    return "Created";
  case 202:
    return "Accepted";
  case 204:
    return "No Content";
  case 400:
    return "Bad Request";
  case 401:
    return "Unauthorized";
  case 403:
    return "Forbidden";
  case 404:
    return "Not Found";
  case 405:
    return "Method Not Allowed";
  case 408:
    return "Request Timeout";
  case 413:
    return "Payload Too Large";
  case 429:
    return "Too Many Requests";
  case 499:
    return "Client Closed Request";
  case 500:
    return "Internal Server Error";
  case 502:
    return "Bad Gateway";
  case 503:
    return "Service Unavailable";
  case 504:
    return "Gateway Timeout";
  default:
    return "Status";
  }
}

class SocketResponseWriter : public ResponseWriter {
public:
  explicit SocketResponseWriter(int fd) : fd_(fd) {}

  HeaderList &header() override { return headers_; }

  void write_header(int status_code) override {
    if (!header_written_) {
      status_code_ = status_code;
      header_written_ = true;
    }
  }

  bool write(std::string_view data) override {
    if (!header_written_) {
      write_header(200);
    }
    write_called_ = true;
    if (headers_sent_to_socket_) {
      if (data.empty()) {
        return !write_error_;
      }
      if (!write_all(fd_, data.data(), data.size())) {
        write_error_ = true;
        return false;
      }
      return true;
    }
    pending_body_.append(data.data(), data.size());
    return true;
  }

  void flush() override {
    if (write_error_) {
      return;
    }
    if (!header_written_) {
      write_header(200);
    }
    if (!headers_sent_to_socket_) {
      headers_sent_to_socket_ = true;
      std::string raw;
      raw.reserve(512 + pending_body_.size());
      raw.append("HTTP/1.1 ");
      raw.append(std::to_string(status_code_));
      raw.push_back(' ');
      raw.append(status_reason_phrase(status_code_));
      raw.append("\r\n");

      for (const auto &kv : headers_.items) {
        raw.append(kv.first);
        raw.append(": ");
        raw.append(kv.second);
        raw.append("\r\n");
      }

      if (write_called_) {
        if (headers_.get("Content-Length").empty()) {
          raw.append("Content-Length: ");
          raw.append(std::to_string(pending_body_.size()));
          raw.append("\r\n");
        }
        if (headers_.get("Connection").empty()) {
          raw.append("Connection: close\r\n");
        }
      } else {
        if (headers_.get("Connection").empty()) {
          raw.append("Connection: close\r\n");
        }
      }

      raw.append("\r\n");
      if (!pending_body_.empty()) {
        raw.append(pending_body_);
        pending_body_.clear();
      }
      if (!write_all(fd_, raw.data(), raw.size())) {
        write_error_ = true;
      }
    }
  }

private:
  int fd_;
  int status_code_ = 200;
  bool header_written_ = false;
  bool write_called_ = false;
  bool headers_sent_to_socket_ = false;
  bool write_error_ = false;
  HeaderList headers_;
  std::string pending_body_;
};

bool is_socket_disconnected(int fd) {
  if (fd < 0) {
    return true;
  }
  struct pollfd pfd;
  pfd.fd = fd;
  pfd.events = POLLIN | POLLHUP | POLLERR;
  pfd.revents = 0;
  int rc = poll(&pfd, 1, 0);
  if (rc < 0) {
    return errno != EINTR;
  }
  if (rc == 0) {
    return false;
  }
  if (pfd.revents & (POLLHUP | POLLERR | POLLNVAL)) {
    return true;
  }
  if (pfd.revents & POLLIN) {
    char ch = 0;
    ssize_t n = recv(fd, &ch, 1, MSG_PEEK);
    if (n == 0) {
      return true;
    }
    if (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK && errno != EINTR) {
      return true;
    }
  }
  return false;
}

bool read_chunked_body(int fd,
                       std::string &buffer,
                       size_t initial_offset,
                       std::string &out_body) {
  size_t pos = initial_offset;
  auto ensure_bytes = [&](size_t needed_end) -> bool {
    char tmp[4096];
    while (buffer.size() < needed_end) {
      ssize_t n = read(fd, tmp, sizeof(tmp));
      if (n > 0) {
        buffer.append(tmp, static_cast<size_t>(n));
        if (buffer.size() > kMaxBodyBytes + kMaxHeaderBytes) {
          return false;
        }
      } else if (n < 0 && errno == EINTR) {
        continue;
      } else {
        return false;
      }
    }
    return true;
  };

  auto ensure_crlf = [&](size_t from, size_t &crlf_pos) -> bool {
    char tmp[4096];
    for (;;) {
      crlf_pos = buffer.find("\r\n", from);
      if (crlf_pos != std::string::npos) {
        return true;
      }
      ssize_t n = read(fd, tmp, sizeof(tmp));
      if (n > 0) {
        buffer.append(tmp, static_cast<size_t>(n));
        if (buffer.size() > kMaxBodyBytes + kMaxHeaderBytes) {
          return false;
        }
      } else if (n < 0 && errno == EINTR) {
        continue;
      } else {
        return false;
      }
    }
  };

  for (;;) {
    size_t line_end = 0;
    if (!ensure_crlf(pos, line_end)) {
      return false;
    }
    std::string size_str = str_trim(
      std::string_view(buffer).substr(pos, line_end - pos));
    size_t semi = size_str.find(';');
    if (semi != std::string::npos) {
      size_str.resize(semi);
    }
    if (size_str.empty()) {
      return false;
    }
    char *endptr = nullptr;
    unsigned long chunk_len = std::strtoul(size_str.c_str(), &endptr, 16);
    if (endptr == size_str.c_str() || chunk_len > kMaxBodyBytes) {
      return false;
    }
    pos = line_end + 2;
    if (chunk_len == 0) {
      size_t trailer_end = 0;
      if (!ensure_crlf(pos, trailer_end)) {
        return false;
      }
      return true;
    }
    if (out_body.size() + chunk_len > kMaxBodyBytes) {
      return false;
    }
    if (!ensure_bytes(pos + chunk_len + 2)) {
      return false;
    }
    out_body.append(buffer.data() + pos, chunk_len);
    pos += chunk_len;
    if (buffer.substr(pos, 2) != "\r\n") {
      return false;
    }
    pos += 2;
  }
}

bool read_http_request(int fd, HttpRequest &req, int &err_status) {
  err_status = 400;
  std::string buffer;
  buffer.reserve(4096);
  char tmp[4096];
  size_t header_end = std::string::npos;

  for (;;) {
    header_end = buffer.find("\r\n\r\n");
    if (header_end != std::string::npos) {
      break;
    }
    if (buffer.size() >= kMaxHeaderBytes) {
      return false;
    }
    ssize_t n = read(fd, tmp, sizeof(tmp));
    if (n > 0) {
      buffer.append(tmp, static_cast<size_t>(n));
    } else if (n < 0 && errno == EINTR) {
      continue;
    } else {
      return false;
    }
  }

  std::string_view header_block(buffer.data(), header_end);
  size_t first_line_end = header_block.find("\r\n");
  std::string_view request_line =
    (first_line_end == std::string_view::npos)
      ? header_block
      : header_block.substr(0, first_line_end);

  size_t sp1 = request_line.find(' ');
  if (sp1 == std::string_view::npos) {
    return false;
  }
  req.method = std::string(request_line.substr(0, sp1));

  std::string_view rest = request_line.substr(sp1 + 1);
  size_t sp2 = rest.find(' ');
  std::string_view raw_target =
    (sp2 == std::string_view::npos) ? rest : rest.substr(0, sp2);

  size_t qpos = raw_target.find('?');
  req.path = std::string(
    (qpos == std::string_view::npos) ? raw_target : raw_target.substr(0, qpos));

  size_t pos = (first_line_end == std::string_view::npos)
                 ? header_block.size()
                 : first_line_end + 2;
  long content_length = 0;
  bool is_chunked = false;
  bool expect_continue = false;

  while (pos < header_block.size()) {
    size_t le = header_block.find("\r\n", pos);
    std::string_view line =
      (le == std::string_view::npos)
        ? header_block.substr(pos)
        : header_block.substr(pos, le - pos);
    pos = (le == std::string_view::npos) ? header_block.size() : le + 2;

    if (line.empty()) {
      continue;
    }
    size_t colon = line.find(':');
    if (colon == std::string_view::npos) {
      continue;
    }
    std::string key = str_trim(line.substr(0, colon));
    std::string val = str_trim(line.substr(colon + 1));
    if (key.empty()) {
      continue;
    }
    if (str_iequals(key, "Content-Length")) {
      content_length = std::atol(val.c_str());
    } else if (str_iequals(key, "Transfer-Encoding")) {
      if (str_to_lower(val).find("chunked") != std::string::npos) {
        is_chunked = true;
      }
    } else if (str_iequals(key, "Expect")) {
      if (str_iequals(val, "100-continue")) {
        expect_continue = true;
      }
    }
    req.headers.add(key, val);
  }

  if (expect_continue) {
    const char *cont = "HTTP/1.1 100 Continue\r\n\r\n";
    write_all(fd, cont, std::strlen(cont));
  }

  size_t body_offset = header_end + 4;
  if (is_chunked) {
    if (!read_chunked_body(fd, buffer, body_offset, req.body)) {
      return false;
    }
    return true;
  }

  if (content_length < 0) {
    return false;
  }
  if (static_cast<size_t>(content_length) > kMaxBodyBytes) {
    err_status = 413;
    return false;
  }

  size_t needed = static_cast<size_t>(content_length);
  size_t avail = buffer.size() - body_offset;
  req.body.reserve(needed);
  size_t to_copy = std::min(avail, needed);
  if (to_copy > 0) {
    req.body.append(buffer.data() + body_offset, to_copy);
  }

  while (req.body.size() < needed) {
    size_t rem = needed - req.body.size();
    ssize_t n = read(fd, tmp, std::min(sizeof(tmp), rem));
    if (n > 0) {
      req.body.append(tmp, static_cast<size_t>(n));
    } else if (n < 0 && errno == EINTR) {
      continue;
    } else {
      return false;
    }
  }

  return true;
}

} // namespace

GatewayServer::GatewayServer(int port,
                             ConfigManager &cm,
                             std::vector<std::string> auth_tokens,
                             std::chrono::milliseconds retry_delay)
  : requested_port_(port),
    cm_(cm),
    auth_tokens_(std::move(auth_tokens)),
    retry_delay_(retry_delay) {}

GatewayServer::~GatewayServer() {
  stop();
}

bool GatewayServer::start(std::string *err_out) {
  ensure_curl_initialized();

  int fd = socket(AF_INET, SOCK_STREAM, 0);
  if (fd < 0) {
    if (err_out) {
      *err_out = "socket() failed: " + std::string(std::strerror(errno));
    }
    return false;
  }

  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_ANY);
  addr.sin_port = htons(static_cast<uint16_t>(requested_port_));

  if (bind(fd, reinterpret_cast<struct sockaddr *>(&addr), sizeof(addr)) < 0) {
    if (err_out) {
      *err_out = "bind() failed: " + std::string(std::strerror(errno));
    }
    close(fd);
    return false;
  }

  if (listen(fd, 128) < 0) {
    if (err_out) {
      *err_out = "listen() failed: " + std::string(std::strerror(errno));
    }
    close(fd);
    return false;
  }

  struct sockaddr_in bound_addr;
  socklen_t addr_len = sizeof(bound_addr);
  std::memset(&bound_addr, 0, sizeof(bound_addr));
  if (getsockname(fd, reinterpret_cast<struct sockaddr *>(&bound_addr),
                  &addr_len) == 0) {
    bound_port_ = ntohs(bound_addr.sin_port);
  } else {
    bound_port_ = requested_port_;
  }

  listen_fd_ = fd;
  stopping_.store(false);
  return true;
}

void GatewayServer::serve() {
  while (!stopping_.load()) {
    int lfd = listen_fd_;
    if (lfd < 0) {
      break;
    }

    struct pollfd pfd;
    pfd.fd = lfd;
    pfd.events = POLLIN;
    pfd.revents = 0;

    int prc = poll(&pfd, 1, 100);
    if (prc < 0) {
      if (errno == EINTR) {
        continue;
      }
      break;
    }
    if (prc == 0 || !(pfd.revents & POLLIN)) {
      continue;
    }

    struct sockaddr_in client_addr;
    socklen_t client_len = sizeof(client_addr);
    int cfd = accept(lfd, reinterpret_cast<struct sockaddr *>(&client_addr),
                     &client_len);
    if (cfd < 0) {
      if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK) {
        continue;
      }
      if (stopping_.load()) {
        break;
      }
      continue;
    }

#ifdef SO_NOSIGPIPE
    int one = 1;
    setsockopt(cfd, SOL_SOCKET, SO_NOSIGPIPE, &one, sizeof(one));
#endif

    {
      std::unique_lock<std::mutex> lock(active_mu_);
      active_cv_.wait(lock, [&]() {
        return stopping_.load() ||
               active_connections_ < kMaxConcurrentConnections;
      });
      if (stopping_.load()) {
        close(cfd);
        break;
      }
      ++active_connections_;
    }

    std::thread([this, cfd]() {
      handle_client(cfd);
      close(cfd);
      {
        std::lock_guard<std::mutex> lock(active_mu_);
        --active_connections_;
      }
      active_cv_.notify_all();
    }).detach();
  }
}

void GatewayServer::stop() {
  bool expected = false;
  if (!stopping_.compare_exchange_strong(expected, true)) {
    return;
  }

  if (listen_fd_ >= 0) {
    shutdown(listen_fd_, SHUT_RDWR);
    close(listen_fd_);
    listen_fd_ = -1;
  }

  active_cv_.notify_all();

  std::unique_lock<std::mutex> lock(active_mu_);
  active_cv_.wait_for(lock, std::chrono::seconds(10), [&]() {
    return active_connections_ == 0;
  });
}

void GatewayServer::handle_client(int client_fd) {
  HttpRequest req;
  req.client_fd = client_fd;
  req.is_cancelled = [this, client_fd]() {
    return stopping_.load() || is_socket_disconnected(client_fd);
  };

  SocketResponseWriter w(client_fd);
  int err_status = 400;
  if (!read_http_request(client_fd, req, err_status)) {
    if (err_status == 413) {
      w.header().set("Content-Type", "text/plain; charset=utf-8");
      w.write_header(413);
      w.write("Payload Too Large\n");
      w.flush();
    }
    return;
  }

  if (req.path == "/v1/models" || req.path == "/models") {
    handle_models(cm_, auth_tokens_, req, w);
  } else if (req.path == "/v1/chat/completions" ||
             req.path == "/chat/completions") {
    handle_chat_completions(cm_, auth_tokens_, retry_delay_, req, w);
  } else if (req.path == "/v1/images/generations" ||
             req.path == "/images/generations") {
    handle_image_generations(cm_, auth_tokens_, retry_delay_, req, w);
  } else if (req.path == "/v1/messages" || req.path == "/messages") {
    handle_messages(cm_, auth_tokens_, retry_delay_, req, w);
  } else if (req.path == "/v1/responses" || req.path == "/responses") {
    handle_responses(cm_, auth_tokens_, retry_delay_, req, w);
  } else {
    w.header().set("Content-Type", "text/plain; charset=utf-8");
    w.write_header(404);
    w.write("404 page not found\n");
    w.flush();
  }
}