// A minimal HTTP/1.1 server for the local web UI. It only listens on the
// loopback interface, answers one request per connection ("Connection:
// close"), and runs each connection on its own thread so that an idle
// browser pre-connection cannot block anything.
//
// POSIX sockets, with a Winsock branch for Windows.
#pragma once

#include "platform.hpp"

#include <atomic>
#include <chrono>
#include <cstdint>
#include <cstring>
#include <functional>
#include <map>
#include <string>
#include <thread>
#include <utility>
#include <vector>

#include "crash.hpp"

#ifndef _WIN32
#include <arpa/inet.h>
#include <netinet/in.h>
#include <signal.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <unistd.h>
#endif

namespace lab {

struct Request {
  std::string method;
  std::string target;  // raw request target
  std::string path;    // decoded path, without the query
  std::map<std::string, std::string> query;
  std::map<std::string, std::string> headers;  // names lower-cased
  std::string body;

  std::string header(const std::string& name) const {
    const auto it = headers.find(name);
    return it == headers.end() ? std::string() : it->second;
  }
};

struct Response {
  int status = 200;
  std::string contentType = "text/plain; charset=utf-8";
  std::string body;
  std::vector<std::pair<std::string, std::string>> headers;
};

using Handler = std::function<Response(const Request&)>;

namespace http {

constexpr size_t kMaxHeaderBytes = 64 * 1024;
constexpr size_t kMaxBodyBytes = 24 * 1024 * 1024;

#ifdef _WIN32
using Sock = SOCKET;
inline bool validSock(Sock s) { return s != INVALID_SOCKET; }
inline void closeSock(Sock s) { closesocket(s); }
constexpr Sock kBadSock = INVALID_SOCKET;
#else
using Sock = int;
inline bool validSock(Sock s) { return s >= 0; }
inline void closeSock(Sock s) { ::close(s); }
constexpr Sock kBadSock = -1;
#endif

inline int hexVal(char c) {
  if (c >= '0' && c <= '9') return c - '0';
  if (c >= 'a' && c <= 'f') return c - 'a' + 10;
  if (c >= 'A' && c <= 'F') return c - 'A' + 10;
  return -1;
}

inline std::string percentDecode(const std::string& s, bool plusIsSpace) {
  std::string o;
  o.reserve(s.size());
  for (size_t i = 0; i < s.size(); ++i) {
    if (s[i] == '%' && i + 2 < s.size() && hexVal(s[i + 1]) >= 0 && hexVal(s[i + 2]) >= 0) {
      o += static_cast<char>(hexVal(s[i + 1]) * 16 + hexVal(s[i + 2]));
      i += 2;
    } else if (s[i] == '+' && plusIsSpace) {
      o += ' ';
    } else {
      o += s[i];
    }
  }
  return o;
}

inline std::string lower(std::string s) {
  for (char& c : s) {
    if (c >= 'A' && c <= 'Z') c = static_cast<char>(c - 'A' + 'a');
  }
  return s;
}

inline std::string trim(const std::string& s) {
  size_t a = 0, b = s.size();
  while (a < b && (s[a] == ' ' || s[a] == '\t')) ++a;
  while (b > a && (s[b - 1] == ' ' || s[b - 1] == '\t' || s[b - 1] == '\r')) --b;
  return s.substr(a, b - a);
}

inline const char* statusText(int code) {
  switch (code) {
    case 100: return "Continue";
    case 200: return "OK";
    case 204: return "No Content";
    case 302: return "Found";
    case 400: return "Bad Request";
    case 403: return "Forbidden";
    case 404: return "Not Found";
    case 405: return "Method Not Allowed";
    case 408: return "Request Timeout";
    case 409: return "Conflict";
    case 411: return "Length Required";
    case 413: return "Payload Too Large";
    case 415: return "Unsupported Media Type";
    case 431: return "Request Header Fields Too Large";
    case 500: return "Internal Server Error";
    case 501: return "Not Implemented";
    default: return "OK";
  }
}

inline bool sendAll(Sock s, const char* data, size_t n) {
  size_t sent = 0;
  while (sent < n) {
#ifdef _WIN32
    const int r = ::send(s, data + sent, static_cast<int>(n - sent), 0);
#elif defined(MSG_NOSIGNAL)
    const ssize_t r = ::send(s, data + sent, n - sent, MSG_NOSIGNAL);
#else
    const ssize_t r = ::send(s, data + sent, n - sent, 0);
#endif
    if (r <= 0) return false;
    sent += static_cast<size_t>(r);
  }
  return true;
}

inline void writeResponse(Sock s, const Response& r) {
  std::string head = "HTTP/1.1 " + std::to_string(r.status) + " " + statusText(r.status) + "\r\n";
  head += "Content-Type: " + r.contentType + "\r\n";
  head += "Content-Length: " + std::to_string(r.body.size()) + "\r\n";
  head += "Connection: close\r\n";
  head += "Cache-Control: no-store\r\n";
  head += "X-Content-Type-Options: nosniff\r\n";
  head += "Referrer-Policy: no-referrer\r\n";
  head += "X-Frame-Options: DENY\r\n";
  head += "Cross-Origin-Resource-Policy: same-origin\r\n";
  for (const auto& h : r.headers) head += h.first + ": " + h.second + "\r\n";
  head += "\r\n";
  if (!sendAll(s, head.data(), head.size())) return;
  if (!r.body.empty()) sendAll(s, r.body.data(), r.body.size());
}

inline Response errorResponse(int code, const std::string& msg) {
  Response r;
  r.status = code;
  r.contentType = "application/json; charset=utf-8";
  r.body = "{\"error\":\"" + msg + "\"}";
  return r;
}

// Reads and parses one request. Returns 0 on success, otherwise the HTTP
// status to answer with (-1: just drop the connection).
inline int readRequest(Sock s, Request& rq) {
  std::string buf;
  char chunk[8192];
  size_t headerEnd = std::string::npos;
  while (headerEnd == std::string::npos) {
    const int n = static_cast<int>(::recv(s, chunk, static_cast<int>(sizeof chunk), 0));
    if (n <= 0) return buf.empty() ? -1 : 408;
    buf.append(chunk, static_cast<size_t>(n));
    headerEnd = buf.find("\r\n\r\n");
    if (headerEnd == std::string::npos && buf.size() > kMaxHeaderBytes) return 431;
  }

  const std::string head = buf.substr(0, headerEnd);
  std::string body = buf.substr(headerEnd + 4);

  size_t eol = head.find("\r\n");
  const std::string line = head.substr(0, eol);
  const size_t sp1 = line.find(' ');
  const size_t sp2 = line.rfind(' ');
  if (sp1 == std::string::npos || sp2 == sp1) return 400;
  rq.method = line.substr(0, sp1);
  rq.target = line.substr(sp1 + 1, sp2 - sp1 - 1);
  if (line.compare(sp2 + 1, 5, "HTTP/") != 0) return 400;

  while (eol != std::string::npos) {
    const size_t start = eol + 2;
    eol = head.find("\r\n", start);
    const std::string h = head.substr(start, eol == std::string::npos ? std::string::npos : eol - start);
    const size_t colon = h.find(':');
    if (colon == std::string::npos || colon == 0) continue;
    rq.headers[lower(h.substr(0, colon))] = trim(h.substr(colon + 1));
  }

  const size_t qm = rq.target.find('?');
  rq.path = percentDecode(rq.target.substr(0, qm), false);
  if (qm != std::string::npos) {
    const std::string qs = rq.target.substr(qm + 1);
    size_t i = 0;
    while (i <= qs.size()) {
      const size_t amp = qs.find('&', i);
      const std::string pair = qs.substr(i, amp == std::string::npos ? std::string::npos : amp - i);
      if (!pair.empty()) {
        const size_t eq = pair.find('=');
        const std::string k = percentDecode(pair.substr(0, eq), true);
        const std::string v = eq == std::string::npos ? "" : percentDecode(pair.substr(eq + 1), true);
        rq.query[k] = v;
      }
      if (amp == std::string::npos) break;
      i = amp + 1;
    }
  }

  if (!rq.header("transfer-encoding").empty()) return 501;

  size_t want = 0;
  const std::string cl = rq.header("content-length");
  if (!cl.empty()) {
    for (char c : cl) {
      if (c < '0' || c > '9' || cl.size() > 12) return 400;
    }
    want = static_cast<size_t>(std::stoull(cl));
    if (want > kMaxBodyBytes) return 413;
  } else if (rq.method == "POST" || rq.method == "PUT") {
    return 411;
  }

  if (want > body.size() && lower(rq.header("expect")) == "100-continue") {
    const char* c100 = "HTTP/1.1 100 Continue\r\n\r\n";
    sendAll(s, c100, std::strlen(c100));
  }
  while (body.size() < want) {
    const int n = static_cast<int>(::recv(s, chunk, static_cast<int>(sizeof chunk), 0));
    if (n <= 0) return 408;
    body.append(chunk, static_cast<size_t>(n));
  }
  body.resize(want);
  rq.body = std::move(body);
  return 0;
}

#ifdef _WIN32
struct WinsockInit {
  WinsockInit() {
    WSADATA d;
    WSAStartup(MAKEWORD(2, 2), &d);
  }
  ~WinsockInit() { WSACleanup(); }
};
#endif

}  // namespace http

class HttpServer {
 public:
  HttpServer() {
#ifdef _WIN32
    static http::WinsockInit init;
#else
    ::signal(SIGPIPE, SIG_IGN);
#endif
  }
  ~HttpServer() {
    if (http::validSock(listen_)) http::closeSock(listen_);
  }
  HttpServer(const HttpServer&) = delete;
  HttpServer& operator=(const HttpServer&) = delete;

