// for u Nightor01

#define uchar unsigned char
#define uint unsigned int
#define ushort unsigned short
#define ulong unsigned long
enum calls {
    SYS_READ,
    SYS_WRITE,
    SYS_OPEN,
    SYS_CLOSE,
    SYS_SOCKET = 41,
    SYS_CONNECT = 42,
    SYS_ACCEPT = 43,
    SYS_SENDTO = 44,
    SYS_RECVFROM = 45,
    SYS_SHUTDOWN = 48,
    SYS_BIND = 49,
    SYS_LISTEN = 50,
    SYS_FORK = 57,
    SYS_EXIT = 60,
};

#define AF_INET 2
#define SOCK_STREAM 1

#define STDIN 0
#define STDOUT 1
#define STDERR 2

#define SHUT_RDWR 2

struct addr {
    unsigned short family;
    unsigned short port;
    unsigned int ip;
    char pad[8];
} typedef addr;

#define itoa(x) ({char __[4]={0};short _=x;for(int i=3;i>=0;i--){__[i]=_%10+'0';_/=10;}__;})
//#define htonl(_) ({uint a = *(uint*)(uchar[]){((uchar*)&_)[3],((uchar*)&_)[2],((uchar*)&_)[1],((uchar*)&_)[0]}; a;})
//ushort htons(ushort _) {ushort a = *(ushort*)(uchar[]){((uchar*)&_)[1],((uchar*)&_)[0]}; return a;}
#define strlen(str) ({int l = 0; while(str[l] != '\0') {l++;} l;})
#define print(s) {syscall(SYS_WRITE, STDERR, (long)s, strlen(s), 0, 0, 0);}
#define exit(e) {syscall(SYS_EXIT, e, 0, 0, 0, 0, 0);}
#define assert(expr) ({if (!(expr)) {print("Assertion failed! "); print(__FILE__); print(":"); print(itoa(__LINE__)); print(" -> "); print(#expr); exit(1);}})

#define FRAGMENT 1024


void server(int socket_fd, char*arg);
void client(int socket_fd, char*arg);
long syscall(int call_eax, long rdi, long rsi, long rdx, long r10, long r8, long r9);

struct {
    ulong rip;
    ulong rbp;
} typedef stack_frame;

__attribute__((force_align_arg_pointer))
void _start() {
    register void * stack_ptr asm("rsp");
    int socket_fd;
    int argc;
    char** argv;
    // STACK GROWS DOWNWARDS!!!

    argc = *(int*)(stack_ptr+sizeof(stack_frame)+sizeof(void*)+sizeof(char**)+sizeof(int)+sizeof(int));
                    // gcc compiles so that every variable is defined first
    argv = (char**)(&argc+11);

    if (argc < 3) {
        print("Usage:\n\tServer: ./hard s file\n\tClient: ./hard c [IP]\n");
        exit(1);
    }
    socket_fd = syscall(SYS_SOCKET, AF_INET, SOCK_STREAM, 0, 0, 0,0);
    assert(socket_fd != -1);

    if (argv[1][0] == 's') {
        server(socket_fd, argv[2]);
    } else {
        client(socket_fd, argv[2]);
    }
    exit(0);
    __builtin_unreachable();
}

long syscall(int call_eax, long rdi, long rsi, long rdx, long r10, long r8, long r9) {
    unsigned long out;
    register int _r10 asm("r10") = r10;
    register int _r8 asm("r8") = r8;
    register int _r9 asm("r9") = r9;
    asm volatile (
        "syscall"
        : "=a" (out)
        : "0"(call_eax), "D"(rdi), "S"(rsi), "d"(rdx), "r"(_r10), "r"(_r8), "r"(_r9)
        : "rcx", "r11", "memory"
    );
    return out;
}

void server(int socket_fd, char* file) {
    addr bind_addr = {
        .family = AF_INET,
        //.port = htons(1025),
        .port = 206, // 1025 in BE
        .ip = 0,
        .pad = {0, 0, 0, 0, 0, 0, 0, 0}
    };
    addr client;
    int socket2, fd, client_len;
    assert(syscall(SYS_BIND, socket_fd, (long)&bind_addr, sizeof(addr), 0, 0, 0) == 0);
    assert(syscall(SYS_LISTEN, socket_fd, 16, 0, 0, 0, 0) == 0);
    while (1) {
        assert((socket2 = syscall(SYS_ACCEPT, socket_fd, (long)&client, (long)&client_len, 0, 0, 0)) != -1);
        switch (syscall(SYS_FORK, 0, 0, 0, 0, 0, 0)) {
            case 0:
                assert((fd = syscall(SYS_OPEN, (long)file, 0, 0, 0, 0, 0 )) != -1 );
                int moved = -1;
                char buf[FRAGMENT] = {0};
                while (1) {
                    assert((moved = syscall(SYS_READ, fd, (long)buf, FRAGMENT, 0, 0, 0)) != -1);
                    if (moved == 0) break;
                    assert(syscall(SYS_SENDTO, socket2, (long)buf, moved, 0, 0, 0) != -1);
                }
                syscall(SYS_SHUTDOWN, socket2, SHUT_RDWR, 0, 0, 0, 0);
                syscall(SYS_CLOSE, socket2, 0, 0, 0, 0, 0);
                syscall(SYS_CLOSE, fd, 0, 0, 0, 0, 0);
                print("Closed connection\n");
                exit(0);
            default:
                continue;
        }
    }
}

void client(int socket_fd, char* ip) {
    assert(strlen(ip) < 16);
    unsigned char ip_octets[4] = {0};
    int counter = 0;
    for (int i = 0; i<strlen(ip); i++) {
        if (ip[i] == '.') {counter++;continue;}
        ip_octets[counter] *= 10;
        ip_octets[counter] += (ip[i]-0x30);
    }
    addr server = {
        .family = AF_INET,
        //.port = htons(1025),
        //.ip = htonl(*(unsigned int*)ip_octets),
        .port = 206, // 1025 in BE
        .ip = *(unsigned int*)ip_octets, // the way the for loop above is written means it's already big endian
        .pad = {0, 0, 0, 0, 0, 0, 0, 0}
    };

    assert(syscall(SYS_CONNECT, socket_fd, (long)&server, sizeof(addr), 0, 0, 0) != -1);
    char buf[FRAGMENT];
    long read_from = -1;
    while(1) {
        assert((read_from = syscall(SYS_RECVFROM, socket_fd, (long)buf, FRAGMENT, 0, 0, 0)) > -1);
        if (read_from == 0) break;
        syscall(SYS_WRITE, STDOUT, (long)buf, read_from, 0,0,0);
    }
    syscall(SYS_CLOSE, socket_fd, 0, 0, 0, 0, 0);
    exit(0);
}