#include "pihttpserversessionauth.h"

#include "pidigest.h"
#include "pijson.h"
#include "piliterals_string.h"


namespace {

PIHTTP::MessageMutable jsonReply(PIHTTP::Code code, const PIJSON & j) {
	return PIHTTP::MessageMutable::fromCode(code)
	    .addHeader(PIHTTP::Header::ContentType, "application/json")
	    .setBody(j.toJSON(PIJSON::Compact).toByteArray());
}


PIHTTP::MessageMutable errorReply(PIHTTP::Code code, const PIString & error) {
	PIJSON j;
	j["error"] = error;
	return jsonReply(code, j);
}


bool parseToken(const PIString & auth_header, PIString & token) {
	if (auth_header.left(7).toLowerCase() != "bearer "_a) return false;
	PIString rest = auth_header.mid(7).trimmed();
	if (rest.isEmpty()) return false;
	token = rest;
	return true;
}

} // namespace


PIHTTPServerSessionAuth::PIHTTPServerSessionAuth()
	: token_length(PIDigest::hashLength(PIDigest::Type::BLAKE2b_512))
	, login_path("/api/login"_a)
	, logout_path("/api/logout"_a) {
	setBearerAuthCallback([this](const PIString & token) -> PIHTTP::AuthInfo { return checkToken(token); });
	registerLoginRoute();
	registerLogoutRoute();
}


void PIHTTPServerSessionAuth::registerLoginRoute() {
	registerPath(login_path, PIHTTP::Method::Post, [this](const PIHTTP::MessageConst & r) { return handleLogin(r); });
}


void PIHTTPServerSessionAuth::registerLogoutRoute() {
	registerProtectedPath(logout_path, PIHTTP::Method::Post, [this](const PIHTTP::MessageConst & r, const PIHTTP::AuthInfo &) {
		return handleLogout(r);
	});
}


void PIHTTPServerSessionAuth::setLoginPath(const PIString & p) {
	if (p.isEmpty() || p == login_path) return;
	unregisterPath(login_path, PIHTTP::Method::Post);
	login_path = p;
	registerLoginRoute();
}


void PIHTTPServerSessionAuth::setLogoutPath(const PIString & p) {
	if (p.isEmpty() || p == logout_path) return;
	unregisterPath(logout_path, PIHTTP::Method::Post);
	logout_path = p;
	registerLogoutRoute();
}


PIHTTP::AuthInfo PIHTTPServerSessionAuth::checkToken(const PIString & token) {
	PIMutexLocker locker(mutex);
	SessionRec s = sessions.value(token);
	if (s.isInvalid()) return false;
	const PISystemTime now = PISystemTime::current();
	if (hasTokenTtl() && now - s.created >= token_ttl) {
		sessions.remove(token);
		return false;
	}
	if (hasSessionIdleTimeout()) {
		const PISystemTime base = s.last_used.isNull() ? s.created : s.last_used;
		if (now - base >= session_idle_timeout) {
			sessions.remove(token);
			return false;
		}
		sessions[token].last_used = now;
	}
	return {true, s.user_id};
}


PIString PIHTTPServerSessionAuth::issueToken(int user_id) {
	PIMutexLocker locker(mutex);
	cleanupExpiredLocked();
	PIByteArray data = token_generator ? token_generator(token_length) : defaultTokenData(token_length);
	if (data.isEmpty()) return PIString();
	PIString token = data.toHex();
	// Refuse an already used token instead of silently merging two sessions.
	if (sessions.contains(token)) return PIString();
	SessionRec & s(sessions[token]);
	s.user_id   = user_id;
	s.created   = PISystemTime::current();
	s.last_used = s.created;
	return token;
}


void PIHTTPServerSessionAuth::cleanupExpired() {
	PIMutexLocker locker(mutex);
	cleanupExpiredLocked();
}