  // Tries each port in turn (0 = let the system choose). Loopback only.
  bool listenOn(const std::vector<int>& ports, std::string& error) {
    for (int p : ports) {
      http::Sock s = ::socket(AF_INET, SOCK_STREAM, 0);
      if (!http::validSock(s)) {
        error = "cannot create a socket";
        return false;
      }
      int yes = 1;
#ifdef _WIN32
      // Without this, another program could bind the same port and share our traffic.
      ::setsockopt(s, SOL_SOCKET, SO_EXCLUSIVEADDRUSE, reinterpret_cast<const char*>(&yes), sizeof yes);
      ::SetHandleInformation(reinterpret_cast<HANDLE>(s), HANDLE_FLAG_INHERIT, 0);  // not for the Java child
#else
      ::setsockopt(s, SOL_SOCKET, SO_REUSEADDR, reinterpret_cast<const char*>(&yes), sizeof yes);
#endif
      sockaddr_in addr;
      std::memset(&addr, 0, sizeof addr);
      addr.sin_family = AF_INET;
      addr.sin_port = htons(static_cast<uint16_t>(p));
      addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
      if (::bind(s, reinterpret_cast<sockaddr*>(&addr), sizeof addr) != 0 || ::listen(s, 32) != 0) {
        http::closeSock(s);
        continue;
      }
      sockaddr_in bound;
#ifdef _WIN32
      int len = sizeof bound;
#else
      socklen_t len = sizeof bound;
#endif
      ::getsockname(s, reinterpret_cast<sockaddr*>(&bound), &len);
      port_ = ntohs(bound.sin_port);
      listen_ = s;
      return true;
    }
    error = ports.size() == 1 ? "port " + std::to_string(ports[0]) + " is not available"
                              : "no free port found";
    return false;
  }

