// sinja: blazing-fast, stable, production-grade JSON templating REST server
// Architecture: Multi-threaded SO_REUSEPORT with robust, high-performance request handling.
// Dependency: inja (which includes nlohmann::json)
// Build: see CMakeLists.txt (ensure you link with -lpthread)
// Run  : ./sinja --templates /path/to/templates

#include <inja/inja.hpp>
#include <atomic>
#include <chrono>
#include <cctype>
#include <cerrno>
#include <csignal>
#include <cstring>
#include <filesystem>
#include <iostream>
#include <map>
#include <mutex>
#include <optional>
#include <stdexcept>
#include <string>
#include <thread>
#include <vector>

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

using json = nlohmann::json;
namespace fs = std::filesystem;

// ---- Tunables ----------------------------------------------------------------
static constexpr const char* kServerName = "sinja/5.0-stable";
static constexpr size_t kMaxBodyBytes   = 8 * 1024 * 1024;
static constexpr size_t kMaxHeaderBytes = 64 * 1024;
static constexpr int    kRecvTimeoutSec = 10;
static constexpr int    kSendTimeoutSec = 10;
static constexpr int    kBacklog        = 8192; // Increased for modern kernels

// ---- Global state for graceful shutdown --------------------------------------
static std::atomic<bool> g_running{true};
static void on_signal(int){ g_running.store(false); }

// ---- Config ------------------------------------------------------------------
struct Settings {
  std::string address = "0.0.0.0";
  int port = 8080;
  size_t threads = std::thread::hardware_concurrency() ? std::thread::hardware_concurrency() : 4;
  fs::path template_root;
};

// ---- Small utils -------------------------------------------------------------
static inline std::string ltrim(std::string s){
  size_t i=0; while(i<s.size() && std::isspace((unsigned char)s[i])) ++i; return s.substr(i);
}
static inline std::string rtrim(std::string s){
  if(s.empty()) return s;
  size_t i=s.size()-1;
  while (i<s.size() && std::isspace((unsigned char)s[i])) { if(i==0) return ""; --i; }
  return s.substr(0, i+1);
}
static inline std::string trim(std::string s){ return rtrim(ltrim(std::move(s))); }
static std::string to_lower(std::string s){
  for (auto& c: s) c = (char)std::tolower((unsigned char)c);
  return s;
}

// ---- Secure template validation ----------------------------------------------
static bool validate_template_path(const std::string& req) {
  return !(req.empty() || req[0] == '/' || req.find("..") != std::string::npos);
}

// ---- HTTP primitives ---------------------------------------------------------
struct Request {
  std::string method, target, version;
  std::map<std::string, std::string> headers;
  std::string body;
  bool keep_alive = true;
};
struct Response {
  int status = 200;
  std::string reason = "OK";
  std::vector<std::pair<std::string,std::string>> headers;
  std::string body;
};

static std::string status_reason(int code){
  switch(code){
    case 200: return "OK"; case 400: return "Bad Request"; case 404: return "Not Found";
    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"; default:  return "Error";
  }
}

static void set_socket_timeouts(int fd, int recv_sec, int send_sec){
  timeval tv{.tv_sec = recv_sec, .tv_usec = 0};
  setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
  tv = timeval{.tv_sec = send_sec, .tv_usec = 0};
  setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));
  int flag = 1; setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &flag, sizeof(flag));
}

static bool recv_append(int fd, std::string& buf){
  char tmp[8192];
  ssize_t n = ::recv(fd, tmp, sizeof(tmp), 0);
  if (n <= 0) return false;
  buf.append(tmp, (size_t)n);
  return true;
}

static bool send_all(int fd, const char* data, size_t len){
  size_t off = 0;
  while (off < len) {
    ssize_t n = ::send(fd, data + off, len - off, MSG_NOSIGNAL);
    if (n < 0) { if (errno == EINTR) continue; return false; }
    off += (size_t)n;
  }
  return true;
}

