#include "CodegenModel.hpp"
#include "CodegenModelArchive.hpp"

#include "CodegenUtils.hpp"

#include <Macros.hpp>
#include <ScopeGuard.hpp>
#include <Utils.hpp>

#include <robin_hood.h>
#include <cassert>
#include <cstdint>
#include <stdexcept>
#include <string>
#include <variant>

using namespace std::literals;

struct SomeDecl {
	std::variant<DeclStruct, DeclFunction, DeclEnum> v;
};

class CodegenModel::Private {
	friend class CodegenModelArchive;

public:
	// We want address stability for everything
	robin_hood::unordered_node_map<std::string, SomeDecl, StringHash, StringEqual> decls;
	robin_hood::unordered_node_map<std::string, DeclNamespace, StringHash, StringEqual> namespaces;
};

// A number for `PRAGMA user_vesrion`, representing the current database version. Increment when the table format changes.
#define CURRENT_DATABASE_VERSION 1
constexpr int64_t kGlobalNamespaceId = 1;

struct SQLiteDatabase {
	sqlite3* database = nullptr;

	~SQLiteDatabase() {
		// NOTE: calling with NULL is a harmless no-op
		int result = sqlite3_close(database);
		assert(result == SQLITE_OK);
	}

	operator sqlite3*() const { return database; }
	sqlite3** operator&() { return &database; }
};

struct SQLiteStatement {
	sqlite3_stmt* stmt = nullptr;

	SQLiteStatement(const SQLiteStatement&) = delete;
	SQLiteStatement& operator=(const SQLiteStatement&) = delete;

	SQLiteStatement() = default;

	~SQLiteStatement() {
		// NOTE: calling with NULL is a harmless no-op
		// NOTE: we don't care about the error code, because they are returned if the statement has errored in the most recent execution
		//       but deleting it will succeeed anyways
		sqlite3_finalize(stmt);
	}

	operator sqlite3_stmt*() const { return stmt; }
	sqlite3_stmt** operator&() { return &stmt; }

	bool InitializeLazily(sqlite3* database, std::string_view sql) {
		if (!stmt) {
			int result = sqlite3_prepare(database, sql.data(), sql.size(), &stmt, nullptr);
			if (result != SQLITE_OK) {
				INPLACE_FMT(msg, "Failed to prepare statement, error message: %s", sqlite3_errmsg(database));
				throw std::runtime_error(msg);
			}

			return true;
		}
		return false;
	}
};

namespace {
void PrintErrMsgIfPresent(char*& errMsg) {
	if (errMsg) {
		printf("SQLite error: %s\n", errMsg);
		sqlite3_free(errMsg);
	}
}
} // namespace

class CodegenModelArchive::Private {
	friend class CodegenModel;

public:
	// NOTE: this must be the first field, because we want it to destruct after all other statement fields
	SQLiteDatabase database;
	/* Core Statements */
	SQLiteStatement beginTransactionStmt;
	SQLiteStatement commitTransactionStmt;
	SQLiteStatement rollbackTransactionStmt;
	/* Component Statements, initalized on demand */
	SQLiteStatement findNamespaceStmt;
	SQLiteStatement findOrStoreNamespaceStmt;
	SQLiteStatement storeEnumStmt;
	SQLiteStatement storeEnumElmStmt;
	SQLiteStatement deleteStructDeclByFilenameStmt;
	SQLiteStatement deleteEnumDeclByFilenameStmt;

