#include "login.hpp"    
#include "string.h"
#include "sodium.h"
#include "unordered_map"
#include "database.hpp"

namespace login
{
struct Session {
    std::string user_id;
    std::string csrf_token;
};

std::unordered_map<std::string, Session> sessions;

bool verifyHashWithPassword(const std::string& hash, std::string const& password)
{
    if (crypto_pwhash_str_verify(hash.c_str(), password.c_str(), password.size()) == 0) {
        return true;
    } else {
        return false;
    }
}

std::string hashPassword(const std::string& password)
{
    // Allocate storage for the hash
    char hash[crypto_pwhash_STRBYTES];

    // Hash the password using Argon2id
    if (crypto_pwhash_str(
            hash,
            password.c_str(),
            password.size(),
            crypto_pwhash_OPSLIMIT_INTERACTIVE,
            crypto_pwhash_MEMLIMIT_INTERACTIVE
        ) != 0) {
        std::cerr << "Out of memory while hashing password!" << std::endl;
        return "";
    }
    return hash;
}

std::string get_session_id(const crow::request& req) {
    auto cookie_header = req.get_header_value("Cookie");
    std::string prefix = "session_id=";
    auto pos = cookie_header.find(prefix);
    if (pos == std::string::npos) return "";
    return cookie_header.substr(pos + prefix.size(), 32);
}

// Utility: generate random string
std::string random_string(size_t length) {
    static const char charset[] =
        "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
    std::string result;
    result.resize(length);
    std::mt19937 rng(std::random_device{}());
    std::uniform_int_distribution<> dist(0, sizeof(charset)-2);
    for (size_t i = 0; i < length; ++i) result[i] = charset[dist(rng)];
    return result;
}

bool is_logged_in(const crow::request& req) {
    std::string session_id = get_session_id(req);
    if (session_id.empty()) return false;
    auto it = sessions.find(session_id);
    if (it == sessions.end()) return false;
    return !it->second.user_id.empty();
}

// lambda to be used by endpoint that requiere login 
auto login_required = [](auto handler){
    return [handler](const crow::request& req){
        if (!is_logged_in(req)) {
            crow::response res;
            if (req.get_header_value("HX-Request") == "true")
                res = crow::response(401, "Login required");
            else {
                res = crow::response(302);
                res.add_header("Location", "/login");
            }
            return res;
        }
        return handler(req);
    };
};

bool loginUser(const std::string& username, const std::string& password)
{
    auto sql = "SELECT id password_hash FROM users WHERE username = '?' LIMIT 1;";
    auto db = Database();
    
    if (!db.open())
        return false;

    auto opt_pair = db.get<int, std::string>(sql, {username});
    if (opt_pair.has_value()) {
        return verifyHashWithPassword(opt_pair.value().second, password); 
    } else {
        return false;
   }
}

bool initDB()
{
    auto db = Database();
    
    if (!db.open()){
        return false;
    }

    // Create a tables
    const char* create_sql_chars = "CREATE TABLE IF NOT EXISTS users (" 
        "id INTEGER PRIMARY KEY," 
        "username TEXT NOT NULL," 
        "password_hash TEXT NOT NULL,"
        "created_at DATETIME DEFAULT CURRENT_TIMESTAMP);";

    if (!db.exec(create_sql_chars)){
        CROW_LOG_ERROR << "Failed to create users table";
        return false;
    }
    
    return true;
}

bool initLogin(crow::SimpleApp& app)
{
    if (sodium_init() < 0) {
        CROW_LOG_ERROR << "Failed to Init Sodium";
        return false;
    }
    if(!initDB())
    {
        return false;
    }
    
    CROW_ROUTE(app, "/login").methods("GET"_method)
    ([](const crow::request& req){
        std::string csrf = random_string(32);
        // store CSRF in a temporary session cookie
        std::string session_id = random_string(32);
        sessions[session_id] = Session{"", csrf};

        crow::response res;
        res.body = "<form hx-post='/login' hx-swap='none'>"
                    "<input type='hidden' name='csrf_token' value='" + csrf + "'>"
                    "<input name='username'>"
                    "<input name='password' type='password'>"
                    "<button>Login</button></form>";
        res.add_header("Set-Cookie", "session_id=" + session_id + "; HttpOnly; Secure; SameSite=Strict");
        return res;
    });

    CROW_ROUTE(app, "/login").methods("POST"_method)
    ([](const crow::request& req){
        auto cookie_it = req.get_header_value("Cookie").find("session_id=");
        if (cookie_it == std::string::npos)
            return crow::response(401, "No session");

        // extract session_id
        std::string session_id = req.get_header_value("Cookie").substr(cookie_it + 11, 32);
        auto it = sessions.find(session_id);
        if (it == sessions.end()) return crow::response(401, "Invalid session");

        auto session = it->second;

        // parse form
        auto body = crow::query_string(req.body);
        std::string csrf_token = body.get("csrf_token");
        std::string username = body.get("username");
        std::string password = body.get("password");

        if (csrf_token != session.csrf_token) return crow::response(403, "CSRF failed");

        bool ok = loginUser(username, password); 

        if (!ok) return crow::response(401, "Invalid credentials");

        // regenerate session, mark as logged in
        sessions[session_id].user_id = "123"; // user ID

        crow::response res;
        res.add_header("HX-Redirect", "/dashboard"); // htmx redirect
        return res;
    });

    return true;
}
}