#pragma once

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

class SessionStore {
private:
	static constexpr size_t SHARD_COUNT = 16;
	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;
		bool dirty = false;
	};

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

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

	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;

		// Load from persistent store
		SessionData session;
		session.last_access = std::chrono::steady_clock::now();

		// Try to load existing session data
		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)) {
					session.data[session_key] = *str_val;
				}
			}
		}

		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];

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

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

	bool needs_session(const Request& req) const {
		// Skip sessions for static content
		if (req.path.starts_with("/static") ||
		    req.path.starts_with("/assets") ||
		    req.path.ends_with(".css") ||
		    req.path.ends_with(".js") ||
		    req.path.ends_with(".png") ||
		    req.path.ends_with(".jpg") ||
		    req.path.ends_with(".jpeg") ||
		    req.path.ends_with(".gif") ||
		    req.path.ends_with(".ico") ||
		    req.path.ends_with(".svg") ||
		    req.path.ends_with(".woff") ||
		    req.path.ends_with(".woff2") ||
		    req.path.ends_with(".ttf")) {
			return false;
		}

		// Need sessions for API and admin routes
		return req.path.starts_with("/api") ||
		       req.path.starts_with("/admin") ||
		       req.method == HttpMethod::POST ||
		       req.method == HttpMethod::PUT ||
		       req.method == HttpMethod::DELETE;
	}

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

		std::lock_guard<std::mutex> lock(mutexes_[shard]);
		SessionData session;
		session.last_access = std::chrono::steady_clock::now();
		session.dirty = true;
		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;

		it->second.last_access = std::chrono::steady_clock::now();
		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) {
		size_t shard = get_shard(id);
		std::lock_guard<std::mutex> lock(mutexes_[shard]);

		lazy_load_session(id, shard);

		auto& session = shards_[shard][id];
		session.data[key] = value;
		session.last_access = std::chrono::steady_clock::now();
		session.dirty = true;
	}

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

		bool found = shards_[shard].erase(id) > 0;

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

		return found;
	}

	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();
		// Sessions loaded lazily on first access
	}

	void save() {
		// Save only dirty sessions
		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();
	}

	// Use existing cookie helpers
	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;
	}
};