#include <cerrno>
#include <cstring>
#include <filesystem>
#include <print>
#include <stdexcept>
#include <string>
#include <string_view>

#include <fcntl.h>
#include <pwd.h>
#include <sys/stat.h>
#include <unistd.h>

#include <security/pam_modules.h>
#include <sqlite3.h>

namespace pam_xdg {

constexpr std::string_view USER_DIRECTORY = "/run/user";
constexpr std::string_view DB_PATH = "/run/pam_xdg.db";

void open_db(sqlite3 **db, const std::filesystem::path &path, int flags) {
  auto rc = sqlite3_open_v2(path.c_str(), db, flags, nullptr);

  if (rc != SQLITE_OK) {
    throw std::runtime_error(std::string("failed to open database"));
  }
}

void init_db(sqlite3 *db) {
  auto rc = sqlite3_exec(db, "CREATE TABLE SESSIONS(USERNAME TEXT, LOGINS INT)",
                         nullptr, nullptr, nullptr);

  if (rc != SQLITE_OK) {
    throw std::runtime_error(
        std::string("failed to init database: {}", sqlite3_errmsg(db)));
  }
}

void init_session(sqlite3 *db, std::string_view username) {
  const auto *query = "INSERT INTO SESSIONS(USERNAME, LOGINS) VALUES(?, 0)";
  sqlite3_stmt *stmt;
  int rc;

  rc = sqlite3_prepare_v2(db, query, -1, &stmt, NULL);

  if (rc != SQLITE_OK) {
    throw std::runtime_error(
        std::format("failed to prepare stmt: {}", sqlite3_errmsg(db)));
  }

  rc = sqlite3_bind_text(stmt, 1, username.data(), username.length(), nullptr);

  if (rc != SQLITE_OK) {
    sqlite3_finalize(stmt);
    throw std::runtime_error(std::string("failed to bind text to stmt"));
  }

  rc = sqlite3_step(stmt);

  if (rc != SQLITE_DONE) {
    sqlite3_finalize(stmt);
    throw std::runtime_error(std::string("failed to step stmt"));
  }

  sqlite3_finalize(stmt);
}

void record_login(sqlite3 *db, std::string_view username) {
  const auto *query = "UPDATE SESSIONS SET LOGINS=LOGINS+1 WHERE USERNAME=?";
  sqlite3_stmt *stmt;
  int rc;

  rc = sqlite3_prepare_v2(db, query, -1, &stmt, nullptr);

  if (rc != SQLITE_OK) {
    throw std::runtime_error(
        std::string("failed to prepare stmt: {}", sqlite3_errmsg(db)));
  }

  rc = sqlite3_bind_text(stmt, 1, username.data(), username.length(), nullptr);

  if (rc != SQLITE_OK) {
    sqlite3_finalize(stmt);
    throw std::runtime_error(std::string("failed to bind text"));
  }

  rc = sqlite3_step(stmt);

  if (rc != SQLITE_DONE) {
    sqlite3_finalize(stmt);
    throw std::runtime_error(std::string("failed to step stmt"));
  }

  sqlite3_finalize(stmt);
}

void record_logout(sqlite3 *db, std::string_view username) {
  const auto *query = "UPDATE SESSIONS SET LOGINS=LOGINS-1 WHERE USERNAME=?";
  sqlite3_stmt *stmt;
  int rc;

  rc = sqlite3_prepare_v2(db, query, -1, &stmt, nullptr);

  if (rc != SQLITE_OK) {
    throw std::runtime_error(
        std::string("failed to prepare stmt: {}", sqlite3_errmsg(db)));
  }

  rc = sqlite3_bind_text(stmt, 1, username.data(), username.length(), nullptr);

  if (rc != SQLITE_OK) {
    sqlite3_finalize(stmt);
    throw std::runtime_error(std::string("failed to bind text"));
  }

  rc = sqlite3_step(stmt);

  if (rc != SQLITE_DONE) {
    sqlite3_finalize(stmt);
    throw std::runtime_error(std::string("failed to step stmt"));
  }

  sqlite3_finalize(stmt);
}

bool is_logged_in(sqlite3 *db, std::string_view username) {
  const auto *query = "SELECT LOGINS FROM SESSIONS WHERE USERNAME=?";
  sqlite3_stmt *stmt;
  int rc;

  rc = sqlite3_prepare_v2(db, query, -1, &stmt, nullptr);

  if (rc != SQLITE_OK) {
    throw std::runtime_error(
        std::format("failed to prepare stmt: {}", sqlite3_errmsg(db)));
  }

  rc = sqlite3_bind_text(stmt, 1, username.data(), username.length(), nullptr);

  if (rc != SQLITE_OK) {
    sqlite3_finalize(stmt);
    throw std::runtime_error(std::string("failed to bind text"));
  }

  rc = sqlite3_step(stmt);

  if (rc != SQLITE_ROW) {
    sqlite3_finalize(stmt);
    throw std::runtime_error(std::string("failed to step stmt"));
  }

  int logins = sqlite3_column_int(stmt, 0);

  sqlite3_finalize(stmt);

  return (logins > 0);
}

bool has_session(sqlite3 *db, std::string_view username) {
  const auto *query = "SELECT COUNT(*) FROM SESSIONS WHERE USERNAME=?";
  sqlite3_stmt *stmt;
  int rc;

  rc = sqlite3_prepare_v2(db, query, -1, &stmt, nullptr);

  if (rc != SQLITE_OK) {
    throw std::runtime_error(
        std::format("failed to prepare stmt: {}", sqlite3_errmsg(db)));
  }

  rc = sqlite3_bind_text(stmt, 1, username.data(), username.length(), nullptr);

  if (rc != SQLITE_OK) {
    sqlite3_finalize(stmt);
    throw std::runtime_error(std::string("failed to bind text"));
  }

  rc = sqlite3_step(stmt);

  if (rc != SQLITE_ROW) {
    sqlite3_finalize(stmt);
    throw std::runtime_error(std::string("failed to step stmt"));
  }

  int count = sqlite3_column_int(stmt, 0);

  sqlite3_finalize(stmt);

  return (count > 0);
}

void touch(const std::filesystem::path &path, int mode) {
  auto fd = creat(path.c_str(), mode);

  if (fd == -1) {
    throw std::runtime_error(std::format("failed to touch file: {} {}",
                                         path.c_str(), std::strerror(errno)));
  }

  close(fd);
}

void create_directory(const std::filesystem::path &path, std::uint16_t owner,
                      int mode) {
  int rc;

  rc = mkdir(path.c_str(), mode);

  if (rc == -1) {
    throw std::runtime_error(std::format("failed to mkdir: {} {}", path.c_str(),
                                         std::strerror(errno)));
  }

  rc = chown(path.c_str(), owner, owner);

  if (rc == -1) {
    throw std::runtime_error(std::format("failed to chown: {} {}", path.c_str(),
                                         std::strerror(errno)));
  }
}
} // namespace pam_xdg

