#include <windows.h>
#include <stdlib.h>
#include <tchar.h>
#include "sqlite3.h"
#include "sha-2/sha-256.h"
#define SHASIZ SIZE_OF_SHA_256_HASH

/* TODO: "\\?\" trick?
 * Windows limits path lengths to 260 characters. Explorer on Windows XP
 * has this issue too.
 *
 * UTF issues?
 */

/*************************************************************************
 *  Compatability defines for non-Unicode systems
 ************************************************************************/

#ifdef UNICODE
# define ENCODING SQLITE3_UTF16
# define sqlite3_open_U sqlite3_open16
# define sqlite3_errmsg_U sqlite3_errmsg16
# define sqlite3_bind_text_U sqlite3_bind_text16
#else
# define ENCODING SQLITE3_UTF8
# define sqlite3_open_U sqlite3_open
# define sqlite3_errmsg_U sqlite3_errmsg
# define sqlite3_bind_text_U sqlite3_bind_text
#endif

/****************
 * Globals
 ***************/

sqlite3 *g_db = NULL;
int g_verbose = 0;
sqlite_int64 g_backup_id = -1;

/********************
 * Utility Functions
 *******************/

static void _log(TCHAR *msg, ...)
{
	va_list va;

	if (g_verbose) {
		va_start(va, emsg);
		_tvprintf(emsg, va);
		va_end(va);
	}
}
/* String literals only */
#define log(msg, ...) _log(TEXT(msg), __VA_ARGS__)

static void _die(TCHAR *emsg, ...)
{
	va_list va;
	va_start(va, emsg);
	_tvprintf(emsg, va);
	va_end(va);

	sqlite3_close(g_db);
	exit(1);
}
/* String literals only */
#define die(emsg, ...) _die(TEXT(emsg), __VA_ARGS__)

/**********************************
 * Database constants
 *********************************/

/* 'wlb' in ASCII */
#define UPPER_VERSION 0x776c62
#define CUR_VERSION 0x776c6201
#define S(s) #s
#define INIT_SCRIPT "\
BEGIN; \
PRAGMA user_version = " S(CUR_VERSION) ";\
CREATE TABLE chunks (\
	rowid INTEGER PRIMARY KEY, \
	sha256 TEXT UNIQUE NOT NULL, \
	data BLOB NOT NULL\
); \
CREATE TABLE backup_ids (\
	ts TEXT UNIQUE NOT NULL,\
	rowid INTEGER PRIMARY KEY\
); \
CREATE TABLE backups (\
	backup INTEGER NOT NULL REFERENCES backup_ids ON DELETE CASCADE, \
	path TEXT NOT NULL, \
	chunk INTEGER NOT NULL REFERENCES chunks ON DELETE RESTRICT\
); \
CREATE INDEX backups_paths ON backups (path); \
CREATE INDEX backups_backup ON backups (backup); \
CREATE INDEX backups_chunk ON backups (chunk);\
COMMIT;"

/* Check if the database version is correct. If it is not correct
 * and the create flag is enabled, then initialize a new database.
 */
static void check_db_version(int did_not_exist)
{
	int user_version;
	sqlite3_stmt *stmt;

	/* DB did not exist beforehand: initialize DB */
	if (did_not_exist) {
		if (sqlite3_exec(g_db, INIT_SCRIPT, NULL, NULL, NULL) != SQLITE_OK)
			die("Error initializing database: %s\n",
			    sqlite3_errmsg_U(g_db));
		log("Intialized database\n");
		return;
	}

	/* DB existed beforehand: do version check */
	if (sqlite3_prepare_v2(g_db, "PRAGMA user_version;", -1, &stmt, NULL)
	                       != SQLITE_OK)
		die("Error preparing user version check: %s\n",
		    sqlite3_errmsg_U(g_db));
	if (sqlite3_step(stmt) != SQLITE_ROW)
		die("Error executing user version check: %s\n",
		    sqlite3_errmsg_U(g_db));
	user_version = sqlite3_column_int(stmt, 0);
	sqlite3_finalize(stmt);

	if (user_version != CUR_VERSION) {
		die("Bad DB version found (expected %d), got %d\n",
		    CUR_VERSION & 0xFF, user_version & 0xFF);
	}
}

/* Initialize the backup ID used in this session.
 * This is only called if the user requests an archive of a directory.
 */