void PIHTTPServerSessionAuth::cleanupExpiredLocked() {
	const PISystemTime now = PISystemTime::current();
	sessions.removeWhere([this, now](const PIString &, const SessionRec & s) {
		if (s.isInvalid()) return true;
		if (!token_ttl.isNull() && now - s.created >= token_ttl) return true;
		if (!session_idle_timeout.isNull()) {
			const PISystemTime base = s.last_used.isNull() ? s.created : s.last_used;
			if (now - base >= session_idle_timeout) return true;
		}
		return false;
	});
}


bool PIHTTPServerSessionAuth::revokeToken(const PIString & token) {
	PIMutexLocker locker(mutex);
	if (!sessions.contains(token)) return false;
	sessions.remove(token);
	return true;
}


void PIHTTPServerSessionAuth::revokeUserTokens(int user_id) {
	PIMutexLocker locker(mutex);
	sessions.removeWhere([user_id](const PIString &, const SessionRec & s) { return s.user_id == user_id; });
}


void PIHTTPServerSessionAuth::setTokenGenerator(TokenGenerator g) {
	PIMutexLocker locker(mutex);
	token_generator = std::move(g);
}


void PIHTTPServerSessionAuth::setTokenTtl(PISystemTime ttl) {
	PIMutexLocker locker(mutex);
	token_ttl = ttl;
}


PISystemTime PIHTTPServerSessionAuth::tokenTtl() const {
	PIMutexLocker locker(mutex);
	return token_ttl;
}


bool PIHTTPServerSessionAuth::hasTokenTtl() const {
	PIMutexLocker locker(mutex);
	return !token_ttl.isNull();
}


void PIHTTPServerSessionAuth::setSessionIdleTimeout(PISystemTime t) {
	PIMutexLocker locker(mutex);
	session_idle_timeout = t;
}


PISystemTime PIHTTPServerSessionAuth::sessionIdleTimeout() const {
	PIMutexLocker locker(mutex);
	return session_idle_timeout;
}


bool PIHTTPServerSessionAuth::hasSessionIdleTimeout() const {
	PIMutexLocker locker(mutex);
	return !session_idle_timeout.isNull();
}


PIHTTP::MessageMutable PIHTTPServerSessionAuth::handleLogin(const PIHTTP::MessageConst & request) {
	PIJSON j = PIJSON::fromJSON(PIString::fromUTF8(request.body()));
	if (!j.isObject() || !j.contains("login") || !j.contains("password") || j["login"].toString().isEmpty())
		return errorReply(PIHTTP::Code::BadRequest, "invalid json");
	PIString login       = j["login"].toString();
	PIString password    = j["password"].toString();
	// The credential check runs without the internal lock held: the subclass owns the user
	// storage and may take its own locks.
	const auto auth_info = checkCredentials(login, password);
	if (!auth_info.authorized) return errorReply(PIHTTP::Code::Unauthorized, "invalid credentials");
	PIString token = issueToken(auth_info.user_id);
	if (token.isEmpty()) return errorReply(PIHTTP::Code::InternalServerError, "token generation failed");
	PIJSON resp;
	resp["token"]   = token;
	resp["user_id"] = auth_info.user_id;
	return jsonReply(PIHTTP::Code::Ok, resp);
}


PIHTTP::MessageMutable PIHTTPServerSessionAuth::handleLogout(const PIHTTP::MessageConst & request) {
	PIString token;
	if (parseToken(request.headers().value(PIHTTP::Header::Authorization, ""), token)) revokeToken(token);
	return PIHTTP::MessageMutable::fromCode(PIHTTP::Code::NoContent);
}


PIByteArray PIHTTPServerSessionAuth::defaultTokenData(uint len) {
	static std::atomic<ullong> counter(0);
	PIByteArray in = piSerialize(PISystemTime::current());
	in << len << counter++;
	PIByteArray out;
	while (out.size() < len) {
		out.append(PIDigest::calculate(in, PIDigest::Type::BLAKE2b_512));
		in = out;
		in << counter++;
	}
	return out.resized(len);
}