extern "C" int pam_sm_open_session(pam_handle_t *pamh,
                                   [[maybe_unused]] int flags,
                                   [[maybe_unused]] int argc,
                                   [[maybe_unused]] const char **argv) {

  const void *data;
  if (pam_get_item(pamh, PAM_USER, &data) != PAM_SUCCESS) {
    std::println(stderr, "failed to get pam item PAM_USER");
    return PAM_SESSION_ERR;
  }
  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 xdg_runtime_dir =
      std::filesystem::path(pam_xdg::USER_DIRECTORY) / std::format("{}", uid);

  sqlite3 *db;

  try {
    if (!std::filesystem::exists(pam_xdg::DB_PATH)) {
      pam_xdg::touch(pam_xdg::DB_PATH, 0600);

      pam_xdg::open_db(&db, pam_xdg::DB_PATH,
                       SQLITE_OPEN_CREATE | SQLITE_OPEN_READWRITE);

      pam_xdg::init_db(db);
    } else {
      pam_xdg::open_db(&db, pam_xdg::DB_PATH, SQLITE_OPEN_READWRITE);
    }

    if (!pam_xdg::has_session(db, username)) {
      pam_xdg::init_session(db, username);
    }

    if (!pam_xdg::is_logged_in(db, username)) {
      if (!std::filesystem::exists(pam_xdg::USER_DIRECTORY)) {
        pam_xdg::create_directory(pam_xdg::USER_DIRECTORY, 0, 0755);
      }
      if (!std::filesystem::is_directory(pam_xdg::USER_DIRECTORY)) {
        throw std::runtime_error(
            std::format("{} is not a directory", pam_xdg::USER_DIRECTORY));
      }

      pam_xdg::create_directory(std::filesystem::path(pam_xdg::USER_DIRECTORY) /
                                    std::format("{}", uid),
                                uid, 0700);
    }

    pam_xdg::record_login(db, username);

    auto envvar = std::format("XDG_RUNTIME_DIR={}", xdg_runtime_dir.c_str());
    if (pam_putenv(pamh, envvar.c_str()) != PAM_SUCCESS) {
      throw std::runtime_error(std::format("failed to putenv: {}", envvar));
    }

  } catch (const std::exception &e) {
    std::println(stderr, "{}", e.what());

    return PAM_SESSION_ERR;
  }

  return PAM_SUCCESS;
}

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

  const void *data;
  if (pam_get_item(pam, PAM_USER, &data) != PAM_SUCCESS) {
    std::println(stderr, "failed to get pam item PAM_USER");
    return PAM_SESSION_ERR;
  }
  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 xdg_runtime_dir =
      std::filesystem::path(pam_xdg::USER_DIRECTORY) / std::format("{}", uid);

  sqlite3 *db;

  try {
    pam_xdg::open_db(&db, pam_xdg::DB_PATH, SQLITE_OPEN_READWRITE);

    pam_xdg::record_logout(db, username);

    if (!pam_xdg::is_logged_in(db, username)) {
      std::filesystem::remove_all(xdg_runtime_dir);
    }
  } catch (const std::exception &e) {
    std::println(stderr, "{}", e.what());

    return PAM_SESSION_ERR;
  }

  return PAM_SUCCESS;
}