static void initialize_backup_id(void)
{
	sqlite3_stmt *stmt;

	if (sqlite3_prepare(g_db,
	                    "INSERT INTO backup_ids (ts) VALUES (datetime()) RETURNING rowid;",
	                    -1, &stmt, NULL) != SQLITE_OK) {
		die("failed to prepare inserting timestamp: %s\n",
		    sqlite3_errmsg_U(g_db));
	}

	if (sqlite3_step(stmt) != SQLITE_ROW)
		die("failed to insert timestamp: %s\n",
		    sqlite3_errmsg_U(g_db));
	g_backup_id = sqlite3_column_int64(stmt, 0);
	sqlite3_finalize(stmt);
}

/* Opens the DB and initializes it if it did not exist. */
static void open_db(TCHAR *fn)
{
	int did_not_exist = GetFileAttributes(fn)
	                    == INVALID_FILE_ATTRIBUTES;
	char *errmsg;

	if (sqlite3_open_U(fn, &g_db) != SQLITE_OK)
		die("Error opening database %s: %s\n", argv[i],
		    sqlite3_errmsg_U(g_db));
	check_db_version(did_not_exist);

	/* Store current timestamp */
	if (sqlite3_exec(g_db, "BEGIN;" NULL, NULL, &errmsg) != SQLITE_OK)
		die("failed to begin transaction: %s\n", errmsg);
}

/* Insert a backup record into the database.
 *
 * This part occurs after the file data has been inserted or identified
 * by SHA256 hash.
 */
static int insert_backup_record(TCHAR *name, sqlite_int64 chunk_id)
{
	sqlite3_stmt *stmt;
	int r = 0;

	if (sqlite3_prepare(g_db,
	                    "INSERT INTO backups (backup, path, chunk) VALUES (?,?,?);",
	                    -1, &stmt, NULL) != SQLITE_OK) {
		_tprintf("Could not prepare backup insert statement for %s: %s\n",
		         name, sqlite3_errmsg_U(g_db));
		return 0;
	}

	if (sqlite3_bind_int64(stmt, 1, g_backup_id) != SQLITE_OK) {
		_tprintf("Could not bind backup id for %s: %s\n", name,
		         sqlite3_errmsg_U(g_db));
		goto end;
	}

	if (sqlite3_bind_text_U(stmt, 2, name) != SQLITE_OK) {
		_tprintf("Could not bind path for %s: %s\n", name,
		         sqlite3_errmsg_U(g_db));
		goto end;
	}

	if (sqlite3_bind_int64(stmt, 3, chunk_id) != SQLITE_OK) {
		_tprintf("Could not bind chunk for %s: %s\n", name,
		         sqlite3_errmsg_U(g_db));
		goto end;
	}

	if (sqlite3_step(stmt) != SQLITE_DONE) {
		_tprintf("Error in inserting chunk for %s: %s\n", name,
		         sqlite3_errmsg_U(g_db));
		goto end;
	}

	r = 1;
end:
	sqlite3_finalize(stmt);
	return r;
}

/* Write a file to a BLOB in ``chunks`` at row ``rowid``.
 *
 * Returns 0 on failure and 1 on success.
 */
static int write_chunk(TCHAR *name, HANDLE f, sqlite_int64 rowid)
{
	sqlite3_blob *blob;
	char buf[4096];
	DWORD read = 0;
	int has_written = 0;
	int r = 0;

	if (sqlite3_blob_open(g_db, "main", "chunks", "data", rowid, 1,
	                      &blob) != SQLITE_OK) {
		_tprintf("Could not open BLOB to write chunk for %s: %s\n",
		         name, sqlite3_errmsg_U(g_db));
		return 0;
	}

	for (;;) {
		if (ReadFile(f, buf, sizeof(buf), &read, NULL) == FALSE) {
			_tprintf("failed to read in %s for sha256 calculation\n",
			         name);
			goto end;
		}

		if (read == 0)
			break;
		if (sqlite3_blob_write(g_db, buf, read, has_written)
		    != SQLITE_OK) {
			_tprintf("Could not write to BLOB for %s: %s\n",
			         name, sqlite3_errmsg_U(g_db));
			goto end;
		}

		has_written += read;
	}

	r = 1;
end:
	sqlite3_blob_close(blob);
	return r;
}

/* Insert a chunk with a specified SHA256 checksum into the database.
 *
 * Returns 1 on success, 0 on failure. On success, ``rowid`` contains
 * the row in ``chunks`` that contains the data of ``f``.
 */