static std::string build_response(const Response& res, bool keep_alive){
  std::string out;
  out.reserve(256 + res.body.size());
  out += "HTTP/1.1 " + std::to_string(res.status) + " " + (res.reason.empty()?status_reason(res.status):res.reason) + "\r\n";
  out += "Server: "; out += kServerName; out += "\r\n";
  out += "Content-Length: " + std::to_string(res.body.size()) + "\r\n";
  out += std::string("Connection: ") + (keep_alive ? "keep-alive" : "close") + "\r\n";
  for (const auto& h : res.headers) { out += h.first; out += ": "; out += h.second; out += "\r\n"; }
  out += "\r\n";
  out += res.body;
  return out;
}

static bool is_json_ct(const std::map<std::string, std::string>& h){
  auto it = h.find("content-type");
  if (it == h.end()) return false;
  return to_lower(it->second).find("application/json") != std::string::npos;
}

static std::optional<std::pair<size_t, std::string>>
parse_chunked_body(const std::string& buf, size_t start_off) {
    size_t p = start_off;
    std::string out;
    while (true) {
        size_t eol = buf.find("\r\n", p);
        if (eol == std::string::npos) return std::nullopt;
        std::string size_line = buf.substr(p, eol - p);
        size_t sc = size_line.find(';'); if (sc != std::string::npos) size_line.resize(sc);
        size_line = trim(size_line);
        if (size_line.empty()) throw std::runtime_error("bad chunk size");
        size_t chunk_size = 0;
        try { chunk_size = std::stoul(size_line, nullptr, 16); } catch (...) { throw std::runtime_error("bad chunk size"); }
        p = eol + 2;
        if (chunk_size == 0) {
            size_t trailer_end = buf.find("\r\n\r\n", p);
            if (trailer_end == std::string::npos) return std::nullopt;
            return std::make_pair((trailer_end + 4) - start_off, std::move(out));
        }
        if (buf.size() < p + chunk_size + 2) return std::nullopt;
        if (out.size() + chunk_size > kMaxBodyBytes) throw std::runtime_error("payload too large");
        out.append(buf.data() + p, chunk_size);
        p += chunk_size;
        if (!(buf[p] == '\r' && buf[p+1] == '\n')) throw std::runtime_error("bad chunk CRLF");
        p += 2;
    }
}

static std::optional<Request> parse_request(std::string& inbuf) {
    size_t hdr_end = inbuf.find("\r\n\r\n");
    if (hdr_end == std::string::npos) return std::nullopt;
    if (hdr_end > kMaxHeaderBytes) throw std::runtime_error("headers too large");

    std::string head = inbuf.substr(0, hdr_end);
    size_t line_end = head.find("\r\n");
    if (line_end == std::string::npos) return std::nullopt;

    Request r;
    std::string req_line = head.substr(0, line_end);
    size_t sp1 = req_line.find(' ');
    size_t sp2 = req_line.find(' ', sp1 + 1);
    if (sp1 == std::string::npos || sp2 == std::string::npos) throw std::runtime_error("bad request line");
    r.method = req_line.substr(0, sp1);
    r.target = req_line.substr(sp1 + 1, sp2 - sp1 - 1);
    r.version = req_line.substr(sp2 + 1);

    size_t off = line_end + 2;
    while (off < head.size()) {
        size_t next_line = head.find("\r\n", off);
        std::string line = head.substr(off, next_line - off);
        if (next_line == std::string::npos) {
            off = head.size();
        } else {
            off = next_line + 2;
        }
        if (line.empty()) break;
        size_t colon_pos = line.find(':');
        if (colon_pos != std::string::npos) {
            r.headers[to_lower(trim(line.substr(0, colon_pos)))] = trim(line.substr(colon_pos + 1));
        }
    }

    const size_t body_start = hdr_end + 4;
    size_t consumed_len = 0;

    auto te_it = r.headers.find("transfer-encoding");
    if (te_it != r.headers.end() && to_lower(te_it->second).find("chunked") != std::string::npos) {
        auto parsed = parse_chunked_body(inbuf, body_start);
        if (!parsed) return std::nullopt;
        r.body = std::move(parsed->second);
        consumed_len = body_start + parsed->first;
    } else {
        size_t content_len = 0;
        auto cl_it = r.headers.find("content-length");
        if (cl_it != r.headers.end()) {
            try { content_len = std::stoull(cl_it->second); } catch (...) { throw std::runtime_error("invalid content-length"); }
            if (content_len > kMaxBodyBytes) throw std::runtime_error("payload too large");
        }
        if (inbuf.size() < body_start + content_len) return std::nullopt;
        r.body = inbuf.substr(body_start, content_len);
        consumed_len = body_start + content_len;
    }
    
    inbuf.erase(0, consumed_len);

    auto conn_it = r.headers.find("connection");
    if (r.version == "HTTP/1.0") {
        r.keep_alive = (conn_it != r.headers.end() && to_lower(conn_it->second) == "keep-alive");
    } else {
        r.keep_alive = (conn_it == r.headers.end() || to_lower(conn_it->second) != "close");
    }
    return r;
}

