#include "stdafx.h"
#include "StorageFactory.h"
#include <chrono>
#include <filesystem>
#include <optional>
#include <mutex>
#include "sqlite3.h"
#include "Util.h"

struct BaseRecord
{
	int64_t id;
	string filePath;
	string dbPath;
	std::string createAt;
};

class SQLiteStorage : public IActionStorage
{
public:
	SQLiteStorage(const std::filesystem::path& backupDir, const std::string& dbName)
		: m_backupDir(backupDir),
		m_dbName(dbName),
		m_baseStorage(MakeBaseStorage())
	{
		std::vector<BaseRecord> records = GetAllBaseRecords();
		for (const auto& record : records)
		{
			m_cacheMap[record.filePath] = record.dbPath;
		}
	}

	~SQLiteStorage()
	{
		for (auto& pair : m_cacheStorage)
		{
			sqlite3_free(pair.second);
		}
	}

	bool InsertFileData(const std::string& filePath, FileData& fileData) override
	{
		if (!ExistFilePath(filePath))
		{
			CreateFileMapInfo(filePath);
		}

		sqlite3* storage = GetStorage(m_cacheMap[filePath]);
		return InsertFileData(storage, fileData);
	}

	std::optional<FileData> RetrieveFileData(const std::string& filePath) override
	{
		auto it = m_cacheMap.find(filePath);
		if (it == m_cacheMap.end())
		{
			return std::nullopt;
		}

		sqlite3* storage = GetStorage(it->second);

		auto fileDatas = GetAllFileData(storage);
		if (fileDatas.empty())
		{
			return std::nullopt;
		}

		return fileDatas[0];
	}

	bool ExistsFileData(const std::string& filePath) override
	{
		auto it = m_cacheMap.find(filePath);
		if (it == m_cacheMap.end())
		{
			return false;
		}

		sqlite3* storage = GetStorage(it->second);
		return CountFileData(storage) > 0;
	}

	void RemoveFileData(const std::string& filePath) override
	{
		auto it = m_cacheMap.find(filePath);

		if (it != m_cacheMap.end())
		{
			std::string dbName = it->second;

			std::error_code ec;
			// 删除存该文件内容的数据库
			std::filesystem::remove(dbName, ec);
		}

		// 删除映射信息
		RemoveBaseRecord(filePath);

		m_cacheMap.erase(filePath);
	}

	bool AddActionRecord(const std::string& filePath, const ActionRecord& record) override
	{
		auto it = m_cacheMap.find(filePath);
		if (it == m_cacheMap.end())
		{
			return false;
		}

		sqlite3* storage = GetStorage(it->second);
		InsertActionRecord(storage, record);

		return true;
	}

	bool ExistsActionRecord(const std::string& filePath, const std::string& uuid) override
	{
		auto it = m_cacheMap.find(filePath);
		if (it == m_cacheMap.end())
		{
			return false;
		}

		sqlite3* storage = GetStorage(it->second);
		return CountActionRecords(storage, uuid) > 0;
	}

	int ActionRecordCount(const std::string& filePath) override
	{
		auto it = m_cacheMap.find(filePath);
		if (it == m_cacheMap.end())
		{
			return false;
		}

		sqlite3* storage = GetStorage(it->second);
		return CountActionRecords(storage);
	}

	std::vector<ActionRecord> RetrieveAllActionRecords(const std::string& filePath) override
	{
		auto it = m_cacheMap.find(filePath);
		if (it == m_cacheMap.end())
		{
			return {};
		}

		sqlite3* storage = GetStorage(it->second);
		return GetAllActionRecords(storage);
	}

