#include "SessionManager.h"

#include <fmt/core.h>
#include "SafeRNG.h"
#ifndef _WIN32
#include "UOHash.h"
#else
#include "../UOHash/UOHash.h"
#endif
#include <fstream>
#include <filesystem>
#include <json/json.h>

const std::chrono::hours SessionManager::Session::max_age = std::chrono::hours(24 * 5);

SessionManager::Session::Session()	//默认构造函数：无效的已过期空Cookie
{
	uid = -1;
	expiry_date = std::chrono::system_clock::now() - std::chrono::seconds(1);
}

SessionManager::Session::Session(const Session& obj)
{
	uid = obj.uid;
	cookie = obj.cookie;
	expiry_date = obj.expiry_date;
}

SessionManager::Session::Session(const std::string& cookie)
{
	uid = -1;
	this->cookie = cookie;
	expiry_date = std::chrono::system_clock::now() - std::chrono::seconds(1);
}

SessionManager::Session::Session(const Json::Value& json_obj)
{
	if(json_obj["UID"].isInt() && json_obj["Cookie"].isString() && (json_obj["ExpiryDate"].isInt64()))
	{
		uid = json_obj["UID"].asInt();
		cookie = json_obj["Cookie"].asString();
		auto ms = std::chrono::milliseconds(json_obj["ExpiryDate"].asInt64());
		expiry_date = systime(ms);
	}
	else
	{
		uid = -1;
		expiry_date = std::chrono::system_clock::now() - std::chrono::seconds(1);
	}
}

SessionManager::Session::Session(int uid, const std::string& cookie)
{
	this->uid = uid;
	this->cookie = cookie;
	expiry_date = std::chrono::system_clock::now() + max_age;
}

int SessionManager::Session::GetUID() const
{
	return uid;
}

bool SessionManager::Session::Expired() const
{
	return (expiry_date < std::chrono::system_clock::now());
}

std::string SessionManager::Session::GetCookie() const
{
	return cookie;
}

SessionManager::systime SessionManager::Session::GetExpiryDate() const
{
	return expiry_date;
}

bool SessionManager::Session::NeedExpiryDateRefresh() const
{
	if (Expired())
		return true;
	auto usage_time = std::chrono::system_clock::now() - (expiry_date - max_age);
	return (usage_time >= std::chrono::hours(24));
}

void SessionManager::Session::SetUID(int uid)
{
	this->uid = uid;
}

void SessionManager::Session::RefreshExpiryDate()
{
	expiry_date = std::chrono::system_clock::now() + max_age;
}

void SessionManager::Session::SetCookie(const std::string& cookie)
{
	this->cookie = cookie;
}

void SessionManager::Session::SetExpiryDate(const systime& expiry_time)
{
	this->expiry_date = expiry_date;
}

bool SessionManager::Session::operator<(const Session& obj) const
{
	return (this->cookie < obj.cookie);
}

SessionManager::Session::operator Json::Value() const
{
	Json::Value root;
	root["UID"] = uid;
	root["Cookie"] = cookie;
	auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(expiry_date.time_since_epoch()).count();
	root["ExpiryDate"] = Json::Int64(ms);
	return root;
}

int SessionManager::Session::GetMaxAgeSeconds()
{
	return (max_age.count() * 3600);
}

SessionManager::SessionManager()
{
}

std::string SessionManager::Random()
{
	return RandomNumberGenerator::SecureRandomHex(32);
}

std::string SessionManager::ISO8601_TimeString()
{
	time_t now = time(nullptr);
	tm tim = {};
#if defined(_WIN32)
	localtime_s(&tim, &now);   // Windows
#else
	localtime_r(&now, &tim);   // POSIX
#endif
	return fmt::format("{}-{}-{}T{}:{}:{}+08:00", (tim.tm_year + 1900), (tim.tm_mon + 1), tim.tm_mday, tim.tm_hour, tim.tm_min, tim.tm_sec);
}