	void InitializeCoreStatements() {
		char* errMsg = nullptr;

		int result = sqlite3_exec(database, "PRAGMA user_version = " STRINGIFY(CURRENT_DATABASE_VERSION), nullptr, nullptr, &errMsg);
		PrintErrMsgIfPresent(errMsg);
		assert(result == SQLITE_OK);

		result = sqlite3_exec(database, R"""(
BEGIN TRANSACTION;
CREATE TABLE Namespaces(
	-- NOTE: SQLite forbids foreign keys referencing the implicit `rowid` column, we have to create an alias for it
	Id INTEGER PRIMARY KEY,
	ParentNamespaceId INTEGER REFERENCES Namespaces(Id),
	Name TEXT,
	UNIQUE (ParentNamespaceId, Name)
);

CREATE TABLE DeclStructs(
	Id INTEGER PRIMARY KEY,
	NamespaceId INTEGER REFERENCES Namespaces(Id),
	ParentStructId INTEGER REFERENCES DeclStructs(Id),
	Name TEXT,
	FileName TEXT,
	UNIQUE (NamespaceId, Name)
);

CREATE TABLE DeclEnums(
	Id INTEGER PRIMARY KEY,
	NamespaceId INTEGER REFERENCES Namespaces(Id),
	Name TEXT,
	UnderlyingType TEXT,
	FileName TEXT,
	UNIQUE (NamespaceId, Name)
);
CREATE TABLE DeclEnumElements(
	EnumId INTEGER REFERENCES DeclEnums(Id) ON DELETE CASCADE,
	Name TEXT,
	Value INTEGER,
	UNIQUE (EnumId, Name)
);

CREATE INDEX Index_DeclStructs_FileName
	ON DeclStructs(FileName);
CREATE INDEX Index_DeclEnums_FileName
	ON DeclEnums(FileName);

INSERT INTO Namespaces(Id, ParentNamespaceId, Name)
VALUES
	-- Special global namespace that has no parent, and Id should always be 1
	(1, NULL, "<global namespace>");

COMMIT TRANSACTION;
)""",
			nullptr,
			nullptr,
			&errMsg);
		PrintErrMsgIfPresent(errMsg);
		assert(result == SQLITE_OK);
	}

	void BeginTransaction() {
		int result = sqlite3_step(beginTransactionStmt);
		assert(result == SQLITE_DONE);
		sqlite3_reset(beginTransactionStmt);
	}

	void CommitTransaction() {
		int result = sqlite3_step(commitTransactionStmt);
		assert(result == SQLITE_DONE);
		sqlite3_reset(commitTransactionStmt);
	}

	void RollbackTransaction() {
		int result = sqlite3_step(rollbackTransactionStmt);
		assert(result == SQLITE_DONE);
		sqlite3_reset(rollbackTransactionStmt);
	}

	/// \return Row ID of the namespace, or 0 if it currently doesn't exist.
	int64_t FindNamespace(const DeclNamespace* ns) {
		if (!ns) {
			return kGlobalNamespaceId;
		}

		InitializeStmt_findNamespaceStmt();
		return FindNamespaceImpl(*ns);
	}

	/// \return Row ID of the namespace.
	int64_t FindOrStoreNamespace(const DeclNamespace* ns) {
		if (!ns) {
			return kGlobalNamespaceId;
		}

		InitializeStmt_findNamespaceStmt();
		InitializeStmt_findOrStoreNamespaceStmt();
		return FindOrStoreNamespaceImpl(*ns);
	}

private:
	void InitializeStmt_findNamespaceStmt() {
		findNamespaceStmt.InitializeLazily(database, R"""(
SELECT Id FROM Namespaces WHERE ParentNamespaceId = ?1 AND Name = ?2;
)"""sv);
	}

	void InitializeStmt_findOrStoreNamespaceStmt() {
		findOrStoreNamespaceStmt.InitializeLazily(database, R"""(
INSERT INTO Namespaces(ParentNamespaceId, Name)
	VALUES (?1, ?2)
	RETURNING Id;
)"""sv);
	}

	int64_t FindNamespaceImpl(const DeclNamespace& ns) {
		int64_t parentNsRowId;
		if (ns.container) {
			parentNsRowId = FindNamespaceImpl(*ns.container);
			if (parentNsRowId == 0) {
				// Parent namespace doesn't exist in database, shortcircuit
				return 0;
			}
		} else {
			parentNsRowId = kGlobalNamespaceId;
		}

		return FindNamespaceImpl(ns, parentNsRowId);
	}

