#include <net/if.h>
#include <stdio.h>
/* tap device */
#include <fcntl.h>
#include <sys/ioctl.h>
#include <linux/if_tun.h>
#include <unistd.h>
#include <sys/time.h>
#include <sys/poll.h>
#include <string.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <pthread.h>
#include <stdlib.h>
#include "femtotcp.h"
static int tap_fd;

//#define DHCP
#define FEMTOTCP_IP "10.10.10.2"
#define LINUX_IP "10.10.10.1"
#define TEST_SIZE (8 * 1024)

void print_buffer(uint8_t *buf, int len)
{
    int i;
    for (i = 0; i < len; i++) {
        if (i % 16 == 0)
            printf("\n");
        printf("%02x ", buf[i]);
    }
    printf("\n");
}

int tap_poll(struct ll *ll, void *buf, int len)
{
    struct pollfd pfd;
    (void)ll;
    int ret;
    pfd.fd = tap_fd;
    pfd.events = POLLIN;
    ret = poll(&pfd, 1, 2);
    if (ret < 0) {
        perror("poll");
        return -1;
    }
    if (ret == 0) {
        return 0;
    }
    return read(tap_fd, buf, len);
}

int tap_send(struct ll *ll, void *buf, int len)
{
    (void)ll;
    //print_buffer(buf, len);
    return write(tap_fd, buf, len);
}

int tap_init(struct ll *ll, const char *ifname)
{
    struct ifreq ifr;
    struct sockaddr_in *addr;
    int sock_fd;

    if ((tap_fd = open("/dev/net/tun", O_RDWR)) < 0) {
        perror("accessing /dev/net/tun");
        return -1;
    }
    memset(&ifr, 0, sizeof(ifr));
    ifr.ifr_flags = IFF_TAP | IFF_NO_PI;
    strncpy(ifr.ifr_name, ifname, IFNAMSIZ);
    if (ioctl(tap_fd, TUNSETIFF, (void *)&ifr) < 0) {
        perror("ioctl TUNSETIFF");
        close(tap_fd);
        return -1;
    }
    /* Get mac address */
    if (ioctl(tap_fd, SIOCGIFHWADDR, &ifr) < 0) {
        perror("ioctl SIOCGIFHWADDR");
        close(tap_fd);
        return -1;
    }
    strncpy(ll->ifname, ifname, sizeof(ll->ifname) - 1);
    memcpy(ll->mac, ifr.ifr_hwaddr.sa_data, 6);
    ll->mac[5] ^= 1;
    ll->poll = tap_poll;
    ll->send = tap_send;


    /* Set up network side */
    sock_fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock_fd < 0) {
        perror("socket");
        close(tap_fd);
        return -1;
    }
    memset(&ifr, 0, sizeof(ifr));
    ifr.ifr_flags = IFF_TAP | IFF_NO_PI;
    strncpy(ifr.ifr_name, ifname, IFNAMSIZ);
    if (ioctl(sock_fd, SIOCGIFFLAGS, &ifr) < 0) {
        perror("ioctl SIOCGIFFLAGS");
        close(sock_fd);
        return -1;
    }
    ifr.ifr_flags |= IFF_UP;
    if (ioctl(sock_fd, SIOCSIFFLAGS, &ifr) < 0) {
        perror("ioctl SIOCSIFFLAGS");
        close(sock_fd);
        return -1;
    }
    addr = (struct sockaddr_in *)&ifr.ifr_addr;
    addr->sin_family = AF_INET;
    inet_pton(AF_INET, LINUX_IP, &addr->sin_addr);
    if (ioctl(sock_fd, SIOCSIFADDR, &ifr) < 0) {
        perror("ioctl SIOCSIFADDR");
        close(sock_fd);
        return -1;
    }
    inet_pton(AF_INET, "255.255.255.0", &addr->sin_addr);
    if (ioctl(sock_fd, SIOCSIFNETMASK, &ifr) < 0) {
        perror("ioctl SIOCSIFNETMASK");
        close(sock_fd);
        return -1;
    }
    if (ioctl(sock_fd, SIOCGIFFLAGS, &ifr) < 0) {
        perror("ioctl SIOCGIFFLAGS");
        close(sock_fd);
        return -1;
    }
    ifr.ifr_flags = IFF_UP | IFF_RUNNING;
    if (ioctl(sock_fd, SIOCSIFFLAGS, &ifr) < 0) {
        perror("ioctl SIOCSIFFLAGS");
        close(sock_fd);
        return -1;
    }
    printf("Successfully initialized tap device %s\n", ifname);
    return 0;
}

