#include <cstdio>
#include <fcntl.h>
#include <netdb.h>
#include <poll.h>
#include <string>
#include <sys/socket.h>
#include <unistd.h>

#include <openssl/err.h>
#include <openssl/ssl.h>

void die(const char *format, ...) {
    va_list args;
    va_start(args, format);
    vfprintf(stderr, format, args);
    va_end(args);
    exit(1);
}

int main(int argc, char **argv) {
    if (argc != 3)
        die("Usage: %s <ip> <port>\n", argv[0]);

    // tworzy gniazdo i nawiązuje połączenie TCP
    addrinfo *res;
    if (getaddrinfo(argv[1], argv[2], 0, &res))
        die("getaddrinfo failed\n");
    int cli = socket(res->ai_family, SOCK_STREAM, 0);
    if (connect(cli, res->ai_addr, res->ai_addrlen) != 0)
        die("connect failed\n");
    freeaddrinfo(res);

    /* ───────────────────────────────────────────────────────────────────── */

    // wybiera konfigurację warstwy szyfrującej
    SSL_CTX *ctx = SSL_CTX_new(TLS_client_method()); // podstawowe ustawienia dla klienta TLS
    SSL_CTX_set_default_verify_paths(ctx); // nakazuje użyć systemowych certyfikatów CA

    // żeby przetestować weryfikację certyfikatu bez instalowania certyfikatu CA w systemie:
    // ustaw zmienną środowiskową MY_ROGUE_CA ma ścieżkę do własnego certyfikatu CA
    if (getenv("MY_ROGUE_CA"))
        SSL_CTX_load_verify_file(ctx, getenv("MY_ROGUE_CA"));

    /* ───────────────────────────────────────────────────────────────────── */

    // tworzy obiekt reprezentujący warstwę szyfrującą i nakazuje mu czytać/pisać do pliku gniazda
    SSL *encryption_layer = SSL_new(ctx);
    SSL_set_fd(encryption_layer, cli);

    // informuje bibliotekę jakiej nazwie powinien odpowiadać certyfikat
    SSL_set1_host(encryption_layer, argv[1]);

    // ustanawia komunikację szyfrowaną TLS (na wskazanym połączeniu TCP) wykonując handshake
    if (SSL_connect(encryption_layer) != 1)
        die("TLS handshake failed!\n");

    // odbiera wyniki weryfikacji poprawności otrzymanego od serwera certyfikatu
    long serverCertStatus = SSL_get_verify_result(encryption_layer);
    if (serverCertStatus != X509_V_OK)
        fprintf(stderr, "[Warning: server certificate verification failed - %s!]\n",
                X509_verify_cert_error_string(serverCertStatus));

    /* ───────────────────────────────────────────────────────────────────── */
    
    // fragment switcha sprawdzającego kody błędów, wspólna cześć dla read i write
    #define MY_ERROR_HANDLING_SWITCH_MISC                          \
    case SSL_ERROR_ZERO_RETURN:                                    \
        fprintf(stderr, "[Server terminated encryption layer]\n"); \
        goto close_connection;                                     \
    default:                                                       \
        fprintf(stderr, "[Some fatal error]\n");                   \
        ERR_print_errors_fp(stderr);                               \
        goto close_connection;

    pollfd pfds[2]{
        {.fd = 0,   .events = POLLIN},
        {.fd = cli, .events = POLLIN},
    };

    // zdarzenia które są potrzebne do odczytania i zapisu danych przez warstwę szyfrującą
    int readNeeds = POLLIN; // uwaga: TLS może potrzebować zrobić zapis do gniazda żeby odebrać dane
    int writeNeeds = 0; // uwaga: TLS może potrzebować zrobić odczyt do gniazda żeby wysłać dane

    // bufor na odbierane dane
    char recvBuffer[255];
    // bufory na wysłane dane
    std::string dataBeingWritten;
    std::string dataWaitingToBeWritten;
    // uwaga: jeśli nieblokujące SSL_write się nie powiodło bo musiałoby czekać,
    //        trzeba je powtórzyć z identycznymi argumentami; stąd dwa bufory

    // ustawia gniazdo w tryb nieblokujący (!)
    fcntl(cli, F_SETFL, fcntl(cli, F_GETFL) | O_NONBLOCK);
    // funkcje SSL_read i SSL_write mogą wielokrotnie wywoływać send i recv na
    // gnieździe, co przy blokującym trybie i użyciu zdarzeń mogłoby zablokować
    // program na obsłudze któregoś zdarzenia; stąd tryb nieblokujący

    while (1) {
        poll(pfds, 2, -1);

        if (pfds[0].revents) {
            char buf[256];
            int c = read(0, buf, 255);
            if (c > 0) {
                buf[c] = 0;
                dataWaitingToBeWritten += buf;
                writeNeeds |= POLLOUT;
                // dane będą wysłane dopiero w następnej iteracji pętli zdarzeń (to upraszcza kod)
            } else {
                // koniec pliku lub błąd na standardowym wejściu
                if (dataBeingWritten.empty() && dataWaitingToBeWritten.empty()) {
                    fcntl(cli, F_SETFL, fcntl(cli, F_GETFL) & ~O_NONBLOCK);
                    SSL_shutdown(encryption_layer);
                    // SSL_shutdown zamyka tylko wysyłanie danych; dane można
                    // odbierać aż druga strona nie zrobi SSL_shutdown
                    fcntl(cli, F_SETFL, fcntl(cli, F_GETFL) | O_NONBLOCK);
                    pfds[0].fd = -1;
                } else 
                    usleep(1000);
                // aktywne czekanie (powtarzana próba odczytu z zamkniętego standardowego
                // wejścia) na wepchnięcie do bufora nadawczego TCP zaszyfrowanych danych
                // przed włożeniem tam wiadomości zamknięciem warstwy szyfrującej jest tutaj
                // pozostawione dla uproszczenia kodu
            }
        }

        /* ───────────────────────────────────────────────────────────────────── */

        bool isThereAnyDataToWrite = (!dataBeingWritten.empty() || !dataWaitingToBeWritten.empty());
        if (pfds[1].revents & writeNeeds && isThereAnyDataToWrite) {
            writeNeeds = 0;

            if (dataBeingWritten.empty())
                std::swap(dataBeingWritten, dataWaitingToBeWritten);

            int result = SSL_write(encryption_layer, dataBeingWritten.c_str(), dataBeingWritten.length());

            // obsługa niepowodzenia funkcji SSL_write
            switch (SSL_get_error(encryption_layer, result)) {
                case SSL_ERROR_NONE:
                    break;
                case SSL_ERROR_WANT_READ:
                    writeNeeds = POLLIN;
                    break;
                case SSL_ERROR_WANT_WRITE:
                    writeNeeds = POLLOUT;
                    break;
                MY_ERROR_HANDLING_SWITCH_MISC
            }
            ERR_clear_error();

            // obsługa jeśli udało się wysłać dane
            if (result > 0) {
                // domyślnie SSL_write kończy się sukcesem dopiero jeśli wysłano
                // wszystkie dane, nawet jeśli gniazdo jest w trybie nieblokującym
                dataBeingWritten.clear();
                if (!dataWaitingToBeWritten.empty())
                    writeNeeds = POLLOUT;
            }
        }

        /* ───────────────────────────────────────────────────────────────────── */

        if (pfds[1].revents & readNeeds) {
            readNeeds = POLLIN;

            int result = SSL_read(encryption_layer, recvBuffer, 255);
            if (result > 0) {
                // odczyt się powiódł == SSL odczytał blok danych; rozmiar bloku
                // może być mniejszy bądź większy niż bufor w SSL_read.
                write(1, recvBuffer, result);
                // Jeżeli blok był większy niż bufor, to reszta danych czeka
                // w buforach warstwy szyfrującej - funkcja SSL_read może zwrócić
                // kolejną porcję danych bez wywoływania recv na gnieździe
                while (int remaining = SSL_pending(encryption_layer)) {
                    result = SSL_read(encryption_layer, recvBuffer, remaining < 255 ? remaining : 255);
                    write(1, recvBuffer, result);
                }
            } else {
                // obsługa niepowodzenia funkcji SSL_read
                switch (SSL_get_error(encryption_layer, result)) {
                case SSL_ERROR_NONE:
                case SSL_ERROR_WANT_READ:
                    break;
                case SSL_ERROR_WANT_WRITE:
                    readNeeds = POLLOUT;
                    break;
                    MY_ERROR_HANDLING_SWITCH_MISC
                }
                ERR_clear_error();
            }
        }

        /* ───────────────────────────────────────────────────────────────────── */

        pfds[1].events = readNeeds | writeNeeds;
    }

close_connection:
    shutdown(cli, SHUT_RDWR);
    close(cli);
    return 0;
}