void SessionManager::AutoCleanCookiePool()
{
	std::lock_guard<std::mutex> lock(pool_mutex);	//自动清理前需要上锁
	for (auto it = cookie_storage.begin(); it != cookie_storage.end(); )	//正向遍历，删除无效/过期Session
	{
		const std::string& cookie = it->first;
		int uid = it->second;

		auto rit = reverse_cookie_storage.find(uid);
		if (rit == reverse_cookie_storage.end())	//有Cookie无Session，删掉
		{
			it = cookie_storage.erase(it);
			continue;
		}

		Session index_session(cookie.substr(prefix.length()));
		auto& rcs = rit->second;
		auto sit = rcs.find(index_session);
		if ((sit == rcs.end()) || sit->Expired())	//无Session或Session已过期，删掉
		{
			if (sit != rcs.end())
				rcs.erase(sit);
			if (rcs.empty())
				reverse_cookie_storage.erase(rit);
			it = cookie_storage.erase(it);
			continue;
		}
		++it;	//前置++性能开销较低
	}

	for (auto rit = reverse_cookie_storage.begin(); rit != reverse_cookie_storage.end(); )	//反向遍历，删除异常Session
	{
		auto& rcs = rit->second;
		for (auto sit = rcs.begin(); sit != rcs.end(); )
		{
			std::string key = prefix + sit->GetCookie();
			if (cookie_storage.find(key) == cookie_storage.end())	//正向没有反向有，异常Session（无法被利用）
			{
				auto sit_next = std::next(sit);
				rcs.erase(sit);
				sit = sit_next;
			}
			else
				++sit;
		}
		if (rcs.empty())
		{
			auto rit_next = std::next(rit);
			reverse_cookie_storage.erase(rit);
			rit = rit_next;
		}
		else
			++rit;
	}
}

std::string SessionManager::GenerateCookieForUser(int uid, bool cookie_only)
{
	//生成Cookie和Session
	std::string str_uid = std::to_string(uid);
	std::string raw_cookie = str_uid + ISO8601_TimeString() + str_uid + Random() + str_uid;
	auto res = uns::UOHash::HashString(uns::HashID::SHA3_224, raw_cookie);
	if (!res)
		return "";
	std::string cookie = res.GetResult();
	Session session(uid, cookie);
	//存储到Cookie池 - 使用互斥体保证线程安全
	{
		std::lock_guard<std::mutex> lock(pool_mutex);
		cookie_storage.insert({ (prefix + cookie), uid });
		if (reverse_cookie_storage.find(uid) == reverse_cookie_storage.end())
			reverse_cookie_storage.insert({ uid, { session } });
		else
			reverse_cookie_storage.at(uid).insert(session);
	}
	//返回用于响应的Cookie串
	if(cookie_only)
		return (prefix + cookie);
	else
		return fmt::format(fmt::runtime(cookie_template), cookie, Session::GetMaxAgeSeconds());
}

std::string SessionManager::DeleteCookie(const std::string& cookie)
{
	std::lock_guard<std::mutex> lock(pool_mutex);
	if (cookie_storage.find(cookie) != cookie_storage.end())
	{
		int uid = cookie_storage.at(cookie);
		cookie_storage.erase(cookie);
		if (reverse_cookie_storage.find(uid) != reverse_cookie_storage.end())
		{
			auto& rcs = reverse_cookie_storage.at(uid);
			Session index_session(cookie.substr(prefix.length()));
			rcs.erase(index_session);
			if (rcs.empty())
				reverse_cookie_storage.erase(uid);
		}
	}
	//返回用于响应的Cookie串（注销浏览器端的Cookie）
	return fmt::format(fmt::runtime(cookie_template.substr(prefix.length())), cookie, 0);
}