  int port() const { return port_; }

  // Blocks until `stop` becomes true.
  void serve(const Handler& handler, const std::atomic<bool>& stop) {
    while (!stop.load()) {
      fd_set fds;
      FD_ZERO(&fds);
      FD_SET(listen_, &fds);
      timeval tv;
      tv.tv_sec = 0;
      tv.tv_usec = 200 * 1000;
#ifdef _WIN32
      const int ready = ::select(0, &fds, nullptr, nullptr, &tv);
#else
      const int ready = ::select(listen_ + 1, &fds, nullptr, nullptr, &tv);
#endif
      if (ready <= 0) continue;
      http::Sock c = ::accept(listen_, nullptr, nullptr);
      if (!http::validSock(c)) continue;
#ifdef _WIN32
      ::SetHandleInformation(reinterpret_cast<HANDLE>(c), HANDLE_FLAG_INHERIT, 0);
#endif
      active_.fetch_add(1);
      std::thread([this, c, &handler]() {
        connection(c, handler);
        active_.fetch_sub(1);
      }).detach();
    }
    for (int i = 0; i < 100 && active_.load() > 0; ++i) {  // let running requests finish
      std::this_thread::sleep_for(std::chrono::milliseconds(50));
    }
  }

 private:
  http::Sock listen_ = http::kBadSock;
  int port_ = 0;
  std::atomic<int> active_{0};

  static void connection(http::Sock c, const Handler& handler) {
    crash::AltStack altStack;
#ifdef _WIN32
    DWORD ms = 15000;
    ::setsockopt(c, SOL_SOCKET, SO_RCVTIMEO, reinterpret_cast<const char*>(&ms), sizeof ms);
    ::setsockopt(c, SOL_SOCKET, SO_SNDTIMEO, reinterpret_cast<const char*>(&ms), sizeof ms);
#else
    timeval tv;
    tv.tv_sec = 15;
    tv.tv_usec = 0;
    ::setsockopt(c, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof tv);
    ::setsockopt(c, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof tv);
#endif
    Request rq;
    const int rc = http::readRequest(c, rq);
    if (rc > 0) {
      http::writeResponse(c, http::errorResponse(rc, http::statusText(rc)));
    } else if (rc == 0) {
      Response r;
      try {
        r = handler(rq);
      } catch (const std::exception& e) {
        r = http::errorResponse(500, "internal error");
        (void)e;
      }
      http::writeResponse(c, r);
    }
    http::closeSock(c);
  }
};

}  // namespace lab