static int insert_chunk(TCHAR *name, HANDLE f, sqlite_int64 *rowid,
                        uint8_t sha256[SHASIZ])
{
	LARGE_INTEGER fsize = 0;
	sqlite3_stmt *stmt;
	int r = 0;

	if (SetFilePointer(f, 0, NULL, FILE_BEGIN)
	    == INVALID_SET_FILE_POINTER) {
		_tprintf("rewind for %s failed\n", name);
		return 0;
	}

	if (GetFileSizeEx(f, &fsize) == 0) {
		_tprintf("Could not get the file size of %s\n", name);
		return 0;
	}

	/* Chunks are fixed size in SQLite. They need to be pre-allocated
	 * with the size of a file before a file can be written to it.
	 */
	if (sqlite3_prepare(g_db,
	                    "INSERT INTO chunks (sha256, data) VALUES (hex(?), zeroblob(?)) RETURNING rowid;",
	                    -1, stmt, NULL) != SQLITE_OK) {
		_tprintf("Could not prepare chunk insertion for %s: %s\n",
		         name, sqlite3_errmsg_U(g_db));
		return 0;
	}

	if (sqlite3_bind_blob(stmt, 1, sha256, SHASIZ, SQLITE_TRANSIENT)
	    != SQLITE_OK) {
		_tprintf("Could not bind sha256 to statement for %s: %s\n",
		         name, sqlite3_errmsg_U(g_db));
		goto finalize;
	}

	if (sqlite3_bind_int64(stmt, 2, fsize) != SQLITE_OK) {
		_tprintf("Could not bind file size to statement for %s: %s\n",
		         name, sqlite3_errmsg_U(g_db));
		goto finalize;
	}

	if (sqlite3_step(&stmt) != SQLITE_ROW)
		_tprintf("Could not step chunk insertion statement for %s: %s\n",
		         name, sqlite3_errmsg_U(g_db));
		goto finalize;
	}

	*rowid = sqlite3_column_int64(stmt, 0);
	r = write_chunk(name, f, *rowid);
finalize:
	sqlite3_finalize(stmt);
	return r;
}

/* Check if the SHA256 checksum already exists in the database.
 *
 * If the checksum exists, ``rowid`` is filled with the row in ``chunks``
 * that contains that sha256 sum. If it does not exist, then ``rowid``
 * contains ``-1``.
 *
 * Returns 0 if an error occured, and 1 if no error occured.
 */
static int check_sha256(TCHAR *name, uint8_t sha256[SHASIZ],
                        sqlite_int64 *rowid)
{
	sqlite3_stmt *stmt;
	int r = 0;

	/* Since ``sha256`` is binary, use ``hex`` to convert the binary
	 * string to a text representation.
	 */
	if (sqlite3_prepare_v2(g_db,
	                       "SELECT rowid FROM chunks WHERE sha256 = hex(?);",
	                       -1, &stmt, NULL) != SQLITE_OK) {
		_tprintf("failed to prepare sha256 statement for %s: %s\n",
		         name, sqlite_errmsg_U(g_db));
		goto end;
	}

	if (sqlite3_bind_blob(&stmt, 1, sha256, SHASIZ, SQLITE_TRANSIENT)
	    != SQLITE_OK) {
		_tprintf("failed to bind sha256 value for %s: %s\n",
		         name, sqlite_errmsg_U(g_db));
		goto end;
	}

	switch (sqlite3_step(&stmt)) {
	case SQLITE_ROW:
		/* there is a sha256 value */
		*rowid = sqlite3_column_int(stmt, 0);
		break;
	case SQLITE_DONE:
		/* The chunk was never entered */
		*rowid = -1;
		break;
	}

	r = 1;
end:
	sqlite3_finalize(stmt);
	return r;
}

/* Calculate the SHA256 checksum of the file in ``f``, and place a binary
 * representation of the checksum into ``shastr``.
 *
 * Returns 0 if there was an error, 1 for success.
 */
static int calculate_sha256(TCHAR *name, HANDLE f, uint8_t sha256[SHASIZ])
{
	DWORD read = 0;
	uint8_t buf[SIZE_OF_SHA_256_CHUNK * 64];
	struct Sha_256 sha_state;
	int i;

	sha_256_init(&sha_state, sha256);
	for (;;) {
		if (ReadFile(f, buf, sizeof(buf), &read, NULL) == FALSE) {
			_tprintf("failed to read in %s for sha256 calculation\n",
			         name);
			return 0;
		}

		/* ReadFile() reads 0 bytes on EOF. */
		if (read == 0)
			break;
		sha_256_write(&sha_state, buf, read);
	}

	sha_256_close(&sha_state);
	return 1;
}