	int64_t FindNamespaceImpl(const DeclNamespace& ns, int64_t parentNsRowId) {
		sqlite3_stmt* stmt = findNamespaceStmt;
		sqlite3_bind_int64(stmt, 1, parentNsRowId);
		sqlite3_bind_text(stmt, 2, ns.name.c_str(), ns.name.size(), nullptr);
		int result = sqlite3_step(stmt);
		if (result == SQLITE_ROW) {
			int64_t rowId = sqlite3_column_int64(stmt, 0);
			sqlite3_reset(stmt);

			return rowId;
		} else {
			return 0;
		}
	}

	int64_t FindOrStoreNamespaceImpl(const DeclNamespace& ns) {
		if (auto rowId = FindNamespaceImpl(ns); rowId != 0) {
			return rowId;
		}

		sqlite3_stmt* stmt = findOrStoreNamespaceStmt;
		int64_t parentRowId = ns.container
			? FindOrStoreNamespaceImpl(*ns.container)
			: kGlobalNamespaceId;
		sqlite3_bind_int64(stmt, 1, parentRowId);
		sqlite3_bind_text(stmt, 2, ns.name.c_str(), ns.name.size(), nullptr);
		int result = sqlite3_step(stmt);
		assert(result == SQLITE_ROW);

		auto rowId = sqlite3_column_int64(stmt, 0);
		sqlite3_reset(stmt);

		return rowId;
	}
};

CodegenModel::CodegenModel()
	: m{ new Private() } //
{
}

CodegenModel::~CodegenModel() {
	delete m;
}

#define STORE_DECL_OF_TYPE(DeclType, fullname, decl)                                                    \
	auto [iter, success] = m->decls.try_emplace(std::move(fullname), SomeDecl{ .v = std::move(decl) }); \
	auto& key = iter->first;                                                                            \
	auto& val = iter->second;                                                                           \
	auto& declRef = std::get<DeclType>(val.v);                                                          \
	declRef.fullname = key;                                                                             \
	return &declRef

DeclEnum* CodegenModel::AddEnum(std::string fullname, DeclEnum decl) {
#if CODEGEN_DEBUG_PRINT
	printf("Committed enum '%s'\n", decl.name.c_str());
	for (auto& elm : decl.elements) {
		printf("  - element %s = %" PRId64 "\n", elm.name.c_str(), elm.value);
	}
#endif

	STORE_DECL_OF_TYPE(DeclEnum, fullname, decl);
}

DeclStruct* CodegenModel::AddStruct(std::string fullname, DeclStruct decl) {
#if CODEGEN_DEBUG_PRINT
	printf("Committed struct '%s'\n", decl.name.c_str());
	printf("  Base classes:\n");
	for (auto& base : decl.baseClasses) {
		printf("  - %.*s\n", PRINTF_STRING_VIEW(base->name));
	}
#endif

	STORE_DECL_OF_TYPE(DeclStruct, fullname, decl);
}

#define FIND_DECL_OF_TYPE(DeclType)                     \
	auto iter = m->decls.find(name);                    \
	if (iter != m->decls.end()) {                       \
		auto& some = iter->second.v;                    \
		if (auto decl = std::get_if<DeclType>(&some)) { \
			return decl;                                \
		}                                               \
	}                                                   \
	return nullptr

const DeclEnum* CodegenModel::FindEnum(std::string_view name) const {
	FIND_DECL_OF_TYPE(DeclEnum);
}

const DeclStruct* CodegenModel::FindStruct(std::string_view name) const {
	FIND_DECL_OF_TYPE(DeclStruct);
}

DeclNamespace* CodegenModel::AddNamespace(DeclNamespace ns) {
	auto path = Utils::MakeFullName(""sv, &ns);
	auto [iter, success] = m->namespaces.try_emplace(std::move(path), std::move(ns));
	auto& nsRef = iter->second;
	if (success) {
		nsRef.fullname = iter->first;
	}
	return &nsRef;
}

