#pragma once

#include <unordered_map>
#include <shared_mutex>
#include <fstream>
#include <optional>
#include <string>
#include <variant>
#include <array>
#include <mutex>

using std::string;

class KeyValueStore {
public:
	using Value = std::variant<int, double, string, bool>;

private:
	static constexpr size_t SHARD_COUNT = 16;

	std::array<std::shared_mutex, SHARD_COUNT> mutexes_;
	std::array<std::unordered_map<string, Value>, SHARD_COUNT> shards_;
	string filename_ = "store.txt";

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

public:
	void set(const string& key, const Value& value) {
		size_t shard = get_shard(key);
		std::unique_lock<std::shared_mutex> lock(mutexes_[shard]);
		shards_[shard][key] = value;
	}

	template<typename T>
	std::optional<T> get(const string& key) {
		size_t shard = get_shard(key);
		std::shared_lock<std::shared_mutex> lock(mutexes_[shard]);
		auto it = shards_[shard].find(key);
		if (it == shards_[shard].end()) return std::nullopt;
		if (auto* ptr = std::get_if<T>(&it->second)) {
			return *ptr;
		}
		return std::nullopt;
	}

	std::unordered_map<string, Value> get_all() {
		std::unordered_map<string, Value> result;

		for (size_t i = 0; i < SHARD_COUNT; ++i) {
			std::shared_lock<std::shared_mutex> lock(mutexes_[i]);
			for (const auto& [key, value] : shards_[i]) {
				result[key] = value;
			}
		}
		return result;
	}

	bool del(const string& key) {
		size_t shard = get_shard(key);
		std::unique_lock<std::shared_mutex> lock(mutexes_[shard]);
		return shards_[shard].erase(key) > 0;
	}

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

	void set_file(const string& filename) {
		filename_ = filename;
	}

	void load() {
		std::ifstream file(filename_);
		if (!file.is_open()) return;

		string line;
		while (std::getline(file, line)) {
			size_t eq = line.find('=');
			if (eq == string::npos || eq < 2) continue;

			string key = line.substr(2, eq - 2);
			string val_str = line.substr(eq + 1);
			char type = line[0];

			Value value;
			switch (type) {
				case 'i': value = std::stoi(val_str); break;
				case 'd': value = std::stod(val_str); break;
				case 's': value = val_str; break;
				case 'b': value = (val_str == "1"); break;
				default: continue;
			}

			size_t shard = get_shard(key);
			std::unique_lock<std::shared_mutex> lock(mutexes_[shard]);
			shards_[shard][key] = std::move(value);
		}
	}

	void save() {
		std::ofstream file(filename_);
		if (!file.is_open()) return;

		for (size_t i = 0; i < SHARD_COUNT; ++i) {
			std::shared_lock<std::shared_mutex> lock(mutexes_[i]);
			for (const auto& [key, value] : shards_[i]) {
				file << serialize_entry(key, value) << "\n";
			}
		}
	}

private:
	string serialize_entry(const string& key, const Value& value) {
		return std::visit([&key](const auto& val) -> string {
			using T = std::decay_t<decltype(val)>;
			if constexpr (std::is_same_v<T, int>) {
				return "i:" + key + "=" + std::to_string(val);
			} else if constexpr (std::is_same_v<T, double>) {
				return "d:" + key + "=" + std::to_string(val);
			} else if constexpr (std::is_same_v<T, string>) {
				return "s:" + key + "=" + val;
			} else if constexpr (std::is_same_v<T, bool>) {
				return "b:" + key + "=" + (val ? "1" : "0");
			}
		}, value);
	}
};