#pragma once

#include "kv_store.hpp"
#include "request.hpp"
#include <random>
#include <unordered_map>
#include <mutex>
#include <array>
#include <chrono>
#include <optional>
#include <atomic>
#include <list>

class SessionStore {
private:
	static constexpr size_t SHARD_COUNT = 16;
	static constexpr size_t MAX_SESSIONS_PER_SHARD = 64;
	static constexpr size_t MAX_SESSION_DATA_SIZE = 16;
	static constexpr size_t MAX_VALUE_SIZE = 1024;
	static constexpr auto SESSION_TIMEOUT = std::chrono::hours(24);

	struct SessionData {
		std::unordered_map<std::string, std::string> data;
		std::chrono::steady_clock::time_point last_access;
		std::list<std::string>::iterator lru_iter;
		bool dirty = false;
	};

	std::array<std::mutex, SHARD_COUNT> mutexes_;
	std::array<std::unordered_map<std::string, SessionData>, SHARD_COUNT> shards_;
	std::array<std::list<std::string>, SHARD_COUNT> lru_lists_;
	KeyValueStore store_;
	std::mt19937 rng_;

	size_t get_shard(const std::string& id) const {
		return std::hash<std::string>{}(id) % SHARD_COUNT;
	}

	void evict_lru_session(size_t shard) {
		auto& shard_map = shards_[shard];
		auto& lru_list = lru_lists_[shard];

		if (lru_list.empty()) return;

		std::string oldest_id = lru_list.back();
		lru_list.pop_back();

		auto it = shard_map.find(oldest_id);
		if (it != shard_map.end()) {
			// Save dirty session before eviction
			if (it->second.dirty) {
				for (const auto& [key, value] : it->second.data) {
					store_.set(oldest_id + ":" + key, value);
				}
			}
			shard_map.erase(it);
		}
	}

	void touch_session(const std::string& id, size_t shard) {
		auto& session = shards_[shard][id];
		auto& lru_list = lru_lists_[shard];

		// Move to front of LRU
		lru_list.erase(session.lru_iter);
		lru_list.push_front(id);
		session.lru_iter = lru_list.begin();
		session.last_access = std::chrono::steady_clock::now();
	}

	void lazy_load_session(const std::string& id, size_t shard) {
		auto& shard_map = shards_[shard];
		if (shard_map.find(id) != shard_map.end()) return;

		// Check session limit
		if (shard_map.size() >= MAX_SESSIONS_PER_SHARD) {
			evict_lru_session(shard);
		}

		SessionData session;
		session.last_access = std::chrono::steady_clock::now();

		// Load from persistent store
		for (const auto& [key, value] : store_.get_all()) {
			if (key.starts_with(id + ":")) {
				std::string session_key = key.substr(id.length() + 1);
				if (auto* str_val = std::get_if<std::string>(&value)) {
					if (session.data.size() < MAX_SESSION_DATA_SIZE) {
						session.data[session_key] = *str_val;
					}
				}
			}
		}

		// Add to LRU
		lru_lists_[shard].push_front(id);
		session.lru_iter = lru_lists_[shard].begin();

		shard_map[id] = std::move(session);
	}

	void cleanup_expired_sessions() {
		auto now = std::chrono::steady_clock::now();

		for (size_t i = 0; i < SHARD_COUNT; ++i) {
			std::lock_guard<std::mutex> lock(mutexes_[i]);
			auto& shard = shards_[i];
			auto& lru_list = lru_lists_[i];

			for (auto it = shard.begin(); it != shard.end();) {
				if (now - it->second.last_access > SESSION_TIMEOUT) {
					lru_list.erase(it->second.lru_iter);
					it = shard.erase(it);
				} else {
					++it;
				}
			}
		}
	}

public:
	SessionStore() : rng_(std::random_device{}()) {
		store_.set_file("sessions.txt");
	}

	std::string create() {
		std::string id = generate_id();
		size_t shard = get_shard(id);

		std::lock_guard<std::mutex> lock(mutexes_[shard]);

		// Check session limit
		if (shards_[shard].size() >= MAX_SESSIONS_PER_SHARD) {
			evict_lru_session(shard);
		}

		SessionData session;
		session.last_access = std::chrono::steady_clock::now();
		session.dirty = true;

		// Add to LRU
		lru_lists_[shard].push_front(id);
		session.lru_iter = lru_lists_[shard].begin();

		shards_[shard][id] = std::move(session);
		return id;
	}