// ---- Core Server Logic -------------------------------------------------------
class Server {
  Settings cfg;
  std::vector<std::thread> workers;
  std::vector<int> listen_fds;
  std::mutex fds_mutex;

public:
  explicit Server(Settings s): cfg(std::move(s)) {}

  void start() {
    std::cout << "Starting " << kServerName << "...\n";
    for (size_t i = 0; i < cfg.threads; ++i) {
      workers.emplace_back([this]{ worker_loop(); });
    }
    std::cout << "  Listening on http://" << cfg.address << ":" << cfg.port << " with " << cfg.threads << " workers.\n";
  }

  void stop() {
    g_running.store(false);
    std::cout << "\nShutting down... finishing active connections.\n";
    std::lock_guard<std::mutex> lock(fds_mutex);
    for (int fd : listen_fds) {
        ::shutdown(fd, SHUT_RDWR);
        ::close(fd);
    }
  }

  void join() {
    for (auto& t : workers) {
      if (t.joinable()) t.join();
    }
    std::cout << "All workers stopped. Shutdown complete.\n";
  }

private:
  void send_json_error(int fd, int code, const std::string& msg) {
    Response res;
    res.status = code; res.reason = status_reason(code);
    res.headers.push_back({"Content-Type","application/json; charset=utf-8"});
    res.headers.push_back({"Cache-Control","no-store"});
    res.body = json({{"error", msg}}).dump();
    auto raw = build_response(res, false);
    send_all(fd, raw.data(), raw.size());
  }

  void handle_connection(int fd, inja::Environment& env) {
    std::string inbuf;
    inbuf.reserve(16*1024);
    bool keep = true;

    while (keep && g_running.load(std::memory_order_relaxed)) {
        std::optional<Request> reqOpt;
        while (g_running.load(std::memory_order_relaxed)) {
            try {
                reqOpt = parse_request(inbuf);
                if (reqOpt) break;
            } catch (const std::exception& e) {
                int code = 400; std::string what = e.what();
                if (what == "payload too large") code = 413;
                else if (what == "headers too large") code = 431;
                send_json_error(fd, code, what);
                return;
            }
            if (!recv_append(fd, inbuf)) return;
        }

        if (!reqOpt) return;
        Request req = std::move(*reqOpt);
        keep = req.keep_alive;

        if (!(req.method == "POST" && req.target == "/render")) {
            send_json_error(fd, 404, "Not Found");
            continue;
        }
        if (!is_json_ct(req.headers)) {
            send_json_error(fd, 415, "Content-Type must be application/json");
            continue;
        }

        json jreq;
        try {
            if (req.body.empty()) throw std::runtime_error("Request body is empty");
            jreq = json::parse(req.body);
        } catch (const std::exception& e) {
            send_json_error(fd, 400, std::string("Invalid JSON: ") + e.what());
            continue;
        }

        if (!jreq.contains("template") || !jreq["template"].is_string()) {
            send_json_error(fd, 400, R"(Missing "template" (string))"); continue;
        }
        if (jreq.contains("context") && !jreq["context"].is_object()) {
            send_json_error(fd, 400, R"("context" must be an object)"); continue;
        }

        std::string template_name = jreq["template"].get<std::string>();
        if (!validate_template_path(template_name)) {
            send_json_error(fd, 400, "Template error: invalid path"); continue;
        }

        json ctx = jreq.value("context", json::object());
        Response res;

        auto t0 = std::chrono::steady_clock::now();
        try {
            res.body = env.render_file(template_name, ctx);
            auto t1 = std::chrono::steady_clock::now();
            auto us = std::chrono::duration_cast<std::chrono::microseconds>(t1 - t0).count();
            res.headers.push_back({"Content-Type","text/plain; charset=utf-8"});
            res.headers.push_back({"X-Render-Time-Us", std::to_string(us)});
        } catch (const std::exception& e) {
            send_json_error(fd, 500, std::string("Render error: ") + e.what());
            continue;
        }

        auto raw = build_response(res, keep);
        if (!send_all(fd, raw.data(), raw.size())) return;
    }
  }