	void ClearBackup(const std::string& filePath) override
	{
		auto it = m_cacheMap.find(filePath);
		if (it == m_cacheMap.end())
		{
			return;
		}

		ClearStorage(it->second);
		RemoveBaseRecord(it->first);
		m_cacheMap.erase(filePath);
	}

private:
	sqlite3* MakeBaseStorage() const
	{
		std::filesystem::path fullPath = m_backupDir / m_dbName;

		std::error_code ec;
		if (!std::filesystem::exists(m_backupDir))
		{
			if (!std::filesystem::create_directory(m_backupDir, ec))
			{
				throw std::runtime_error("Create backup directory failed");
			}
		}

		sqlite3* db = nullptr;
		if (sqlite3_open(fullPath.string().c_str(), &db) != SQLITE_OK)
		{
			throw std::runtime_error("Failed to open database");
		}

		const char* createTableSQL = R"(
			CREATE TABLE IF NOT EXISTS map (
				id INTEGER PRIMARY KEY AUTOINCREMENT,
				file_path TEXT NOT NULL,
				db_path TEXT NOT NULL,
				create_at TEXT NOT NULL
			);
		)";

		char* errMsg = nullptr;
		if (sqlite3_exec(db, createTableSQL, nullptr, nullptr, &errMsg) != SQLITE_OK)
		{
			std::string err = errMsg ? errMsg : "Unknown error";
			sqlite3_free(errMsg);
			sqlite3_close(db);
			throw std::runtime_error("Failed to create table: " + err);
		}