const DeclNamespace* CodegenModel::FindNamespace(std::string_view fullname) const {
	auto iter = m->namespaces.find(fullname);
	if (iter != m->namespaces.end()) {
		return &iter->second;
	} else {
		return nullptr;
	}
}

DeclNamespace* CodegenModel::FindNamespace(std::string_view name) {
	return const_cast<DeclNamespace*>(const_cast<const CodegenModel*>(this)->FindNamespace(name));
}

CodegenModelArchive::CodegenModelArchive(std::string_view dbPath)
	: m{ new Private() } //
{
	std::string zstrPath(dbPath);
	int reuslt = sqlite3_open(zstrPath.c_str(), &m->database);
	if (reuslt != SQLITE_OK) {
		std::string msg;
		msg += "Failed to open SQLite3 database, error message:\n";
		msg += sqlite3_errmsg(m->database);
		throw std::runtime_error(msg);
	}

	{
		SQLiteStatement readVersionStmt;
		readVersionStmt.InitializeLazily(m->database, "PRAGMA user_version"sv);

		int result = sqlite3_step(readVersionStmt);
		assert(result == SQLITE_ROW);
		int currentDatabaseVersion = sqlite3_column_int(readVersionStmt, 0);

		result = sqlite3_step(readVersionStmt);
		assert(result == SQLITE_DONE);

		if (currentDatabaseVersion == 0) {
			// Newly created database, initialize it
			m->InitializeCoreStatements();
		} else if (currentDatabaseVersion == CURRENT_DATABASE_VERSION) {
			// Same version, no need to do anything
		} else {
			INPLACE_FMT(msg, "Incompatbile database versions %d (in file) vs %d (expected).", currentDatabaseVersion, CURRENT_DATABASE_VERSION);
			throw std::runtime_error(msg);
		}
	}

	// NOTE: These pragmas are not persistent, so we need to set them every time
	// SQLite3 as of 2022-06-24 defaults to foreign_keys = OFF, we need this to be on for ON DELETE CASCADE and etc. to work
	sqlite3_exec(m->database, "PRAGMA foreign_keys = ON", nullptr, nullptr, nullptr);
	// This database is used for a buildsystem and can be regenerated at any time. We don't care for the slightest about data integrity, we just want fast updates
	sqlite3_exec(m->database, "PRAGMA synchronous = OFF", nullptr, nullptr, nullptr);
	sqlite3_exec(m->database, "PRAGMA journal_mode = MEMORY", nullptr, nullptr, nullptr);

	m->beginTransactionStmt.InitializeLazily(m->database, "BEGIN TRANSACTION");
	m->commitTransactionStmt.InitializeLazily(m->database, "COMMIT TRANSACTION");
	m->rollbackTransactionStmt.InitializeLazily(m->database, "ROLLBACK TRANSACTION");
}

CodegenModelArchive::~CodegenModelArchive() {
	delete m;
}

void CodegenModelArchive::DeleteDeclsRelatedToFile(std::string_view filename) {
	// -Argument-	-Description-
	// ?1			The filename to delete
	m->deleteStructDeclByFilenameStmt.InitializeLazily(m->database, "DELETE FROM DeclStructs WHERE FileName = ?1;"sv);
	m->deleteEnumDeclByFilenameStmt.InitializeLazily(m->database, "DELETE FROM DeclEnums WHERE FileName = ?1;"sv);

	m->BeginTransaction();
	auto stmtList = {
		m->deleteStructDeclByFilenameStmt.stmt,
		m->deleteEnumDeclByFilenameStmt.stmt,
	};
	for (auto& stmt : stmtList) {
		sqlite3_bind_text(stmt, 1, filename.data(), filename.size(), nullptr);
		int result = sqlite3_step(stmt);
		assert(result == SQLITE_DONE);
		sqlite3_reset(stmt);
		sqlite3_clear_bindings(stmt);
	}
	m->CommitTransaction();
}