bool SessionManager::CheckRequestCookie(uns::RequestPtr request, int& uid)
{
	uid = -1;
	if (!request->HasHeader("Cookie"))	//未找到Cookie
		return false;
	std::string cookie = request->GetHeader("Cookie");
	std::lock_guard<std::mutex> lock(pool_mutex);	//使用互斥体保证Cookie池的线程安全
	auto cit = cookie_storage.find(cookie);
	if (cit == cookie_storage.end())	//不正确的Cookie
		return false;
	uid = cit->second;
	if (reverse_cookie_storage.find(uid) == reverse_cookie_storage.end())
		return false;	//理论上不会出现有Cookie没有UID的情况，仅作兜底处理
	Session index_session(cookie.substr(prefix.length()));
	auto& rcs = reverse_cookie_storage.at(uid);
	auto it = rcs.find(index_session);
	if (it == rcs.end())	//无有效Session
		return false;
	if (it->Expired())	//Cookie过期，删除该Session及对应的cookie。
	{
		rcs.erase(it);	//删除Session
		if (rcs.empty())	//如果一个用户没有任何有效的Cookie，删除该用户的记录
			reverse_cookie_storage.erase(uid);
		cookie_storage.erase(cookie);	//删除Cookie
		return false;
	}
	if (it->NeedExpiryDateRefresh())	//根据需要决定是否刷新Cookie
	{
		auto session = rcs.extract(it);
		session.value().RefreshExpiryDate();
		rcs.insert(std::move(session));
	}
	return true;	//到达此处意味着找到有效的Cookie，并已完成必要的更新工作
}

std::string SessionManager::DeleteAllCookieForUser(int uid, const std::string & current_cookie)
{
	std::lock_guard<std::mutex> lock(pool_mutex);
	if (reverse_cookie_storage.find(uid) != reverse_cookie_storage.end())
	{
		auto sessions = reverse_cookie_storage.at(uid);
		reverse_cookie_storage.erase(uid);
		for(const auto& session : sessions)
		{
			if(cookie_storage.find(session.GetCookie()) != cookie_storage.end())
				cookie_storage.erase(session.GetCookie());
		}
	}
	//返回用于响应的Cookie串（注销浏览器端的Cookie）
	return fmt::format(fmt::runtime(cookie_template.substr(prefix.length())), current_cookie, 0);
}

bool SessionManager::HasDumpedCookie(const std::string& path)
{
	namespace fs = std::filesystem;
	fs::path fn = (fs::path(path) / "unsc_sessions.json");
	std::error_code ec;
	return fs::is_regular_file(fn, ec);
}

bool SessionManager::LoadDumpedCookie(const std::string& path)
{
	namespace fs = std::filesystem;
	fs::path fn = (fs::path(path) / "unsc_sessions.json");
	std::fstream fin(fn.string(), std::ios::in);
	if(!fin.is_open())
		return false;
	try
	{
		Json::Reader reader;
		Json::Value cookies;
		if(!reader.parse(fin, cookies, false))
		{
			fin.close();
			return false;
		}
		fin.close();
		if(!cookies.isArray())
			return false;
		std::lock_guard<std::mutex> lock(pool_mutex);	//使用互斥体保证Cookie池的线程安全
		for(const auto& cookie : cookies)
		{
			Session session(cookie);
			if(session.Expired())
				continue;
			int uid = session.GetUID();
			cookie_storage.insert({ (prefix + session.GetCookie()), uid });
			if (reverse_cookie_storage.find(uid) == reverse_cookie_storage.end())
				reverse_cookie_storage.insert({ uid, { session } });
			else
				reverse_cookie_storage.at(uid).insert(session);
		}
		std::error_code ec;
		fs::remove(fn, ec);
		return true;
	}
	catch(...)
	{
		return false;
	}
}

bool SessionManager::DumpAllValidCookies(const std::string& path)
{
	namespace fs = std::filesystem;
	fs::path fn = (fs::path(path) / "unsc_sessions.json");
	std::fstream fout(fn.string(), std::ios::out | std::ios::trunc);
	if(!fout.is_open())
		return false;
	Json::Value cookies(Json::arrayValue);
	{
		std::lock_guard<std::mutex> lock(pool_mutex);	//使用互斥体保证Cookie池的线程安全
		for(const auto& [uid, sessions] : reverse_cookie_storage)
			for(const auto& session : sessions)
				cookies.append(session);
	}
	Json::FastWriter writer;
	fout << writer.write(cookies) << std::endl;
	fout.close();
	return true;
}

SessionManager GlobalSessionManager;