#pragma once

#include <iostream>
#include <string>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <cstring>
#include <map>
#include <functional>

template < typename T_RpcTag >
struct RpcService {
private:
    uint32_t port;
    std::map< T_RpcTag, std::function<void( char const * input, size_t input_len, int client_sock )> > handler;

public:
    RpcService(uint32_t port) :port(port)
    {}

    template < typename F, typename Request, typename Response >
    void register_handler(T_RpcTag tag, std::function<F(Request const&, Response&)> f) {
        handler.emplace(
            tag,
            std::function([f]( char const * request_buffer, size_t len, int client_sock ) {
                Request * request = (Request*) (request_buffer + sizeof(T_RpcTag));
                Response response;
                f( *request, response );
                send(client_sock, (void const*)&response, sizeof(response), 0);
            })
        );
    }

    void handle_client(int client_sock) {
        char request_buffer[1024];
        ssize_t bytes_received;

        // Receive data from client
        bytes_received = recv(client_sock, request_buffer, sizeof(request_buffer), 0);
        if (bytes_received == -1) {
            std::cerr << "Failed to receive data" << std::endl;
            return;
        }

        printf("received %lu bytes:\n", bytes_received);
        for( size_t i = 0; i < bytes_received;)
        {
            printf("%4lx : ", i);
            for( size_t col = 0; col < 8; ++col ) {
                printf("%2x ", request_buffer[i++]);
            }
            printf("\n");
        }

        if (bytes_received >= sizeof(T_RpcTag) ) {
            T_RpcTag * tag = (T_RpcTag*) request_buffer;
            if( handler.contains(*tag) ){
                (handler[*tag])( request_buffer, bytes_received, client_sock );
            } else {
                std::string response("invalid RPC");
                send(client_sock, response.c_str(), response.length(), 0);
            }
        }
    }

    int run() {
        int server_sock, client_sock;
            struct sockaddr_in server_addr, client_addr;
            socklen_t addr_len = sizeof(client_addr);

            // Create socket
            server_sock = socket(AF_INET, SOCK_STREAM, 0);
            if (server_sock == -1) {
                std::cerr << "Failed to create socket" << std::endl;
                return 1;
            }

            // Set up server address
            server_addr.sin_family = AF_INET;
            server_addr.sin_addr.s_addr = INADDR_ANY;
            server_addr.sin_port = htons(port);

            // Bind socket to the address and port
            if (bind(server_sock, (struct sockaddr*)&server_addr, sizeof(server_addr)) == -1) {
                std::cerr << "Failed to bind socket" << std::endl;
                return 1;
            }

            // Listen for incoming connections
            if (listen(server_sock, 3) == -1) {
                std::cerr << "Failed to listen on socket" << std::endl;
                return 1;
            }
            std::cout << "Server listening on port " << port << "..." << std::endl;

            // Accept and handle client connections
            while (true) {
                client_sock = accept(server_sock, (struct sockaddr*)&client_addr, &addr_len);
                if (client_sock == -1) {
                    std::cerr << "Failed to accept connection" << std::endl;
                    continue;
                }

                std::cout << "Client connected!" << std::endl;
                handle_client(client_sock);

                // Close the client socket after handling the request
                close(client_sock);
            }

            // Close the server socket
            close(server_sock);
    }
};