#include <SQLiteCpp/SQLiteCpp.h>
#include <pwd.h>
#include <security/pam_modules.h>
#include <sys/stat.h>
#include <unistd.h>

#include <cerrno>
#include <cstring>
#include <filesystem>
#include <format>
#include <print>

const std::filesystem::path DB_PATH = "/run/xdg_pam.db";
const std::filesystem::path RUNTIME_DIRECTORY = "/run/user/";

namespace pam_xdg {

void create_directory(const std::filesystem::path &path, mode_t chmod,
                      int owner) {
  if (mkdir(path.c_str(), chmod) != 0) {
    throw std::runtime_error(std::format("failed to create directory {}: {}",
                                         path.c_str(), std::strerror(errno)));
  }

  if (chown(path.c_str(), owner, owner) != 0) {
    throw std::runtime_error(
        std::format("failed to chown directory {}", path.c_str()));
  }
}

void create_table(SQLite::Database &db) {
  db.exec("CREATE TABLE IF NOT EXISTS SESSIONS(USERNAME STRING, LOGINS INT)");
}

void create_session(SQLite::Database &db, const std::string &username) {
  SQLite::Statement stmt(
      db, "INSERT OR IGNORE INTO SESSIONS(USERNAME, LOGINS) VALUES(?, 0)");

  stmt.bind(1, username);

  stmt.exec();
}

void login(SQLite::Database &db, const std::string &username) {
  SQLite::Statement stmt(
      db, "UPDATE SESSIONS SET LOGINS=LOGINS+1 WHERE USERNAME=?");

  stmt.bind(1, username);

  stmt.exec();
}

void logout(SQLite::Database &db, const std::string &username) {
  SQLite::Statement stmt(
      db, "UPDATE SESSIONS SET LOGINS=LOGINS-1 WHERE USERNAME=?");

  stmt.bind(1, username);

  stmt.exec();
}

bool is_logged_in(SQLite::Database &db, const std::string &username) {
  SQLite::Statement stmt(db, "SELECT LOGINS FROM SESSIONS WHERE USERNAME=?");

  stmt.bind(1, username);

  stmt.executeStep();

  int logins = stmt.getColumn(0);

  return (logins > 0);
}
} // namespace pam_xdg

extern "C" int pam_sm_open_session(pam_handle_t *pam, [[gnu::unused]] int flags,
                                   [[gnu::unused]] int argc,
                                   [[gnu::unused]] const char **argv) {

  const void *data;
  pam_get_item(pam, PAM_USER, &data);
  auto *username = static_cast<const char *>(data);

  if (std::string_view(username) == "root") {
    return 0;
  }

  auto passwd = getpwnam(username);
  auto uid = passwd->pw_uid;

  auto user_directory = RUNTIME_DIRECTORY / std::format("{}", uid);

  try {
    SQLite::Database db(DB_PATH, SQLite::OPEN_READWRITE | SQLite::OPEN_CREATE);

    pam_xdg::create_table(db);
    pam_xdg::create_session(db, username);

    if (!pam_xdg::is_logged_in(db, username)) {
      if (!std::filesystem::exists(RUNTIME_DIRECTORY)) {
        pam_xdg::create_directory(RUNTIME_DIRECTORY, 0755, 0);
      }

      pam_xdg::create_directory(user_directory, 0700, uid);
    }

    pam_xdg::login(db, username);

    auto envvar = std::format("XDG_RUNTIME_DIR={}", user_directory.c_str());

    if (pam_putenv(pam, envvar.c_str()) != PAM_SUCCESS) {
      throw std::runtime_error(
          std::string("failed to set XDG_RUNTIME_DIR environment variable"));
    }
  } catch (std::exception &e) {
    std::print("{}", e.what());
    return PAM_SESSION_ERR;
  }

  return PAM_SUCCESS;
}

extern "C" int pam_sm_close_session(pam_handle_t *pam,
                                    [[gnu::unused]] int flags,
                                    [[gnu::unused]] int argc,
                                    [[gnu::unused]] const char **argv) {

  const void *data;
  pam_get_item(pam, PAM_USER, &data);
  auto *username = static_cast<const char *>(data);

  if (std::string_view(username) == "root") {
    return 0;
  }

  auto passwd = getpwnam(username);
  auto uid = passwd->pw_uid;

  auto user_directory = std::filesystem::path(
      std::filesystem::path(RUNTIME_DIRECTORY) / std::format("{}", uid));

  try {
    SQLite::Database db(DB_PATH, SQLite::OPEN_READWRITE | SQLite::OPEN_CREATE);

    pam_xdg::logout(db, username);

    if (!pam_xdg::is_logged_in(db, username)) {
      std::filesystem::remove_all(user_directory);
    }

  } catch (std::exception &e) {
    std::print("{}", e.what());
    return PAM_SESSION_ERR;
  }

  return PAM_SUCCESS;
}