#include <sys/random.h>
uint32_t ipstack_getrandom(void)
{
    uint32_t ret;
    getrandom(&ret, sizeof(ret), 0);
    return ret;
}
#define BUFFER_SIZE TEST_SIZE
static int test_echoserver(struct ipstack *s, int active_close)
{
    int fd, ret;
    int client_fd = -1;
    int tot_sent = 0;
    int exit_ok = 0, exit_count = 0;
    struct ipstack_sockaddr_in local_sock = {
        .sin_family = AF_INET,
        .sin_port = ee16(8), /* Echo */
        .sin_addr.s_addr = 0
    };
    fd = ft_socket(s, AF_INET, IPSTACK_SOCK_STREAM, 0);
    printf("socket: %04x\n", fd);
    ret = ft_bind(s, fd, (struct ipstack_sockaddr *)&local_sock, sizeof(local_sock));
    printf("bind: %d\n", ret);
    ret = ft_listen(s, fd, 1);
    printf("listen: %d\n", ret);

    while(1) {
        uint32_t ms_next;
        uint8_t buf[BUFFER_SIZE];
        struct timeval tv;
        gettimeofday(&tv, NULL);
        ms_next = ipstack_poll(s, tv.tv_sec * 1000 + tv.tv_usec / 1000);
        usleep(ms_next * 1000);
        if (exit_ok > 0) {
            if (exit_count++ < 10)
                continue;
            else break;
        }
        if (client_fd < 0) {
            client_fd = ft_accept(s, fd, NULL, NULL);
            if (client_fd > 0) {
                printf("accept: %04x\n", client_fd);
            }
        } else {
            if (ret <= 0) {
                ret = ft_recvfrom(s, client_fd, buf, sizeof(buf), 0, NULL, NULL);
                if (ret == -11)
                    continue; /* Call again */
                if (ret < 0) {
                    printf("Recv error: %d\n", ret);
                    ft_close(s, client_fd);
                    return -1;
                } else if (ret == 0) {
                    printf("Client side closed the connection.\n");
                    ft_close(s, client_fd);
                    client_fd = -1;
                    printf("Server: Exiting.\n");
                    exit_ok = 1;
                    continue;
                } else if (ret > 0) {
                    printf("recv: %d, echoing back\n", ret);
                }
            }
            if (ret > 0) {
                int snd_ret;
                snd_ret = ft_sendto(s, client_fd, buf + tot_sent, ret - tot_sent, 0, NULL, 0);
                if (snd_ret == -11)
                    continue; /* Call again */
                if (snd_ret < 0) {
                    printf("Send error: %d\n", snd_ret);
                    ft_close(s, client_fd);
                    return 0;
                }
    
                tot_sent += snd_ret;
                printf("sent %d bytes\n", snd_ret);
                if (ret == tot_sent) {
                    tot_sent = 0;
                    ret = 0;
                }
            }
        }
        if ((tot_sent >= 8192) && active_close) {
            ft_close(s, client_fd);
            client_fd = -1;
            printf("Server: Exiting.\n");
            exit_ok = 1;
            continue;
        }
    }
    return 0;
}