  void worker_loop() {
    inja::Environment env(cfg.template_root.string());
    env.set_trim_blocks(true);
    env.set_lstrip_blocks(true);

    int listen_fd = ::socket(AF_INET, SOCK_STREAM, 0);
    if (listen_fd < 0) { perror("socket"); return; }

    int yes=1;
    setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));
#ifdef SO_REUSEPORT
    setsockopt(listen_fd, SOL_SOCKET, SO_REUSEPORT, &yes, sizeof(yes));
#endif

    sockaddr_in addr{};
    addr.sin_family = AF_INET;
    addr.sin_port = htons(cfg.port);
    if (::inet_pton(AF_INET, cfg.address.c_str(), &addr.sin_addr) != 1) {
        std::cerr << "Invalid address\n"; ::close(listen_fd); return;
    }
    if (::bind(listen_fd, (sockaddr*)&addr, sizeof(addr)) < 0) {
        perror("bind"); ::close(listen_fd); return;
    }
    if (::listen(listen_fd, kBacklog) < 0) {
        perror("listen"); ::close(listen_fd); return;
    }

    {
      std::lock_guard<std::mutex> lock(fds_mutex);
      listen_fds.push_back(listen_fd);
    }

    while (g_running.load(std::memory_order_relaxed)) {
        int client_fd = ::accept(listen_fd, nullptr, nullptr);
        if (client_fd < 0) {
            if (g_running.load(std::memory_order_relaxed)) perror("accept");
            break;
        }
        set_socket_timeouts(client_fd, kRecvTimeoutSec, kSendTimeoutSec);
        handle_connection(client_fd, env);
        ::shutdown(client_fd, SHUT_WR);
        ::close(client_fd);
    }

    ::close(listen_fd);
  }
};

// ---- Args & main -------------------------------------------------------------
static Settings parse_args(int argc, char** argv){
  Settings s;
  for (int i=1;i<argc;i++){
    std::string a = argv[i];
    auto next = [&]()->std::string{ if (i+1>=argc) throw std::runtime_error("Missing value for "+a); return std::string(argv[++i]); };
    if (a=="--templates"||a=="-t") s.template_root = next();
    else if (a=="--address"||a=="-a") s.address = next();
    else if (a=="--port"||a=="-p") s.port = std::stoi(next());
    else if (a=="--threads"||a=="-w") s.threads = (size_t)std::stoul(next());
    else if (a=="--help"||a=="-h") {
      std::cout << "Usage: sinja --templates DIR [options]\n"
                   "  --address, -a          Bind address (default 0.0.0.0)\n"
                   "  --port,    -p          Port (default 8080)\n"
                   "  --threads, -w          Worker threads (default HW concurrency)\n";
      std::exit(0);
    } else throw std::runtime_error("Unknown arg: " + a);
  }
  if (s.template_root.empty()) throw std::runtime_error("Missing required --templates DIR");
  if (!fs::is_directory(s.template_root)) throw std::runtime_error("Template directory invalid: " + s.template_root.string());
  (s.template_root) = fs::weakly_canonical(s.template_root);
  return s;
}

int main(int argc, char** argv){
  try {
    auto cfg = parse_args(argc, argv);
    std::signal(SIGPIPE, SIG_IGN);
    std::signal(SIGINT,  on_signal);
    std::signal(SIGTERM, on_signal);

    Server srv(std::move(cfg));
    srv.start();
    
    while (g_running.load()) {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }

    srv.stop();
    srv.join();

  } catch (const std::exception& e) {
    std::cerr << "Fatal: " << e.what() << "\n";
    return 1;
  }
  return 0;
}