#include <wolfssl/options.h>
#include <wolfssl/ssl.h>
#include <wolfssl/error-ssl.h>
#include <stdio.h>
#include <stdlib.h>

#include "transport.h"
#include "agent_cert.h"
#include "agent_key.h"
#include "ca_cert.h"

typedef struct {
    WOLFSSL_CTX* ctx;
    WOLFSSL* ssl;
    SOCKET sockfd;
} WolfSSLConnection;

int wolfssl_send(void* handle, char* data, size_t len) {
    if (!handle) {
        LOG_ERROR("Error: wolfssl_send called with null handle\n");
        return -1;
    }
    WolfSSLConnection* conn = (WolfSSLConnection*)handle;
    if (!conn->ssl) {
        LOG_ERROR("Error: SSL object is null in wolfssl_send\n");
        return -1;
    }

    int result = wolfSSL_write(conn->ssl, data, len);
    if (result <= 0) {
        int err = wolfSSL_get_error(conn->ssl, result);
        LOG_ERROR("Error: wolfSSL_write failed with result %d, error code %d: %s\n",
                result, err, wolfSSL_ERR_reason_error_string(err));
    }
    return result;
}

int wolfssl_recv(void* handle, char* buffer, size_t len) {
    if (!handle || !buffer) {
        LOG_ERROR("Error: wolfssl_recv called with null handle or buffer\n");
        return -1;
    }
    WolfSSLConnection* conn = (WolfSSLConnection*)handle;
    if (!conn->ssl) {
        LOG_ERROR("Error: SSL object is null in wolfssl_recv\n");
        return -1;
    }

    int result = wolfSSL_read(conn->ssl, buffer, len);
    if (result <= 0) {
        int err = wolfSSL_get_error(conn->ssl, result);
        LOG_ERROR("Error: wolfSSL_read failed with result %d, error code %d: %s\n",
                result, err, wolfSSL_ERR_reason_error_string(err));
    }
    return result;
}

void wolfssl_cleanup(void* handle) {
    if (!handle) {
        LOG_ERROR("Error: wolfssl_cleanup called with null handle\n");
        return;
    }
    WolfSSLConnection* conn = (WolfSSLConnection*)handle;
    if (conn->ssl) {
        wolfSSL_shutdown(conn->ssl);
        wolfSSL_free(conn->ssl);
        conn->ssl = NULL;
    }
    if (conn->ctx) {
        wolfSSL_CTX_free(conn->ctx);
        conn->ctx = NULL;
    }
    if (conn->sockfd != INVALID_SOCKET) {
        shutdown(conn->sockfd, SD_BOTH);
        closesocket(conn->sockfd);
        conn->sockfd = INVALID_SOCKET;
    }
    // WSACleanup();
    wolfSSL_Cleanup();
    free(conn);
}

Transport* InitWolfSSLTransport(char* domain, unsigned short port) {
    // Testing
    // WSADATA wsaData;
    // if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
    //     LOG_ERROR("Error: WSAStartup failed, error %d\n", WSAGetLastError());
    //     return NULL;
    // }

    WolfSSLConnection* conn = (WolfSSLConnection*)malloc(sizeof(WolfSSLConnection));
    if (!conn) {
        LOG_ERROR("Error: Failed to allocate memory for WolfSSLConnection\n");
        // WSACleanup();
        return NULL;
    }
    conn->ctx = NULL;
    conn->ssl = NULL;
    conn->sockfd = INVALID_SOCKET;

    conn->sockfd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (conn->sockfd == INVALID_SOCKET) {
        LOG_ERROR("Error: Failed to create socket, error %d\n", WSAGetLastError());
        // WSACleanup();
        free(conn);
        return NULL;
    }

    struct sockaddr_in serverAddr = {
        .sin_family = AF_INET,
        .sin_port = htons(port)
    };
    if (InetPtonA(AF_INET, domain, &serverAddr.sin_addr) <= 0) {
        LOG_ERROR("Error: Failed to resolve domain %s, error %d\n", domain, WSAGetLastError());
        closesocket(conn->sockfd);
        // WSACleanup();
        free(conn);
        return NULL;
    }

    if (connect(conn->sockfd, (struct sockaddr*)&serverAddr, sizeof(serverAddr)) < 0) {
        LOG_ERROR("Error: Failed to connect to %s:%d, error %d\n", domain, port, WSAGetLastError());
        closesocket(conn->sockfd);
        // WSACleanup();
        free(conn);
        return NULL;
    }

    wolfSSL_Init();
    conn->ctx = wolfSSL_CTX_new(wolfTLSv1_3_client_method());
    if (!conn->ctx) {
        LOG_ERROR("Error: Failed to create WolfSSL context\n");
        closesocket(conn->sockfd);
        // WSACleanup();
        free(conn);
        return NULL;
    }

    if (wolfSSL_CTX_use_certificate_buffer(conn->ctx, certificates_agent_crt, certificates_agent_crt_len, WOLFSSL_FILETYPE_PEM) != WOLFSSL_SUCCESS) {
        LOG_ERROR("Error: Failed to load agent certificate\n");
        wolfSSL_CTX_free(conn->ctx);
        closesocket(conn->sockfd);
        // WSACleanup();
        free(conn);
        return NULL;
    }
    if (wolfSSL_CTX_use_PrivateKey_buffer(conn->ctx, certificates_agent_key, certificates_agent_key_len, WOLFSSL_FILETYPE_PEM) != WOLFSSL_SUCCESS) {
        LOG_ERROR("Error: Failed to load agent private key\n");
        wolfSSL_CTX_free(conn->ctx);
        closesocket(conn->sockfd);
        // WSACleanup();
        free(conn);
        return NULL;
    }
    if (wolfSSL_CTX_load_verify_buffer(conn->ctx, certificates_ca_crt, certificates_ca_crt_len, WOLFSSL_FILETYPE_PEM) != WOLFSSL_SUCCESS) {
        LOG_ERROR("Error: Failed to load CA certificate\n");
        wolfSSL_CTX_free(conn->ctx);
        closesocket(conn->sockfd);
        // WSACleanup();
        free(conn);
        return NULL;
    }

    wolfSSL_CTX_set_verify(conn->ctx, WOLFSSL_VERIFY_PEER, NULL);
    conn->ssl = wolfSSL_new(conn->ctx);
    if (!conn->ssl) {
        LOG_ERROR("Error: Failed to create WolfSSL object\n");
        wolfSSL_CTX_free(conn->ctx);
        closesocket(conn->sockfd);
        // WSACleanup();
        free(conn);
        return NULL;
    }

    wolfSSL_set_fd(conn->ssl, conn->sockfd);
    if (wolfSSL_connect(conn->ssl) != WOLFSSL_SUCCESS) {
        int err = wolfSSL_get_error(conn->ssl, 0);
        LOG_ERROR("Error: TLS handshake failed, error code %d: %s\n",
                err, wolfSSL_ERR_reason_error_string(err));
        wolfSSL_free(conn->ssl);
        wolfSSL_CTX_free(conn->ctx);
        closesocket(conn->sockfd);
        // WSACleanup();
        free(conn);
        return NULL;
    }

    Transport* transport = (Transport*)malloc(sizeof(Transport));
    if (!transport) {
        LOG_ERROR("Error: Failed to allocate memory for Transport\n");
        wolfssl_cleanup(conn);
        return NULL;
    }

    transport->handle = conn;
    transport->send = wolfssl_send;
    transport->recv = wolfssl_recv;
    transport->cleanup = wolfssl_cleanup;
    LOG("WolfSSL transport initialized successfully\n");
    return transport;
}