void *pt_echoclient_closing(void *arg)
{
    int fd, ret;
    unsigned total_r = 0;
    unsigned i;
    uint8_t buf[BUFFER_SIZE];
    uint8_t test_pattern[16] = "Test pattern - -";
    uint32_t *srv_addr = (uint32_t *)arg;
    struct sockaddr_in remote_sock = {
        .sin_family = AF_INET,
        .sin_port = ntohs(8), /* Echo */
    };
    remote_sock.sin_addr.s_addr = *srv_addr;
    fd = socket(AF_INET, IPSTACK_SOCK_STREAM, 0);
    if (fd < 0) {
        printf("test client socket: %d\n", fd);
        return (void *)-1;
    }
    sleep(1);
    printf("Connecting to echo server\n");
    ret = connect(fd, (struct sockaddr *)&remote_sock, sizeof(remote_sock));
    if (ret < 0) {
        printf("test client connect: %d\n", ret);
        perror("connect");
        return (void *)-1;
    }
    for (i = 0; i < sizeof(buf); i += sizeof(test_pattern)) {
        memcpy(buf + i, test_pattern, sizeof(test_pattern));
    }
    ret = write(fd, buf, sizeof(buf));
    if (ret < 0) {
        printf("test client write: %d\n", ret);
        return (void *)-1;
    }
    while (total_r < sizeof(buf)) {
        ret = read(fd, buf + total_r, sizeof(buf) - total_r);
        if (ret < 0) {
            printf("failed test client read: %d\n", ret);
            return (void *)-1;
        }
        if (ret == 0) {
            printf("test client read: server has closed the connection.\n");
            return (void *)-1;
        }
        total_r += ret;
    }
    for (i = 0; i < sizeof(buf); i += sizeof(test_pattern)) {
        if (memcmp(buf + i, test_pattern, sizeof(test_pattern))) {
            printf("test client: pattern mismatch\n");
            printf("at position %d\n", i);
            buf[i + 16] = 0;
            printf("%s\n", &buf[i]);
            return (void *)-1;
        }
    }
    close(fd);
    printf("Test client: success\n");
    return (void *)0;
}

int main(int argc, char **argv)
{
    struct ipstack *s;
    struct ll *tapdev;
    pthread_t pt;
    struct timeval tv;
    ip4 ip = 0, nm = 0, gw = 0;
    uint32_t srv_ip;

    int ret, test_ret = 0;
    (void)argc;
    (void)argv;
    (void)ip;
    (void)nm;
    (void)gw;
    (void)tv;
    (void)pt;
    (void)srv_ip;
    ipstack_init_static(&s);
    tapdev = ipstack_getdev(s);
    if (!tapdev)
        return 1;

    
    if (tap_init(tapdev, "femt0") < 0) {
        perror("tap init");
        return 2;
    }

    system("tcpdump -i femt0 -w test.pcap &");

#ifdef DHCP
    gettimeofday(&tv, NULL);
    ipstack_poll(s, tv.tv_sec * 1000 + tv.tv_usec / 1000);
    dhcp_client_init(s);
    do {
        gettimeofday(&tv, NULL);
        ipstack_poll(s, tv.tv_sec * 1000 + tv.tv_usec / 1000);
        usleep(1000);
        ipstack_ipconfig_get(s, &ip, &nm, &gw);
    } while (!dhcp_bound(s));
    printf("DHCP: obtained IP address.\n");
    ipstack_ipconfig_get(s, &ip, &nm, &gw);
    srv_ip = htonl(ip);
#else
    ipstack_ipconfig_set(s, atoip4(FEMTOTCP_IP), atoip4("255.255.255.0"),
            atoip4(LINUX_IP));
    printf("IP: manually configured\n");
    inet_pton(AF_INET, FEMTOTCP_IP, &srv_ip);
#endif


    pthread_create(&pt, NULL, pt_echoclient_closing, &srv_ip);
    printf("Starting test: echo server close-wait\n");
    ret = test_echoserver(s, 0);
    pthread_join(pt, (void **)&test_ret);
    printf("Test echo server close-wait: %d\n", ret);
    printf("Test linux client: %d\n", test_ret);
    sleep(1);
    system("killall tcpdump");
    return 0;
}