#include <arpa/inet.h> // For inet_addr(), htons()
#include <stdio.h> // For printf()
#include <stdlib.h> // For exit()
#include <string.h> // For memset() 
#include <sys/socket.h> // For socket(), connect(), read(), write()
#include <time.h>
#include <unistd.h> // For close()
#include "sockets.h"

#define BUFFER_SIZE 1024

static void error_handling(char *message, int sock) {
    perror(message);
    if (sock != -1)
        close(sock);
}

int tcp_create_socket() {
    // PF_INET: IPv4, SOCK_STREAM: TCP
    int sock = socket(PF_INET, SOCK_STREAM, 0); 
    if (sock == -1)
        error_handling("socket() error", sock);
    return sock;
}

void connect_to_server(int sock, char *ip, int port) {
    struct sockaddr_in serv_addr;
    // Initialize serv_addr to zero
    memset(&serv_addr, 0, sizeof(serv_addr));
    // Set the address family to IPv4
    serv_addr.sin_family = AF_INET;
    // Convert IP address to binary form
    serv_addr.sin_addr.s_addr = inet_addr(ip);
    // Convert port number to network byte order
    serv_addr.sin_port = htons(port);

    if (connect(sock, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) == -1)
        error_handling("connect() error", sock);
}

int tcp_write(int sock, char* message) {
    // Sending the ASCII string
    if (write(sock, message, strlen(message)) == -1)
        error_handling("write() error", sock);
    return 0;
}

char *tcp_read(int sock) {
    char* response = malloc(BUFFER_SIZE);
    if (!response) 
        error_handling("Failed to allocate memory for response", sock);

    int str_len = read(sock, response, BUFFER_SIZE - 1);
    if (str_len == -1){
        free(response);
        error_handling("read() error", sock);
    } else if (str_len == 0) {
        printf("Server closed the connection.\n");
        free(response);
        return NULL;
    } else {
        response[str_len] = '\0'; // Null-terminate the received data
    }
    return response;
}

int tcp_listen(int port) {
    int sock = tcp_create_socket();

    struct sockaddr_in serv_addr;
    memset(&serv_addr, 0, sizeof(serv_addr));
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    serv_addr.sin_port = htons(port);

    if (bind(sock, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) == -1)
        error_handling("bind() error", sock);

    if (listen(sock, 5) == -1)
        error_handling("listen() error", sock);

    return sock;
}

int tcp_accept(int serv_sock) {
    struct sockaddr_in clnt_addr;
    socklen_t clnt_addr_size = sizeof(clnt_addr);
    int clnt_sock = accept(serv_sock, (struct sockaddr*)&clnt_addr, &clnt_addr_size);
    if (clnt_sock == -1)
        error_handling("accept() error", clnt_sock);
    return clnt_sock;
}

void process_request(int clnt_sock) {
    char message[BUFFER_SIZE];
    int str_len = read(clnt_sock, message, BUFFER_SIZE - 1);
    if (str_len == -1)
        // TODO: clnt_sock might be a bug
        error_handling("read() error", clnt_sock); 

    message[str_len] = '\0'; // Null-terminate the string

    time_t t = time(NULL);
    struct tm *timeinfo;

    if (strcmp(message, "1") == 0) {
        timeinfo = localtime(&t); // Local time
    } else {
        timeinfo = gmtime(&t);    // GMT time
    }

    strftime(message, BUFFER_SIZE, "%Y-%m-%d %H:%M:%S", timeinfo);
    write(clnt_sock, message, strlen(message));
}

int udp_create_socket() {
    int sock = socket(PF_INET, SOCK_DGRAM, 0);
    if (sock == -1) {
        error_handling("UDP socket creation failed", sock);
    }
    return sock;
}

void udp_bind_socket(int sock, int port) {
    struct sockaddr_in serv_addr;
    memset(&serv_addr, 0, sizeof(serv_addr));
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    serv_addr.sin_port = htons(port);

    if (bind(sock, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) == -1) {
        error_handling("UDP bind failed", sock);
    }
}

void udp_send_message(int sock, char *server_ip, int port, char *message) {
    struct sockaddr_in dest_addr;
    memset(&dest_addr, 0, sizeof(dest_addr));
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_addr.s_addr = inet_addr(server_ip);
    dest_addr.sin_port = htons(port);

    sendto(sock, message, strlen(message), 0,
           (struct sockaddr *)&dest_addr, sizeof(dest_addr));
}

char *udp_receive_message(int sock) {
    struct sockaddr_in from_addr;
    socklen_t from_addr_size = sizeof(from_addr);
    char *buffer = malloc(BUFFER_SIZE);
    if (buffer == NULL) {
        error_handling("Failed to allocate memory for buffer", sock);
    }

    int str_len = recvfrom(sock, buffer, BUFFER_SIZE, 0,
                           (struct sockaddr *)&from_addr, &from_addr_size);
    if (str_len == -1) {
        free(buffer);
        error_handling("recvfrom() failed", sock);
    }

    buffer[str_len] = '\0';  // Null-terminate the string
    return buffer;
}

void handle_udp_client(int serv_sock) {
    struct sockaddr_in clnt_addr;
    socklen_t clnt_addr_size = sizeof(clnt_addr);
    char buffer[BUFFER_SIZE];

    while (1) {
        int str_len = recvfrom(serv_sock, buffer, BUFFER_SIZE - 1, 0,
                               (struct sockaddr*)&clnt_addr, &clnt_addr_size);
        if (str_len == -1)
            error_handling("recvfrom() failed", serv_sock);

        buffer[str_len] = '\0';  // Null-terminate the string

        time_t t = time(NULL);
        struct tm *timeinfo;

        if (strcmp(buffer, "1") == 0) {
            timeinfo = localtime(&t); // Local time
        } else {
            timeinfo = gmtime(&t);    // GMT time
        }

        strftime(buffer, BUFFER_SIZE, "%Y-%m-%d %H:%M:%S", timeinfo);

        // Prepare a response - could be customized based on the message
        sendto(serv_sock, buffer, strlen(buffer), 0,
               (struct sockaddr*)&clnt_addr, clnt_addr_size);
    }
}