#include <sqlite3.h>
#include <iostream>
#include <limits>
#include <sstream>

#define RESET       "\033[0m"
#define RED         "❌ \033[31m"
#define GREEN       "✅ \033[32m"
#define BOLDWHITE   "\033[1m\033[37m"

constexpr const char* DB_FILE = "users.db";
constexpr const char* MENU_MSG = "=== Database Management System ===\n1. Connect to database\n2. Execute custom SQL query\n3. Execute prepared SQL query\n4. Setup demo database\n5. Exit\n";
constexpr const char* MENU_PREP_MSG = "=== Prepared Statement Menu: Select statement type ===\n1. SELECT\n2. INSERT\n3. EXIT\n";
constexpr const char* MENU_PREP_COL_MSG = "=== Prepared Statement Menu: Select filter culumn ===\n1. Connect to database\n2. Execute custom SQL query\n3. Execute prepared SQL query\n4. Setup demo database\n5. Exit\n";
constexpr const char* MENU_PREP_TYPE_MSG = "=== Prepared Statement Menu: Select filter type ===\n1. Connect to database\n2. Execute custom SQL query\n3. Execute prepared SQL query\n4. Setup demo database\n5. Exit\n";

constexpr const char* PREP_STATEMENT = "INSERT INTO users (FirstName, LastName, Age) VALUES (?, ?, ?)";

constexpr const char* DEMO_CREATE_TABLE = "CREATE TABLE IF NOT EXISTS users (ID INTEGER PRIMARY KEY AUTOINCREMENT, FirstName varchar(255) NOT NULL, LastName varchar(255) NOT NULL, Age int);";
constexpr const char* DEMO_INSERT1 = "INSERT INTO users (FirstName, LastName, Age) VALUES ('Tedd', 'Chadwick', 23)";
constexpr const char* DEMO_INSERT2 = "INSERT INTO users (FirstName, LastName, Age) VALUES ('Kathe', 'Claris', 31)";
constexpr const char* DEMO_INSERT3 = "INSERT INTO users (FirstName, LastName, Age) VALUES ('Meryl', 'Chadwick', NULL)";
constexpr const char* DEMO_SELECT1 = "SELECT * FROM users;";
constexpr const char* DEMO_SELECT2 = "SELECT * FROM users WHERE LastName is 'Chadwick';";

class DatabaseException : public std::exception {
public:
    explicit DatabaseException() : errMsg(nullptr) {} 
    DatabaseException(const char *errMsg) : errMsg(errMsg) {}
    const char* what() const noexcept override {
            return errMsg ? errMsg : "Internal error!";
    }
private:
    const char *errMsg;
};
class InvalidArgument : public DatabaseException {
    const char* what() const noexcept override {
        return "Invalid argument!";
    }
};
class SQLiteException : public DatabaseException {
public:
    SQLiteException(char *errMsg) : DatabaseException(), errMsg(errMsg) {}
    const char* what() const noexcept override {
            return errMsg;
    }
    ~SQLiteException() {
        sqlite3_free((void*)(errMsg));
    }
private:
    char *errMsg;
};

enum class MenuState {
    MENU,
    CONNECT           = 1,
    QUERY_CUSTOM      = 2,
    QUERY_PREPARED    = 3,
    SETUP_DEMO        = 4,
    EXIT              = 5,
};

std::istream& operator>>(std::istream& is, MenuState& state) {
    int a;
    if (!(is >> a)) {
        is.clear();
        is.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
        throw InvalidArgument();
    }

    state = static_cast<MenuState>(a);

    return is;
}

bool dbConnect(sqlite3** db, std::ostream& os) {
    int rc;

    os << "Connecting to SQLite...";
    if (*db) {
        constexpr const char* errMsg = "Already connected!";
        throw DatabaseException(errMsg);
    }

    rc = sqlite3_open(DB_FILE, db);
    if (rc != SQLITE_OK) {
        constexpr const char* errMsg = "Connecting failed!";
        throw DatabaseException(errMsg);
    }

    os << GREEN << "Connected\n" << RESET;
    return true;
}

