firedis.🔥

from sys.info import sizeof
from runtime.llcl import Runtime
from algorithm import parallelize, num_cores

from string_utils import to_repr
from protocol import FiredisParser
from table import Table

from libc import AF_INET, AF_INET6, SOCK_STREAM, SOL_SOCKET, SO_REUSEADDR, SHUT_RDWR
from libc import c_char, c_int, c_uint, c_void
from libc import (
    accept,
    bind,
    c_charptr_to_string,
    exit,
    htons,
    # inet_ntoa,
    # inet_ntop,
    inet_pton,
    listen,
    recv,
    send,
    setsockopt,
    shutdown,
    sockaddr,
    sockaddr_in,
    socket,
    socklen_t,
    to_char_ptr,
)

from config import load_config


fn reply_to_message(inout db: Table, msg: String) -> String:
    var parser = FiredisParser(db, msg)

    try:
        parser.parse()
    except e:
        print("> failed to parse message:", e.value)

    return parser.result


fn respond_to_client(inout db: Table, new_sockfd: c_int, workers: Int):
    with Runtime() as rt:

        @always_inline
        @parameter
        fn handle_messages(n: Int):
            while True:
                let buf_size = 1024
                let buf = Pointer[c_char]().alloc(buf_size)
                let bytes_recv = recv(new_sockfd, buf, buf_size, 0)

                print("> bytes received:", bytes_recv)

                if bytes_recv <= 0:
                    print("> failed to receive message, client disconnected?")
                    break

                let msg = c_charptr_to_string(buf, bytes_recv)

                print("> received message:", to_repr(msg))

                let reply = reply_to_message(db, msg)
                if send(new_sockfd, to_char_ptr(reply), len(reply), 0) == -1:
                    print("> failed to send reply:", to_repr(reply))
                    break

                print("> message sent successfully:", to_repr(reply))

        parallelize[handle_messages](rt, workers)


fn wait_for_clients(inout db: Table, sockfd: c_int, workers: Int):
    print("> waiting for connections...")

    while True:
        with Runtime() as rt:

            @always_inline
            @parameter
            fn handle_client(n: Int):
                let addr_ptr = Pointer[sockaddr].alloc(1)
                var sin_size = socklen_t(sizeof[socklen_t]())

                let new_sockfd = accept(
                    sockfd, addr_ptr, Pointer[socklen_t].address_of(sin_size)
                )

                if new_sockfd == -1:
                    print("> accept failed")
                    return
                else:
                    print("> new connection accepted with fd:", new_sockfd)

                respond_to_client(db, new_sockfd, workers)

            parallelize[handle_client](rt, workers)


fn main():
    let config = load_config()
    var db = Table.create()

    print("> starting firedis at ip addr:", config.host, "port:", config.port, "🔥")

    let address_family = AF_INET
    var ip_buf_size = 4
    if address_family == AF_INET6:
        ip_buf_size = 16

    let ip_buf = Pointer[c_void].alloc(ip_buf_size)
    let conv_status = inet_pton(address_family, to_char_ptr(config.host), ip_buf)
    let raw_ip = ip_buf.bitcast[c_uint]().load()

    print("> inet_pton:", raw_ip, "status:", conv_status)

    let bin_port = htons(config.port)

    print("> htons:", bin_port)

    var ai = sockaddr_in(address_family, bin_port, raw_ip, StaticTuple[8, c_char]())
    let ai_ptr = Pointer[sockaddr_in].address_of(ai).bitcast[sockaddr]()

    let sockfd = socket(address_family, SOCK_STREAM, 0)
    if sockfd == -1:
        print("> socket creation error")

    print("> sockfd:", sockfd)

    var yes: Int = 1
    if (
        setsockopt(
            sockfd,
            SOL_SOCKET,
            SO_REUSEADDR,
            Pointer[Int].address_of(yes).bitcast[c_void](),
            sizeof[Int](),
        )
        == -1
    ):
        print("> set socket options failed")

    if bind(sockfd, ai_ptr, sizeof[sockaddr_in]()) == -1:
        _ = shutdown(sockfd, SHUT_RDWR)
        print("> binding socket failed")

    if listen(sockfd, c_int(128)) == -1:
        print("> listen", sockfd, "failed")

    print("> started at:", config.host, "port:", config.port, "with fd:", sockfd)

    print(
        "> using",
        config.workers,
        "workers",
        "out of",
        config.max_clients,
        "max clients",
    )
    wait_for_clients(db, sockfd, config.workers)

    _ = shutdown(sockfd, SHUT_RDWR)