#include "main.h"
#include "lwip/opt.h"
#include "lwip/arch.h"
#include "lwip/api.h"
#include "netif/ethernet.h"
#include "ethernetif.h"
#include "ethernetif.h"
#include "lwip/init.h"
#include "lwip/netif.h"
#include "lwip/tcp.h"
#include "lwip/dhcp.h"
#include "stdio.h"
#include "stdlib.h"
#include "string.h"
#include "cyclone_ssh.h"

#define SSH_PORT 22
#define MAX_CLIENTS 5

struct netif gnetif;

void Error_Handler(void);

void handle_ssh_connection(int newsock) {
    char buffer[1024];
    ssize_t len;

    CY_SECURE_SOCKET socket = CY_SECURE_SOCKET_INVALID;
    cy_rslt_t result = cy_secure_socket_create(CY_SECURE_SOCKET_TYPE_TCP, &socket);
    if (result != CY_RSLT_SUCCESS) {
        printf("Error creating secure socket\n");
        close(newsock);
        return;
    }

    result = cy_secure_socket_set_identity(socket, CY_SECURE_SOCKET_IDENTITY_TYPE_PSK, NULL, 0);
    if (result != CY_RSLT_SUCCESS) {
        printf("Error setting socket identity\n");
        close(newsock);
        cy_secure_socket_delete(socket);
        return;
    }

    result = cy_secure_socket_set_credentials(socket, CY_SECURE_SOCKET_ROOTCA_CERT, NULL, 0);
    if (result != CY_RSLT_SUCCESS) {
        printf("Error setting socket credentials\n");
        close(newsock);
        cy_secure_socket_delete(socket);
        return;
    }

    result = cy_secure_socket_connect(socket, "0.0.0.0", SSH_PORT);
    if (result != CY_RSLT_SUCCESS) {
        printf("Error connecting to server\n");
        close(newsock);
        cy_secure_socket_delete(socket);
        return;
    }

    // Authenticate the user here if needed

    while ((len = recv(newsock, buffer, sizeof(buffer), 0)) > 0) {
        // Process input and send output here
        // For simplicity, we'll just echo back the received data
        send(newsock, buffer, len, 0);
    }

    close(newsock);
    cy_secure_socket_delete(socket);
}

err_t accept_callback(void *arg, struct tcp_pcb *newpcb, err_t err) {
    static int client_count = 0;
    if (client_count >= MAX_CLIENTS) {
        printf("Maximum number of clients reached\n");
        tcp_abort(newpcb);
        return ERR_OK;
    }

    printf("Incoming connection\n");
    tcp_accepted((void *)newpcb);

    // Start SSH connection handling on a new thread or task
    handle_ssh_connection((int)newpcb->callback_arg);

    client_count++;
    return ERR_OK;
}

int main(void) {
    HAL_Init();
    SystemClock_Config();

    MX_GPIO_Init();
    MX_ETH_Init();
    MX_LWIP_Init();

    // Initialize Cyclone SSH
    int ssh_init_result = ssh_initialize();
    if (ssh_init_result != SSH_OK) {
        printf("Error initializing SSH: %s\n", ssh_get_error());
        Error_Handler();
    }

    struct tcp_pcb *pcb = tcp_new();
    if (pcb == NULL) {
        printf("Error creating PCB\n");
        Error_Handler();
    }

    if (tcp_bind(pcb, IP_ADDR_ANY, SSH_PORT) != ERR_OK) {
        printf("Error binding to port %d\n", SSH_PORT);
        Error_Handler();
    }

    pcb = tcp_listen(pcb);
    if (pcb == NULL) {
        printf("Error listening on port %d\n", SSH_PORT);
        Error_Handler();
    }

    tcp_accept(pcb, accept_callback);

    while (1) {
        MX_LWIP_Process();
        osDelay(1);
    }
}