#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;
};

class CodegenModelArchive::Private {
	friend class CodegenModel;

public:
	sqlite3* database = nullptr;
	sqlite3_stmt* beginTransactionStmt = nullptr;
	sqlite3_stmt* commitTransactionStmt = nullptr;
	sqlite3_stmt* rollbackTransactionStmt = nullptr;
	sqlite3_stmt* storeEnumStmt = nullptr;
	sqlite3_stmt* storeEnumElmStmt = nullptr;

	// See below for definition
	void InitializeDatabase();

	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);
	}
};

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));
}

sqlite3_stmt* Utils::NewSqliteStatement(sqlite3* database, std::string_view sql) {
	sqlite3_stmt* 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 stmt;
}

void CodegenModelArchive::Private::InitializeDatabase() {
	int result = sqlite3_exec(database, "PRAGMA user_version = 1", nullptr, nullptr, nullptr);
	assert(result == SQLITE_OK);

	result = sqlite3_exec(database, R"""(
BEGIN TRANSACTION;
CREATE TABLE DeclEnums(
    Id INT PRIMARY KEY,
    Name TEXT,
    UnderlyingType TEXT
);
CREATE TABLE DeclEnumElements(
    EnumId INT,
    Name TEXT,
    Value INT,
    FOREIGN KEY (EnumId) REFERENCES DeclEnums(Id)
);
COMMIT TRANSACTION;
)""",
		nullptr,
		nullptr,
		nullptr);
	assert(result == SQLITE_OK);
}

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);
	}

	// Keep the same as `PRAGMA user_version` in the CodegenMetadataArchive::Private::InitializeDatabase()
	constexpr int kDatabaseVersion = 1;
	int currentDatabaseVersion;
	{
		auto readVersionStmt = Utils::NewSqliteStatement(m->database, "PRAGMA user_version");
		DEFER { sqlite3_finalize(readVersionStmt); };

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

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

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

	m->beginTransactionStmt = Utils::NewSqliteStatement(m->database, "BEGIN TRANSACTION");
	m->commitTransactionStmt = Utils::NewSqliteStatement(m->database, "COMMIT TRANSACTION");
	m->rollbackTransactionStmt = Utils::NewSqliteStatement(m->database, "ROLLBACK TRANSACTION");

	m->storeEnumStmt = Utils::NewSqliteStatement(m->database, "INSERT INTO DeclEnums(Name, UnderlyingType) VALUES(?, ?)"sv);
	m->storeEnumElmStmt = Utils::NewSqliteStatement(m->database, "INSERT INTO DeclEnumElements(EnumId, Name, Value) VALUES(?, ?, ?)"sv);
}

CodegenModelArchive::~CodegenModelArchive() {
	sqlite3_close(m->database);
	sqlite3_finalize(m->beginTransactionStmt);
	sqlite3_finalize(m->commitTransactionStmt);
	sqlite3_finalize(m->rollbackTransactionStmt);
	sqlite3_finalize(m->storeEnumStmt);
	sqlite3_finalize(m->storeEnumElmStmt);

	delete m;
}

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, 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) {
	// TOOD only update
	sqlite3_bind_text(m->storeEnumStmt, 1, decl.name.c_str(), decl.name.size(), nullptr);
	sqlite3_bind_text(m->storeEnumStmt, 2, decl.underlyingTypeStr.c_str(), decl.underlyingTypeStr.size(), nullptr);
	int result = sqlite3_step(m->storeEnumStmt);
	assert(result == SQLITE_DONE);
	sqlite3_reset(m->storeEnumStmt);
	sqlite3_clear_bindings(m->storeEnumStmt);

	auto enumRowId = sqlite3_last_insert_rowid(m->database);

	for (auto& elm : decl.elements) {
		sqlite3_bind_int64(m->storeEnumElmStmt, 1, enumRowId);
		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);
	}
}