#include <botan/argon2fmt.h>
#include <botan/hex.h>
#include <botan/mac.h>
#include <botan/base32.h>
#include "server_internal.hxx"

std::string hash_password(const std::string &password) {
    return Botan::argon2_generate_pwhash(password.c_str(), password.size(), *auth_rng, 1, 64*1024, 4);
}

std::optional<std::string> Server::Auth_signup(std::string &&username, std::string &&password) {
    std::unique_lock lock{user_lock};
    for (const auto &user : users) {
        if (user.second->name == username)
            return "User already exists";
    }

    if (password.length() < 6)
        return "Password must be 6 characters or longer";

    auto hash = hash_password(password);
    auto id = next_user_id++;

    std::shared_ptr<Node> root{new Node{
        .id = 0,
        .name = "",
        .file = false,
        .preview = false,
        .parent = nullptr,
        .size = 0
    }};

    std::shared_ptr<User> user{new User{
        .id = id,
        .name = username,
        .password = hash,
        .enabled = false,
        .admin = false,
        .tfa_enabled = false
    }};
    user->nodes.emplace(0, std::move(root));
    user->user_dir = files_dir / std::to_string(user->id);

    std::filesystem::create_directory(user->user_dir);

    users.emplace(id, std::move(user));
    save();
    return std::nullopt;
}

mrpc::Response<mrpc::LoginResponse> Server::Auth_login(std::string &&username, std::string &&password, std::optional<std::string> &&otp) {
    std::shared_lock lock{user_lock};
    std::shared_ptr<User> user = nullptr;
    for (const auto &u : users) {
        if (u.second->name == username) {
            user = u.second;
            break;
        }
    }
    if (user == nullptr)
        return { .e = "Invalid username or password" };
    if (!Botan::argon2_check_pwhash(password.c_str(), password.size(), user->password))
        return { .e = "Invalid username or password" };
    if (!user->enabled)
        return { .e = "User is disabled" };
    if (user->tfa_enabled) {
        if (otp.has_value()) {
            bool ok = user->tfa_mail ? check_mail_code(user, otp.value()) : check_tfa_code(user, otp.value());
            if (!ok)
                return { .e = "Invalid code" };
        } else {
            if (user->tfa_mail)
                send_tfa_mail(user);
            mrpc::LoginResponse response { .otp_needed = true };
            return { .o = std::move(response) };
        }
    }

    auto token = Botan::hex_encode(auth_rng->random_vec(16), false);

    {
        std::unique_lock tlock{token_lock};
        tokens.emplace(token, std::make_unique<Token>(user));
    }

    mrpc::LoginResponse response {
        .otp_needed = false,
        .token = token
    };

    return { .o = std::move(response) };
}


void Server::Auth_logout(std::string &&token) {
    std::unique_lock lock{token_lock};
    tokens.erase(token);
}

std::optional<std::string> Server::Auth_logout_all(std::string &&token) {
    check_user_optional();
    logout_user(user->id);
    return std::nullopt;
}

void Server::Auth_send_recovery_key(std::string &&username) {
    std::shared_lock lock{user_lock};
    User *user = nullptr;
    for (const auto &u : users) {
        if (u.second->name == username) {
            user = u.second.get();
            break;
        }
    }
    if (user == nullptr)
        return;

    auto code = Botan::hex_encode(auth_rng->random_vec(5), true);
    {
        std::lock_guard rlock{recovery_lock};
        recovery_keys.emplace(code, std::make_pair(user->id, Token::clock::now() + std::chrono::minutes{5}));
    }

    send_mail(user->name, "MFileserver - Password recovery", "Your recovery key is: " + code + "\r\nIt is valid for 5 minutes");
}

std::optional<std::string> Server::Auth_reset_password(std::string &&key, std::string &&password) {
    std::lock_guard lock{recovery_lock};
    auto now = Token::clock::now();
    for (auto it = recovery_keys.begin(); it != recovery_keys.end();) {
        if (now >= it->second.second)
            recovery_keys.erase(it++);
        else
            ++it;
    }

    auto entry = recovery_keys.find(key);
    if (entry == recovery_keys.end())
        return "Invalid key";
    auto user = get_user(entry->second.first);
    if (!user)
        return "Invalid key";
    if (password.length() < 6)
        return "Password must be 6 characters or longer";
    recovery_keys.erase(key);

    logout_user(user->id);
    {
        std::shared_lock ulock{user_lock};
        user->password = hash_password(password);
    }
    save();

    return std::nullopt;
}

std::optional<std::string> Server::Auth_change_password(std::string &&token, std::string &&old_pw, std::string &&new_pw) {
    check_user_optional();
    std::shared_lock lock{user_lock};
    if (!Botan::argon2_check_pwhash(old_pw.c_str(), old_pw.size(), user->password))
        return "Old password is wrong";
    if (new_pw.length() < 6)
        return "Password must be 6 characters or longer";
    user->password = hash_password(new_pw);
    save();
    return std::nullopt;
}

std::optional<std::string> Server::Auth_tfa_setup_mail(std::string &&token) {
    check_user_optional();
    if (user->tfa_enabled)
        return "Tfa is already enabled";
    user->tfa_mail = true;
    send_tfa_mail(user);
    return std::nullopt;
}

mrpc::Response<std::string> Server::Auth_tfa_setup_totp(std::string &&token) {
    check_user_response();
    if (user->tfa_enabled)
        return { .e = "Tfa is already enabled" };
    user->tfa_mail = false;
    Botan::OctetString secret{*auth_rng, 16};
    user->tfa_secret = secret.to_string();

    return { .o = Botan::base32_encode(secret) };
}

std::optional<std::string> Server::Auth_tfa_complete(std::string &&token, std::string &&otp) {
    check_user_optional();
    bool ok = user->tfa_mail ? check_mail_code(user, otp) : check_tfa_code(user, otp);
    if (!ok)
        return "Invalid code";
    user->tfa_enabled = true;
    logout_user(user->id);
    save();
    return std::nullopt;
}

std::optional<std::string> Server::Auth_tfa_disable(std::string &&token) {
    check_user_optional();
    user->tfa_enabled = false;
    logout_user(user->id);
    save();
    return std::nullopt;
}

std::optional<std::string> Server::Auth_delete_user(std::string &&token) {
    check_user_optional();
    delete_user(user);
    save();
    return std::nullopt;
}

mrpc::Response<mrpc::Session> Server::Auth_session_info(std::string &&token) {
    check_user_response();
    auto t = get_token(token);
    if (!t)
        return { .e = "Invalid token" };

    mrpc::Session info {
        .name = user->name,
        .tfa_enabled = user->tfa_enabled,
        .admin = user->admin,
        .sudo = t->sudo_original_user != nullptr
    };

    return { .o = info };
}