int printQuery(void *os_vp, int argc, char **argv, char **colName) {
    std::ostream& os = *static_cast<std::ostream*>(os_vp); // oh god...
    for (int i = 0; i < argc; i++) {
        os << colName[i] << " = " << (argv[i] ? argv[i] : "NULL") << '\n';
    }
    os << '\n';
    return 0;
}

bool sendQuery(sqlite3* db, std::ostream& os, const char* q, int (*callback)(void *, int, char **, char **)) {
    int rc;
    char* errMsg;
    void* os_p = &os; // oh god...

    os << "Sending query: '" << q << "'...";
    rc = sqlite3_exec(db, q, callback, os_p, &errMsg);
    if (rc != SQLITE_OK) {
        throw SQLiteException(errMsg);
    }
    os << GREEN << "Successful\n\n" << RESET;
    return true;
}

bool setupDemo(sqlite3** db, std::ostream& os) {
    dbConnect(db, os);

    os << "Creating table 'users':\n";
    sendQuery(*db, os, DEMO_CREATE_TABLE, nullptr);

    os << "Adding records:\n";
    sendQuery(*db, os, DEMO_INSERT1, nullptr);
    sendQuery(*db, os, DEMO_INSERT2, nullptr);
    sendQuery(*db, os, DEMO_INSERT3, nullptr);

    os << "Testing queries:\n";
    sendQuery(*db, os, DEMO_SELECT1, printQuery);
    sendQuery(*db, os, DEMO_SELECT2, printQuery);

    return true;
}

bool executePrepared(sqlite3* db, std::ostream& os, std::istream& is) {
    sqlite3_stmt *ppStmt;
    int rc;

    std::string fn;
    std::string ln;
    int age;

    os << "First name: ";
    is >> fn;
    os << "Last name: ";
    is >> ln;
    os << "Age (-1 for NULL): ";
    if (!(is >> age)) age = -1; 

    os << "Executing prepared statement: " << PREP_STATEMENT << "\nwith values: " 
        << fn << ", " << ln << ", " << (age >= 0 ? std::to_string(age) : std::string("NULL")) << "\n";
    rc = sqlite3_prepare(db, PREP_STATEMENT, 1024, &ppStmt, nullptr);
    if (rc != SQLITE_OK) throw DatabaseException("Statement cant be prepared!");

    sqlite3_bind_text(ppStmt, 1, fn.c_str(), -1, SQLITE_STATIC);
    sqlite3_bind_text(ppStmt, 2, ln.c_str(), -1, SQLITE_STATIC);
    if (age >= 0) sqlite3_bind_int(ppStmt, 3, age);

    rc = sqlite3_step(ppStmt);
    if (rc != SQLITE_DONE) {
        throw DatabaseException(sqlite3_errmsg(db));
    }

    rc = sqlite3_finalize(ppStmt);
    if (rc != SQLITE_OK) throw DatabaseException("Statement failed!");
    return rc;
}

int main(int ragc, char* argv[]) {
    std::ostream& os = std::cout;
    std::istream& is = std::cin;

    MenuState state = MenuState::MENU;

    sqlite3* db = nullptr;
    std::string userInput;

    for (;;) {
        userInput.clear();
        try {
            os << BOLDWHITE << MENU_MSG << RESET;
            is >> state;
            is.clear();
            is.ignore(std::numeric_limits<std::streamsize>::max(), '\n');

            switch (state) {
            case MenuState::CONNECT :
                dbConnect(&db, os);
                break;

            case MenuState::QUERY_CUSTOM :
                os << "> ";
                std::getline(is, userInput);
                if(is.bad()) throw DatabaseException("Input error!");
                sendQuery(db, os, userInput.c_str(), printQuery);
            break;

            case MenuState::QUERY_PREPARED :
                executePrepared(db, os, is);
                os << "Done\n";
            break;

            case MenuState::SETUP_DEMO :
                setupDemo(&db, os);
            break;

            case MenuState::EXIT :
                os << "Exiting...\n";
                sqlite3_close(db);
                return EXIT_SUCCESS;
            break;

            case MenuState::MENU :
            default:
                throw InvalidArgument();
            }
        } catch (DatabaseException& e) {
            os << RED << e.what() << "\n" << RESET;
            // return EXIT_FAILURE;
        }
        state = MenuState::MENU;
        os << "\n";
    }
    return EXIT_SUCCESS;
}