	std::optional<std::string> get(const std::string& id, const std::string& key) {
		size_t shard = get_shard(id);
		std::lock_guard<std::mutex> lock(mutexes_[shard]);

		lazy_load_session(id, shard);

		auto it = shards_[shard].find(id);
		if (it == shards_[shard].end()) return std::nullopt;

		touch_session(id, shard);

		auto data_it = it->second.data.find(key);
		return data_it != it->second.data.end() ?
			std::optional<std::string>{data_it->second} : std::nullopt;
	}

	void set(const std::string& id, const std::string& key, const std::string& value) {
		if (value.size() > MAX_VALUE_SIZE) return;

		size_t shard = get_shard(id);
		std::lock_guard<std::mutex> lock(mutexes_[shard]);

		lazy_load_session(id, shard);

		auto& session = shards_[shard][id];

		// Check data limit
		if (session.data.size() >= MAX_SESSION_DATA_SIZE && session.data.find(key) == session.data.end()) {
			return;
		}

		session.data[key] = value;
		touch_session(id, shard);
		session.dirty = true;
	}

	bool del(const std::string& id) {
		size_t shard = get_shard(id);
		std::lock_guard<std::mutex> lock(mutexes_[shard]);

		auto it = shards_[shard].find(id);
		if (it == shards_[shard].end()) return false;

		lru_lists_[shard].erase(it->second.lru_iter);
		shards_[shard].erase(it);

		// Remove from persistent store
		for (const auto& [key, value] : store_.get_all()) {
			if (key.starts_with(id + ":")) {
				store_.del(key);
			}
		}

		return true;
	}

	std::string refresh(const std::string& old_id) {
		size_t old_shard = get_shard(old_id);
		std::lock_guard<std::mutex> lock(mutexes_[old_shard]);

		auto it = shards_[old_shard].find(old_id);
		if (it == shards_[old_shard].end()) return old_id;

		// Generate new ID
		std::string new_id = generate_id();
		size_t new_shard = get_shard(new_id);

		// Handle cross-shard case
		if (new_shard != old_shard) {
			std::lock_guard<std::mutex> new_lock(mutexes_[new_shard]);
			
			// Check session limit on new shard
			if (shards_[new_shard].size() >= MAX_SESSIONS_PER_SHARD) {
				evict_lru_session(new_shard);
			}

			// Move session data
			SessionData session = std::move(it->second);
			lru_lists_[old_shard].erase(session.lru_iter);
			
			// Add to new shard
			lru_lists_[new_shard].push_front(new_id);
			session.lru_iter = lru_lists_[new_shard].begin();
			session.last_access = std::chrono::steady_clock::now();
			session.dirty = true;
			
			shards_[new_shard][new_id] = std::move(session);
			shards_[old_shard].erase(it);
		} else {
			// Same shard - just update the key
			SessionData session = std::move(it->second);
			lru_lists_[old_shard].erase(session.lru_iter);
			
			lru_lists_[old_shard].push_front(new_id);
			session.lru_iter = lru_lists_[old_shard].begin();
			session.last_access = std::chrono::steady_clock::now();
			session.dirty = true;
			
			shards_[old_shard][new_id] = std::move(session);
			shards_[old_shard].erase(it);
		}

		return new_id;
	}

	size_t size() {
		size_t total = 0;
		for (size_t i = 0; i < SHARD_COUNT; ++i) {
			std::lock_guard<std::mutex> lock(mutexes_[i]);
			total += shards_[i].size();
		}
		return total;
	}

	void set_file(const std::string& filename) {
		store_.set_file(filename);
	}

	void load() {
		store_.load();
	}

	void save() {
		for (size_t i = 0; i < SHARD_COUNT; ++i) {
			std::lock_guard<std::mutex> lock(mutexes_[i]);

			for (auto& [id, session] : shards_[i]) {
				if (!session.dirty) continue;

				for (const auto& [key, value] : session.data) {
					store_.set(id + ":" + key, value);
				}
				session.dirty = false;
			}
		}

		cleanup_expired_sessions();
		store_.save();
	}

	static std::string_view extract_session_id(const Request& req) {
		return req.get_cookie("session_id");
	}

private:
	std::string generate_id() {
		static const char chars[] = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
		std::string id;
		id.reserve(16);

		for (int i = 0; i < 16; ++i) {
			id += chars[rng_() % (sizeof(chars) - 1)];
		}

		return id;
	}
};