/* Archive a file. */
static void archive_file(TCHAR *name)
{
	HANDLE f;
	sqlite_int64 chunks_rowid = -1;
	uint8_t sha256[SHASIZ];
	char *errstr;
	int success = 0;

	if (sqlite3_exec(g_db, "SAVEPOINT file;", NULL, NULL, &errstr)
	    != SQLITE_OK)
		die("failed to initialize file savepoint: %s\n", errstr);

	f = CreateFile(name, GENERIC_READ, FILE_SHARE_READ, NULL,
	                     OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
	if (f == INVALID_HANDLE_VALUE) {
		_tprintf("could not open %s\n", name);
		goto end;
	}

	if (!calculate_sha256(name, f, sha256))
		goto end;
	if (!check_sha256(sha256, &chunks_rowid))
		goto end;

	if (chunks_rowid < 0) {
		if (!insert_chunk(name, f, &chunks_rowid, sha256))
			goto end;
	}

	if (!insert_backup_record(name, chunks_rowid))
		goto end;
	success = 1;

	log("Archived %s\n", name);
end:
	if (sqlite3_exec(g_db, success ? "RELEASE file;" : "ROLLBACK TO file;", NULL, NULL,
	                 &errstr) != SQLITE_OK)
		die("failed to rollback file savepoint: %s\n", errstr);
	CloseHandle(f);
}

/* Recursively archive a directory. */
static void archive_directory(TCHAR *dirname)
{
	WIN32_FIND_DATA fdata;
	HANDLE dhandle;
	TCHAR pathname[PATH_MAX];
	dhandle = FindFirstFile(dirname, &fdata);
	if (dhandle == INVALID_HANDLE_VALUE) {
		_tprintf("Failed to open directory %s\n", dirname);
		return;
	}

	do {
		/* Windows versions prior to 11 (!) include _snprintf() as
		 * a non-standard version of snprintf(). _snprintf() will
		 * return -1 if the string does not fit in the buffer.
		 */
		if (_sntprintf(pathname, PATH_MAX, "%s\\%s", dirname,
		               fdata.cFileName) < 0) {
			_tprintf("Pathname for %s\\%s too long\n", dirname,
			         fdata.cFileName);
			continue;
		}

		if (fdata.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
			archive_directory(pathname);
		else
			archive_file(pathname);
	} while (FindNextFile(dhandle, &fdata) != 0);

	FindClose(dhandle);
}

/* Check if the path is a file or a directory, and either archive
 * the file or archive the entire directory recursively. */
static void archive_file_or_directory(TCHAR *name)
{
	DWORD attr = GetFileAttributes(fn);
	if (attr == INVALID_FILE_ATTRIBUTES) {
		_tprintf(TEXT("failed to open %s\n"), name);
		return;
	}

	if (g_backup_id < 0)
		initialize_backup_id();

	if (attr & FILE_ATTRIBUTE_DIRECTORY)
		archive_directory(name);
	else
		archive_file(name);
}

/* Print instructions to console and exit. */
static void usage(void)
{
	_tprintf(TEXT("\nwlb [\\V] [\\H] [\\D DBNAME] [\\A DIRECTORIES...]\n"));
	_tprintf(TEXT("\\V: Verbose\n"));
	_tprintf(TEXT("\\H: Display the help\n"));
	_tprintf(TEXT("\\D: Database (create if does not exist)\n"));
	_tprintf(TEXT("\\A: Archive directories\n"));
	exit(0);
}

int _tmain(int argc, TCHAR *argv[])
{
	int i;

	for (i = 1; i < argc; i++) {
		if (_tcscmp(argv[i], TEXT("\\H")) == 0) {
			usage();
		} else if (_tcscmp(argv[i], TEXT("\\A")) == 0) {
			i++;
			archive_file_or_directory(argv[i]);
		} else if (_tcscmp(argv[i], TEXT("\\D")) == 0) {
			i++;
			open_db(argv[i]);
		} else if (_tcscmp(argv[i], TEXT("\\V")) == 0) {
			g_verbose = 1;
		} else {
			usage();
		}
	}

	if (g_db)
		sqlite3_exec(g_db, "COMMIT;", NULL, NULL, NULL);
	sqlite3_close(g_db);
	return 0;
}