#include "SQLiteX/SQLiteX.h"
#include "SQLiteX/Database.h"
#include "SQLiteX/Transaction.h"
#include "SQLiteX/Query.h"
#include "SQLiteX/Errors.h"
#include <stdexcept>
#include <iostream>

using namespace SQLiteX;
using namespace std;

// ############################################################################################# //
// # simpleTest()                                                                              # //
// ############################################################################################# //
/// A very simple demonstration
void simpleTest() {
  // Delete the old database if it exists
  ::unlink("test.sqlite3.db");

  // Open or create a database
  Database TestDB("test.sqlite3.db");

  // Create a new table in the database
  TestDB.execute(
    "CREATE TABLE Employees(\n"
    "  ID INTEGER PRIMARY KEY NOT NULL,\n"
    "  Name TEXT,\n"
    "  Salary INTEGER\n"
    ")"
  );
  
  // Destroy the table again
  TestDB.execute("DROP TABLE Employees");
  
  cout << "Test table created and destroyed again!" << endl
       << endl;
}

// ############################################################################################# //
// # queryTest()                                                                               # //
// ############################################################################################# //
/// Demonstration of a database query
void queryTest() {
  // Delete the old database if it exists
  ::unlink("test.sqlite3.db");

  // Open or create a database
  Database TestDB("test.sqlite3.db");

  // Create a table and put some entries in it
  TestDB.execute(
    "CREATE TABLE RomanNumbers(\n"
    "  Value INTEGER PRIMARY KEY,\n"
    "  Roman TEXT\n"
    ");\n"
    "INSERT INTO RomanNumbers(Value, Roman)\n"
    "  VALUES(1, 'I');\n"
    "INSERT INTO RomanNumbers(Value, Roman)\n"
    "  VALUES(2, 'II');\n"
    "INSERT INTO RomanNumbers(Value, Roman)\n"
    "  VALUES(3, 'III')"
  );

  // Create a query which lists all roman numbers and execute it
  Query ListNumbers("SELECT Value, Roman FROM RomanNumbers");
  TestDB.execute(ListNumbers);

  // Print out the query results
  cout << "Entries in Table 'RomanNumbers':" << endl;
  for(
    Query::ResultType::const_iterator ResultIt = ListNumbers.getResults().begin();
    ResultIt != ListNumbers.getResults().end();
    ++ResultIt
  )
    cout << static_cast<string>(ResultIt->at(0)) << "\t=\t"
         << static_cast<string>(ResultIt->at(1)) << endl;

  cout << endl;
}

// ############################################################################################# //
// # transactionTest()                                                                         # //
// ############################################################################################# //
/// Demonstrates the usage of transactions
void transactionTest() {

  // Delete the old database if it exists
  ::unlink("test.sqlite3.db");

  // Open or create a database
  Database TestDB("test.sqlite3.db");

  // Create some tables for testing
  TestDB.execute(
    "CREATE TABLE Test1(ID INTEGER PRIMARY KEY);\n"
    "CREATE TABLE Test2(ID INTEGER PRIMARY KEY);\n"
    "CREATE TABLE Test3(ID INTEGER PRIMARY KEY)"
  );

  // Create a transaction which will result in a failure
  Transaction DestroyTables;
  DestroyTables.add(
    "DROP TABLE Test1;\n"
    "DROP TABLE Test2;\n"
    "DROP TABLE MyFunnyTableName" // < Error after 2 tables are already dropped
  );

  // Execute the transaction and let it fail
  cout << "Executing erroneous transaction" << endl;
  try {
    TestDB.execute(DestroyTables);
  }
  catch(const Error &TheError) {
    cout << "Transaction failed on purpose: " << TheError.what() << endl;
  }

  // Because the transaction failed, all changes should have been rolled back
  cout << "Testing for successfull rollback operation" << endl;
  TestDB.execute(
    "DROP TABLE Test1;\n"
    "DROP TABLE Test2;\n"
    "DROP TABLE Test3"
  );

  // If we reach this place (no exception occured), the three tables
  // could be deleted, thus, they did still exist and the rollback worked
  cout << "Rollback was properly performed" << endl
       << endl;
}

// ############################################################################################# //
// # bindTest()                                                                                # //
// ############################################################################################# //
/// Demonstrates variable bindings to circumvent SQL injection attacks
void bindTest() {

  // Delete the old database if it exists
  ::unlink("test.sqlite3.db");

  // Open or create a database
  Database TestDB("test.sqlite3.db");

  // Create a test database which keeps user access rights
  TestDB.execute(
    "CREATE TABLE Users(\n"
    "  Name STRING PRIMARY KEY NOT NULL,\n"
    "  Password STRING,\n"
    "  AccessLevel STRING\n"
    ");\n"
    "INSERT INTO Users(Name, Password, AccessLevel)\n"
    "  VALUES('Fred', '123', 'full')"
  );

  cout << "Demonstrating SQL injection attack and prevention" << endl;

  // This is an SQL injection string to disable the password check
  std::string sTypedPassword = "' OR ''='";
  
  // Method 1: Validate the user login in a blatandly stupid
  //           manner open for SQL injection attacks
  { cout << "Doing unsafe login check..." << endl;
  
    Query CheckUser(
      std::string("SELECT AccessLevel FROM Users\n") +
        "WHERE Name='Fred' AND Password='" + sTypedPassword + "'"
    );
    TestDB.execute(CheckUser);

    if(CheckUser.getResults().size())
      cout << "  User logged in with AccessLevel: " 
           << static_cast<string>(CheckUser.getResults().front().at(0)) << endl;
    else
      cout << "  Unknown username or wrong password" << endl;
  }
  
  // Method 2: Use variable binding to circumvent any kind of SQL injection
  //           attack since no strings from the user appear in actual SQL
  { cout << "The same login check with variable binding..." << endl;
  
    Query CheckUser(
      "SELECT AccessLevel FROM Users\n"
      "WHERE Name='Fred' AND Password=:001\n"
    );
    CheckUser.bind(1, sTypedPassword);
    TestDB.execute(CheckUser);

    if(CheckUser.getResults().size())
      cout << "  User logged in with AccessLevel: " 
           << static_cast<string>(CheckUser.getResults().front().at(0)) << endl;
    else
      cout << "  Unknown username or wrong password" << endl;
  }
}

// ############################################################################################# //
// # main()                                                                                    # //
// ############################################################################################# //
/// Program entry point
int main() {

  // Run the simple test
  try {
    simpleTest();
  }
  catch(const Error &TheError) {
    cerr << "simpleTest() failed: " << TheError.what() << endl;
  }

  cout << string(79, '-') << endl;

  // Run the query test
  try {
    queryTest();
  }
  catch(const Error &TheError) {
    cerr << "queryTest() failed: " << TheError.what() << endl;
  }

  cout << string(79, '-') << endl;

  // Run the transaction test
  try {
    transactionTest();
  }
  catch(const Error &TheError) {
    cerr << "transactionTest() failed: " << TheError.what() << endl;
  }

  cout << string(79, '-') << endl;

  // Run the bindings test
  try {
    bindTest();
  }
  catch(const Error &TheError) {
    cerr << "bindTest() failed: " << TheError.what() << endl;
  }
  
  return 0;
}