#include <random>
#include <algorithm>
#include <UserDB/UserDB.hpp>
#include <Utilities/Utilities.hpp>
#include <Log/Log.hpp>
#include <Settings/Settings.hpp>

UserDB userDB("server/users.db");

UserDB::UserCache::iterator UserDB::tryGetUser(const std::string &username)
{
    auto it = users.find(username);
    if (it != users.end())
    {
        return it;
    }

    std::vector<uint8_t> bytes;
    if (!store->read(username, bytes))
    {
        return users.end();
    }

    memcpy(reinterpret_cast<void*>(&users[username]), bytes.data(), std::min(sizeof(User), bytes.size()));
    return users.find(username);
}

void UserDB::SessionID::generate()
{
    std::independent_bits_engine<std::default_random_engine, 8, unsigned short> rbe;
    std::generate(sid, sid + sizeof(sid), rbe);
}

bool UserDB::SessionID::operator==(UserDB::SessionID other)
{
    return std::memcmp(sid, other.sid, sizeof(sid)) == 0;
}

void UserDB::flushUser(std::string username)
{
    User &user = users[username];

    AtomicDataStore::Plan plan;
    auto &data = plan[username];
    data.resize(sizeof(User));
    memcpy(data.data(), &user, sizeof(user));
    store->update(plan, 0);
}

bool UserDB::isUsernameValid(std::string username)
{
    std::string::iterator invalidCharIter = std::find_if_not(username.begin(), username.end(), [](unsigned char c) -> bool
    {
        return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
            || c == '-' || c == '_';
    });

    return invalidCharIter == username.end();
}

UserDB::UserDB(std::string dbpath) : dbpath(dbpath) {}

void UserDB::open()
{
    store = std::make_unique<AtomicDataStore>();
    store->open(dbpath);
    const std::string &error = store->getError();
    if (!error.empty())
    {
        crash(formatString("Failed to open the user database: $+", error));
    }
}

void UserDB::create()
{
    store = std::make_unique<AtomicDataStore>();
    store->create(dbpath);
    const std::string &error = store->getError();
    if (!error.empty())
    {
        crash(formatString("Failed to create the user database: $+", error));
    }
}

UserDB::CreateUserResult UserDB::createUser(std::string username, std::string password, bool createAsAdmin)
{
    if (!isUsernameValid(username))
    {
        return CreateUserResult{false, "This username contains illegal characters! Only digits, letters, and -_ allowed."};
    }

    std::unique_lock<std::shared_mutex> lock(mutex);

    if (tryGetUser(username) != users.end())
    {
        return CreateUserResult{false, "This username is already taken."};
    }

    std::string message;
    if (createAsAdmin)
    {
        users[username].perms.setAll();

        message = "Created the admin account.";
    }
    else if (!settings.allowRegisterRequests)
    {
        return CreateUserResult{false, "This server is not accepting new accounts at the moment."};
    }
    else if (!settings.autoApproveRegisterRequests)
    {
        users[username].state = UserState::Pending;

        message = "An admin will now have to approve it before you can log in.";
    }
    else
    {
        message = "You can now log in.";
    }

    users[username].passwordHash = password;

    flushUser(username);

    return CreateUserResult{true, message};
}

bool UserDB::getUserInfo(std::string username, UserDB::User &outUser)
{
    if (!isUsernameValid(username))
    {
        return false;
    }

    std::unique_lock<std::shared_mutex> lock(mutex);

    auto it = tryGetUser(username);
    if (it == users.end())
    {
        return false;
    }

    outUser = it->second;
    return true;
}

bool UserDB::startSession(std::string username, std::string password, UserDB::User &outUser)
{
    std::unique_lock<std::shared_mutex> lock(mutex);

    auto it = tryGetUser(username);
    if (it == users.end())
    {
        return false;
    }

    User &user = it->second;
    if (user.passwordHash != password)
    {
        return false;
    }

    // Create a session if not banned and not pending.
    if (user.state != UserState::Banned && user.state != UserState::Pending)
    {
        user.state = UserState::KnownSessionID;
        user.expiryTime = (uint64_t) time(NULL) + 7 * 24 * 60 * 60;
        user.sid.generate();

        flushUser(username);
    }

    outUser = user;
    return true;
}

bool UserDB::approveUser(std::string username)
{
    std::unique_lock<std::shared_mutex> lock(mutex);

    auto it = tryGetUser(username);
    if (it == users.end())
    {
        return false;
    }

    User &user = it->second;

    // Approve the user if pending.
    if (user.state == UserState::Pending)
    {
        user.state = UserState::NoSessionID;
        flushUser(username);
        return true;
    }

    return false;
}

bool UserDB::grantPermission(std::string username, UserPermission perm)
{
    std::unique_lock<std::shared_mutex> lock(mutex);

    auto it = tryGetUser(username);
    if (it == users.end())
    {
        return false;
    }

    User &user = it->second;

    user.perms += perm;
    flushUser(username);

    return true;
}

bool UserDB::revokePermission(std::string username, UserPermission perm)
{
    std::unique_lock<std::shared_mutex> lock(mutex);

    auto it = tryGetUser(username);
    if (it == users.end())
    {
        return false;
    }

    User &user = it->second;

    user.perms -= perm;
    flushUser(username);

    return true;
}

std::string UserDB::banUser(std::string username)
{
    std::unique_lock<std::shared_mutex> lock(mutex);
    auto it = tryGetUser(username);
    if (it == users.end())
    {
        return "User does not exist.";
    }

    User &user = it->second;

    user.state = UserState::Banned;
    flushUser(username);

    return "";
}

std::string UserDB::unbanUser(std::string username)
{
    std::unique_lock<std::shared_mutex> lock(mutex);

    auto it = tryGetUser(username);
    if (it == users.end())
    {
        return "User does not exist.";
    }

    User &user = it->second;

    if (user.state != UserState::Banned)
    {
        return "User is not currently banned.";
    }

    user.state = UserState::NoSessionID;
    flushUser(username);

    return "";
}