		return db;
	}

	void CreateFileMapInfo(const std::string& filePath)
	{
		// 对应的数据库文件路径： {备份目录}/{文件名}.{UUID}.sqlite

		// 获取文件名（不带后缀）
		std::string fileNameStem = std::filesystem::path(filePath).stem().string();

		// 生成带 UUID 和新后缀的文件名
		std::string newFileName = fileNameStem + "." + GenerateUUID() + ".sqlite";

		// 构建数据库文件路径
		std::filesystem::path dbPath = m_backupDir / newFileName;

		std::string dbNameStr = dbPath.string();
		auto createAt = FormatTime(std::chrono::system_clock::now());

		const char* insertSQL = R"(
			INSERT INTO map (file_path, db_path, create_at) VALUES (?, ?, ?);
		)";

		sqlite3_stmt* stmt;
		if (sqlite3_prepare_v2(m_baseStorage, insertSQL, -1, &stmt, nullptr) != SQLITE_OK)
		{
			throw std::runtime_error("Failed to prepare insert statement");
		}

		sqlite3_bind_text(stmt, 1, filePath.c_str(), -1, SQLITE_TRANSIENT);
		sqlite3_bind_text(stmt, 2, dbNameStr.c_str(), -1, SQLITE_TRANSIENT);
		sqlite3_bind_text(stmt, 3, createAt.c_str(), -1, SQLITE_TRANSIENT);

		if (sqlite3_step(stmt) != SQLITE_DONE)
		{
			sqlite3_finalize(stmt);
			throw std::runtime_error("Failed to execute insert statement");
		}

		sqlite3_finalize(stmt);

		m_cacheMap[filePath] = dbNameStr;
	}

	bool ExistFilePath(const std::string& filePath)
	{
		return m_cacheMap.find(filePath) != m_cacheMap.end();
	}

	sqlite3* GetStorage(const std::string& dbPath)
	{
		if (m_cacheStorage.find(dbPath) == m_cacheStorage.end())
		{
			sqlite3* db = nullptr;
			if (sqlite3_open(dbPath.c_str(), &db) != SQLITE_OK)
			{
				throw std::runtime_error("Failed to open database");
			}

			const char* createFileDataTableSQL = R"(
				CREATE TABLE IF NOT EXISTS filedata (
					id INTEGER PRIMARY KEY AUTOINCREMENT,
					filepath TEXT NOT NULL,
					data BLOB NOT NULL,
					create_at TEXT NOT NULL
				);
			)";

			const char* createActionTableSQL = R"(
				CREATE TABLE IF NOT EXISTS actions (
					id INTEGER PRIMARY KEY AUTOINCREMENT,
					uuid TEXT NOT NULL,
					class_type TEXT NOT NULL,
					type TEXT NOT NULL,
					data BLOB NOT NULL,
					timestamp TEXT NOT NULL,
					create_at TEXT NOT NULL
				);
			)";

			char* errMsg = nullptr;
			if (sqlite3_exec(db, createFileDataTableSQL, nullptr, nullptr, &errMsg) != SQLITE_OK)
			{
				std::string err = errMsg ? errMsg : "Unknown error";
				sqlite3_free(errMsg);
				sqlite3_close(db);
				throw std::runtime_error("Failed to create filedata table: " + err);
			}

			if (sqlite3_exec(db, createActionTableSQL, nullptr, nullptr, &errMsg) != SQLITE_OK)
			{
				std::string err = errMsg ? errMsg : "Unknown error";
				sqlite3_free(errMsg);
				sqlite3_close(db);
				throw std::runtime_error("Failed to create actions table: " + err);
			}

			m_cacheStorage.insert({ dbPath, db });
		}

		return m_cacheStorage.at(dbPath);
	}

	void ClearStorage(const std::string& dbPath)
	{
		auto it = m_cacheStorage.find(dbPath);
		if (it != m_cacheStorage.end())
		{
			sqlite3_close(it->second);
			m_cacheStorage.erase(it);
		}

		std::error_code ec;
		CString ansiDbpath = Utf8StringToCString(dbPath);
		std::filesystem::remove(ansiDbpath.GetBuffer(), ec);
	}

	std::vector<BaseRecord> GetAllBaseRecords()
	{
		std::vector<BaseRecord> records;
		const char* selectSQL = "SELECT id, file_path, db_path, create_at FROM map";

		sqlite3_stmt* stmt;
		if (sqlite3_prepare_v2(m_baseStorage, selectSQL, -1, &stmt, nullptr) != SQLITE_OK)
		{
			throw std::runtime_error("Failed to prepare select statement");
		}

		while (sqlite3_step(stmt) == SQLITE_ROW)
		{
			BaseRecord record;
			record.id = sqlite3_column_int64(stmt, 0);
			record.filePath = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 1));
			record.dbPath = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 2));
			record.createAt = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 3));
			records.push_back(record);
		}

		sqlite3_finalize(stmt);
		return records;
	}

	std::vector<FileData> GetAllFileData(sqlite3* db)
	{
		std::vector<FileData> fileDatas;
		const char* selectSQL = "SELECT id, filepath, data, create_at FROM filedata";

		sqlite3_stmt* stmt;
		if (sqlite3_prepare_v2(db, selectSQL, -1, &stmt, nullptr) != SQLITE_OK)
		{
			throw std::runtime_error("Failed to prepare select statement");
		}

		while (sqlite3_step(stmt) == SQLITE_ROW)
		{
			FileData fileData;
			fileData.id = sqlite3_column_int64(stmt, 0);
			fileData.filepath = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 1));
			const void* data = sqlite3_column_blob(stmt, 2);
			int dataSize = sqlite3_column_bytes(stmt, 2);
			fileData.data.assign(static_cast<const char*>(data), static_cast<const char*>(data) + dataSize);
			fileData.createAt = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 3));
			fileDatas.push_back(fileData);
		}

		sqlite3_finalize(stmt);
		return fileDatas;
	}

	bool InsertFileData(sqlite3* db, FileData& fileData)
	{
		const char* deleteSQL = "DELETE FROM filedata";
		char* errMsg = nullptr;
		if (sqlite3_exec(db, deleteSQL, nullptr, nullptr, &errMsg) != SQLITE_OK)
		{
			std::string err = errMsg ? errMsg : "Unknown error";
			sqlite3_free(errMsg);
			throw std::runtime_error("Failed to delete existing filedata: " + err);
		}

		const char* insertSQL = "INSERT INTO filedata (filepath, data, create_at) VALUES (?, ?, ?)";
		sqlite3_stmt* stmt;
		if (sqlite3_prepare_v2(db, insertSQL, -1, &stmt, nullptr) != SQLITE_OK)
		{
			throw std::runtime_error("Failed to prepare insert statement");
		}

		sqlite3_bind_text(stmt, 1, fileData.filepath.c_str(), -1, SQLITE_TRANSIENT);
		sqlite3_bind_blob(stmt, 2, fileData.data.data(), static_cast<int>(fileData.data.size()), SQLITE_TRANSIENT);
		sqlite3_bind_text(stmt, 3, fileData.createAt.c_str(), -1, SQLITE_TRANSIENT);

		if (sqlite3_step(stmt) != SQLITE_DONE)
		{
			sqlite3_finalize(stmt);
			throw std::runtime_error("Failed to execute insert statement");
		}

		sqlite3_finalize(stmt);

		return true;
	}

	int CountFileData(sqlite3* db)
	{
		const char* countSQL = "SELECT COUNT(*) FROM filedata";
		sqlite3_stmt* stmt;
		if (sqlite3_prepare_v2(db, countSQL, -1, &stmt, nullptr) != SQLITE_OK)
		{
			throw std::runtime_error("Failed to prepare count statement");
		}

		int count = 0;
		if (sqlite3_step(stmt) == SQLITE_ROW)
		{
			count = sqlite3_column_int(stmt, 0);
		}

		sqlite3_finalize(stmt);
		return count;
	}

	void RemoveBaseRecord(const std::string& filePath)
	{
		const char* deleteSQL = "DELETE FROM map WHERE file_path = ?";
		sqlite3_stmt* stmt;
		if (sqlite3_prepare_v2(m_baseStorage, deleteSQL, -1, &stmt, nullptr) != SQLITE_OK)
		{
			throw std::runtime_error("Failed to prepare delete statement");
		}

		sqlite3_bind_text(stmt, 1, filePath.c_str(), -1, SQLITE_TRANSIENT);

		if (sqlite3_step(stmt) != SQLITE_DONE)
		{
			sqlite3_finalize(stmt);
			throw std::runtime_error("Failed to execute delete statement");
		}

		sqlite3_finalize(stmt);
	}

	void InsertActionRecord(sqlite3* db, const ActionRecord& record)
	{
		const char* insertSQL = R"(
			INSERT INTO actions (uuid, class_type, type, data, timestamp, create_at)
			VALUES (?, ?, ?, ?, ?, ?)
		)";

		sqlite3_stmt* stmt;
		if (sqlite3_prepare_v2(db, insertSQL, -1, &stmt, nullptr) != SQLITE_OK)
		{
			throw std::runtime_error("Failed to prepare insert statement");
		}

		sqlite3_bind_text(stmt, 1, record.uuid.c_str(), -1, SQLITE_TRANSIENT);
		sqlite3_bind_text(stmt, 2, record.classType.c_str(), -1, SQLITE_TRANSIENT);
		sqlite3_bind_text(stmt, 3, record.type.c_str(), -1, SQLITE_TRANSIENT);
		sqlite3_bind_blob(stmt, 4, record.data.data(), static_cast<int>(record.data.size()), SQLITE_TRANSIENT);
		sqlite3_bind_text(stmt, 5, record.timestamp.c_str(), -1, SQLITE_TRANSIENT);
		sqlite3_bind_text(stmt, 6, record.createAt.c_str(), -1, SQLITE_TRANSIENT);

		if (sqlite3_step(stmt) != SQLITE_DONE)
		{
			sqlite3_finalize(stmt);
			throw std::runtime_error("Failed to execute insert statement");
		}

		sqlite3_finalize(stmt);
	}

	int CountActionRecords(sqlite3* db)
	{
		const char* countSQL = "SELECT COUNT(*) FROM actions";
		sqlite3_stmt* stmt;
		if (sqlite3_prepare_v2(db, countSQL, -1, &stmt, nullptr) != SQLITE_OK)
		{
			throw std::runtime_error("Failed to prepare count statement");
		}

		int count = 0;
		if (sqlite3_step(stmt) == SQLITE_ROW)
		{
			count = sqlite3_column_int(stmt, 0);
		}

		sqlite3_finalize(stmt);
		return count;
	}

	int CountActionRecords(sqlite3* db, const std::string& uuid)
	{
		const char* countSQL = "SELECT COUNT(*) FROM actions WHERE uuid = ?";
		sqlite3_stmt* stmt;
		if (sqlite3_prepare_v2(db, countSQL, -1, &stmt, nullptr) != SQLITE_OK)
		{
			throw std::runtime_error("Failed to prepare count statement");
		}

		sqlite3_bind_text(stmt, 1, uuid.c_str(), -1, SQLITE_TRANSIENT);

		int count = 0;
		if (sqlite3_step(stmt) == SQLITE_ROW)
		{
			count = sqlite3_column_int(stmt, 0);
		}

		sqlite3_finalize(stmt);
		return count;
	}

	std::vector<ActionRecord> GetAllActionRecords(sqlite3* db)
	{
		std::vector<ActionRecord> records;
		const char* selectSQL = "SELECT id, uuid, class_type, type, data, timestamp, create_at FROM actions";

		sqlite3_stmt* stmt;
		if (sqlite3_prepare_v2(db, selectSQL, -1, &stmt, nullptr) != SQLITE_OK)
		{
			throw std::runtime_error("Failed to prepare select statement");
		}

		while (sqlite3_step(stmt) == SQLITE_ROW)
		{
			ActionRecord record;
			record.id = sqlite3_column_int64(stmt, 0);
			record.uuid = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 1));
			record.classType = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 2));
			record.type = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 3));
			const void* data = sqlite3_column_blob(stmt, 4);
			int dataSize = sqlite3_column_bytes(stmt, 4);
			record.data.assign(static_cast<const char*>(data), static_cast<const char*>(data) + dataSize);
			record.timestamp = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 5));
			record.createAt = reinterpret_cast<const char*>(sqlite3_column_text(stmt, 6));
			records.push_back(record);
		}

		sqlite3_finalize(stmt);
		return records;
	}