void CodegenModelArchive::Store(const CodegenModel& cgInput) {
	auto& cgm = cgInput.GetPimpl();

	struct Visiter {
		CodegenModelArchive* self;

		void operator()(const DeclStruct& decl) const {
			self->StoreStruct(decl);
		}
		void operator()(const DeclFunction& decl) const {
			self->StoreFunction(decl);
		}
		void operator()(const DeclEnum& decl) const {
			self->StoreEnum(decl);
		}
	} visiter;
	visiter.self = this;

	m->BeginTransaction();

	for (auto&& [DISCARD, ns] : cgm.namespaces) {
		// This will insert the namespace if it doesn't exist, or no-op (fetches data) if it already exists
		m->FindOrStoreNamespace(&ns);
	}
	for (auto&& [DISCARD, value] : cgm.decls) {
		std::visit(visiter, value.v);
	}

	m->CommitTransaction();
}

void CodegenModelArchive::StoreStruct(const DeclStruct& decl) {
	// TODO
}

void CodegenModelArchive::StoreFunction(const DeclFunction& decl) {
	// TODO
}

void CodegenModelArchive::StoreEnum(const DeclEnum& decl) {
	// -Argument-	-Description-
	// ?1			Namespace ID
	// ?2			Enum name
	// ?3			Enum underlying type
	// ?4			File containing the enum
	m->storeEnumStmt.InitializeLazily(m->database, R"""(
INSERT INTO DeclEnums(NamespaceId, Name, UnderlyingType, FileName)
	VALUES (?1, ?2, ?3, ?4)
	ON CONFLICT DO UPDATE SET
		UnderlyingType = ?3,
		FileName = ?4
	RETURNING Id;
)"""sv);

	// -Argument-	-Description-
	// ?1			Container enum's id
	// ?2			Enum element name
	// ?3			Enum element value
	m->storeEnumElmStmt.InitializeLazily(m->database, R"""(
INSERT INTO DeclEnumElements(EnumId, Name, Value)
	VALUES (?1, ?2, ?3)
	ON CONFLICT DO UPDATE SET Value=?3;
)"""sv);

	sqlite3_bind_int(m->storeEnumStmt, 1, m->FindNamespace(decl.container));
	sqlite3_bind_text(m->storeEnumStmt, 2, decl.name.c_str(), decl.name.size(), nullptr);
	sqlite3_bind_text(m->storeEnumStmt, 3, decl.underlyingTypeStr.c_str(), decl.underlyingTypeStr.size(), nullptr);
	if (decl.sourceFile) {
		sqlite3_bind_text(m->storeEnumStmt, 4, decl.sourceFile->filename.data(), decl.sourceFile->filename.size(), nullptr);
	} else {
		sqlite3_bind_text(m->storeEnumElmStmt, 4, "", 0, nullptr);
	}
	int result = sqlite3_step(m->storeEnumStmt);
	assert(result == SQLITE_ROW);
	auto EnumId = sqlite3_column_int64(m->storeEnumStmt, 0);
	sqlite3_reset(m->storeEnumStmt);
	sqlite3_clear_bindings(m->storeEnumStmt);

	for (auto& elm : decl.elements) {
		sqlite3_bind_int64(m->storeEnumElmStmt, 1, EnumId);
		sqlite3_bind_text(m->storeEnumElmStmt, 2, elm.name.c_str(), elm.name.size(), nullptr);
		sqlite3_bind_int64(m->storeEnumElmStmt, 3, elm.value);
		int result = sqlite3_step(m->storeEnumElmStmt);
		assert(result == SQLITE_DONE);
		sqlite3_reset(m->storeEnumElmStmt);
		sqlite3_clear_bindings(m->storeEnumElmStmt);
	}
}