#include <iostream>
#include <cstring>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <unistd.h>
#include <map>
#include <chrono>
#include <thread>
#include <csignal>
#include <fstream>  
#include <sstream>  
#include <sqlite3.h>


void clearClient(int client_socket){
  const char* clearCmd = "\033c";
  send(client_socket, clearCmd, strlen(clearCmd), 0);
}

void handleSIGPIPE(int) {
    std::cerr << "Received SIGPIPE signal" << std::endl;
}

void setSignalHandlers() {
    signal(SIGPIPE, handleSIGPIPE);
}

void sendTelnetOption(int client_socket, unsigned char option, unsigned char suboption) {
    unsigned char buffer[3];
    buffer[0] = 0xFF;  
    buffer[1] = option;
    buffer[2] = suboption;
    int bytes_sent = send(client_socket, buffer, 3, 0);
    if (bytes_sent == -1) {
        perror("Failed to send telnet option");
    }
}

void sendBusyBoxBanner(int client_socket) {
    sendTelnetOption(client_socket, 0xFD, 0x18);
    sendTelnetOption(client_socket, 0xFD, 0x1F);
    sendTelnetOption(client_socket, 0xFB, 0x03);
    sendTelnetOption(client_socket, 0xFE, 0x01);

    const char* banner = "BusyBox (built-in shell)\r\n";
    int bytes_sent = send(client_socket, banner, strlen(banner), 0);
    if (bytes_sent == -1) {
        perror("Failed to send banner");
    }
}

void sendFrames(int client_socket) {
    std::ifstream framesFile("frames.txt");
    std::string frame;
    std::string delimiter = "fStart";
    std::ostringstream currentFrame;
    while (std::getline(framesFile, frame)) {
        if (frame == delimiter) {
            std::string frameToSend = currentFrame.str();
            int bytes_sent = send(client_socket, frameToSend.c_str(), frameToSend.size(), 0);
            if (bytes_sent == -1) {
                perror("Failed to send frame");
                break;
            }

            currentFrame.str("");
            std::this_thread::sleep_for(std::chrono::milliseconds(42));
	    clearClient(client_socket);
        } else {
            currentFrame << frame << "\r\n";
        }
    }
}

void initializeDatabase() {
    sqlite3 *db;
    int rc = sqlite3_open("connections.db", &db);
    if (rc) {
        std::cerr << "Can't open database: " << sqlite3_errmsg(db) << std::endl;
        sqlite3_close(db);
        return;
    }

    const char* createTableSQL = "CREATE TABLE IF NOT EXISTS connections (ip TEXT PRIMARY KEY, count INT);";
    rc = sqlite3_exec(db, createTableSQL, 0, 0, 0);
    if (rc != SQLITE_OK) {
        std::cerr << "SQL error: " << sqlite3_errmsg(db) << std::endl;
        sqlite3_close(db);
        return;
    }

    sqlite3_close(db);
}


void updateConnectionCount(const std::string& ip) {
    sqlite3* db;
    int rc = sqlite3_open("connections.db", &db);
    if (rc) {
        std::cerr << "Can't open database: " << sqlite3_errmsg(db) << std::endl;
        sqlite3_close(db);
        return;
    }

    const char* insertSQL = "INSERT OR REPLACE INTO connections (ip, count) VALUES (?, COALESCE((SELECT count FROM connections WHERE ip = ?), 0) + 1);";
    sqlite3_stmt* stmt;
    rc = sqlite3_prepare_v2(db, insertSQL, -1, &stmt, 0);
    if (rc != SQLITE_OK) {
        std::cerr << "SQL error: " << sqlite3_errmsg(db) << std::endl;
        sqlite3_close(db);
        return;
    }

    rc = sqlite3_bind_text(stmt, 1, ip.c_str(), -1, SQLITE_STATIC);
    rc = sqlite3_bind_text(stmt, 2, ip.c_str(), -1, SQLITE_STATIC);
    rc = sqlite3_step(stmt);
    sqlite3_finalize(stmt);

    sqlite3_close(db);
}

void ftel() {
    int server_socket = socket(AF_INET, SOCK_STREAM, 0);
    initializeDatabase();
    if (server_socket == -1) {
        perror("Socket creation error");
        return;
    }

    sockaddr_in server_address;
    server_address.sin_family = AF_INET;
    server_address.sin_addr.s_addr = INADDR_ANY;
    server_address.sin_port = htons(23);

    if (bind(server_socket, (struct sockaddr*)&server_address, sizeof(server_address)) == -1) {
        perror("Binding error");
        return;
    }

    if (listen(server_socket, 1) == -1) {
        perror("Listening error");
        return;
    }

    std::cout << "Waiting for a connection..." << std::endl;

    while (true) {
        sockaddr_in client_address;
        socklen_t client_address_len = sizeof(client_address);
        int client_socket = accept(server_socket, (struct sockaddr*)&client_address, &client_address_len);

        if (client_socket == -1) {
            perror("Accepting error");
            continue;
        }

        std::string connectedClient = inet_ntoa(client_address.sin_addr);

        std::cout << "Connection established with: " << connectedClient << std::endl;

        sendBusyBoxBanner(client_socket);

        sendFrames(client_socket);

        close(client_socket);
        std::cout << "Message sent and connection closed." << std::endl;

	updateConnectionCount(connectedClient);
    }

    close(server_socket);
}

int main() {
    setSignalHandlers();
    ftel();
    return 0;
}