private:
	std::filesystem::path m_backupDir;
	std::string m_dbName;
	sqlite3* m_baseStorage;

	std::unordered_map<std::string, std::string> m_cacheMap;
	std::unordered_map<std::string, sqlite3*> m_cacheStorage;
};

// 工厂类用于创建存储实例
StorageFactory& StorageFactory::GetInstance()
{
	static StorageFactory factory;
	return factory;
}

/**
 * 尽量放到 %appdata% 目录，这是现代的主流玩法，放工作目录，工作目录是变动的，放 exe 所在目录，则有可能没权限，
 * 这个目录本身就是用来给 app 存放数据的
 * 如果获取失败(按理说不应该失败)，则放到工作目录下
 * 
 * \return 
 */
static std::filesystem::path GetBackupPath()
{
	PWSTR RoamingPath = nullptr;
	HRESULT result = SHGetKnownFolderPath(FOLDERID_RoamingAppData, 0, nullptr, &RoamingPath);
	FinalAction finalAction([RoamingPath]() { CoTaskMemFree(RoamingPath); });

	if (result == S_OK)
	{
		std::filesystem::path path = RoamingPath;
		path /= "KEVisualization";
		path /= "backup";

		if (!std::filesystem::exists(path))
		{
			std::filesystem::create_directories(path);
		}

		return path;
	}
	else
	{
		return "backup";
	}
}

std::shared_ptr<IActionStorage> StorageFactory::CreateDefaultStorage()
{
	std::filesystem::path backup = GetBackupPath();

	return StorageFactory::GetInstance().CreateStorage(backup, "base.sqlite");
}

std::shared_ptr<IActionStorage> StorageFactory::CreateStorage(const std::filesystem::path& backupDir, const std::string& dbName)
{
	try
	{
		std::filesystem::path dbPath = (backupDir / dbName).string();

		auto it = m_map.find(dbPath.string());
		if (it != m_map.end())
		{
			return it->second;
		}

		auto storage = std::make_shared<SQLiteStorage>(backupDir, dbName);
		m_map[dbPath.string()] = storage;

		return storage;
	}
	catch (std::runtime_error &e)
	{
		TRACE("Create storage failed: %s\n", e.what());
		return nullptr;
	}
}