module fasthttp

import net
import sync.stdatomic
import time

#include <sys/epoll.h>
#include <sys/sendfile.h>
#include <sys/stat.h>
#include <netinet/tcp.h>

const epoll_wait_timeout_ms = 100
const status_408_response = 'HTTP/1.1 408 Request Timeout\r\nContent-Type: text/plain\r\nContent-Length: 19\r\nConnection: close\r\n\r\n408 Request Timeout'.bytes()

fn C.accept4(sockfd i32, addr &net.Addr, addrlen &u32, flags i32) i32

fn C.epoll_create1(__flags i32) i32

fn C.epoll_ctl(__epfd i32, __op i32, __fd i32, __event &C.epoll_event) i32

fn C.epoll_wait(__epfd i32, __events &C.epoll_event, __maxevents i32, __timeout i32) i32

fn C.sendfile(out_fd i32, in_fd i32, offset &i64, count usize) i32

fn C.fstat(fd i32, buf &C.stat) i32

@[typedef]
union C.epoll_data_t {
mut:
    ptr voidptr
    fd  int
    u32 u32
    u64 u64
}

struct C.epoll_event {
mut:
    events u32
    data   C.epoll_data_t
}

pub struct Server {
pub:
    family                  net.AddrFamily = .ip6
    port                    int            = 3000
    max_request_buffer_size int            = 8192
    timeout_in_seconds      int            = 30
    user_data               voidptr
mut:
    listen_fds      []int    = []int{len: max_thread_pool_size, cap: max_thread_pool_size, init: -1}
    epoll_fds       []int    = []int{len: max_thread_pool_size, cap: max_thread_pool_size, init: -1}
    threads         []thread = []thread{len: max_thread_pool_size, cap: max_thread_pool_size}
    request_handler fn (HttpRequest) !HttpResponse @[required]
    running         &stdatomic.AtomicVal[bool] = stdatomic.new_atomic(false)
    shutting_down   &stdatomic.AtomicVal[bool] = stdatomic.new_atomic(false)
    stopped         &stdatomic.AtomicVal[bool] = stdatomic.new_atomic(true)
    active_requests &stdatomic.AtomicVal[int]  = stdatomic.new_atomic(0)
}

// new_server creates and initializes a new Server instance.
pub fn new_server(config ServerConfig) !&Server {
    if config.max_request_buffer_size <= 0 {
        return error('max_request_buffer_size must be greater than 0')
    }
    mut server := &Server{
        family:                  config.family
        port:                    config.port
        max_request_buffer_size: config.max_request_buffer_size
        timeout_in_seconds:      config.timeout_in_seconds
        user_data:               config.user_data
        request_handler:         config.handler
        running:                 stdatomic.new_atomic(false)
        shutting_down:           stdatomic.new_atomic(false)
        stopped:                 stdatomic.new_atomic(true)
        active_requests:         stdatomic.new_atomic(0)
    }
    unsafe {
        server.listen_fds.flags.set(.noslices | .noshrink | .nogrow)
        server.epoll_fds.flags.set(.noslices | .noshrink | .nogrow)
        server.threads.flags.set(.noslices | .noshrink | .nogrow)
    }
    return server
}

fn set_blocking(fd int, blocking bool) {
    flags := C.fcntl(fd, C.F_GETFL, 0)
    if flags == -1 {
        // TODO: better error handling
        eprintln(@LOCATION)
        return
    }
    if blocking {
        // This removes the O_NONBLOCK flag from flags and set it.
        C.fcntl(fd, C.F_SETFL, flags & ~C.O_NONBLOCK)
    } else {
        // This adds the O_NONBLOCK flag from flags and set it.
        C.fcntl(fd, C.F_SETFL, flags | C.O_NONBLOCK)
    }
}

fn close_socket(fd int) bool {
    ret := C.close(fd)
    if ret == -1 {
        if C.errno == C.EINTR {
            // Interrupted by signal, retry is safe
            return close_socket(fd)
        }
        eprintln('ERROR: close(fd=${fd}) failed with errno=${C.errno}')
        return false
    }
    return true
}

fn create_server_socket(server Server) int {
    // Create a socket with non-blocking mode
    server_fd := C.socket(i32(server.family), i32(net.SocketType.tcp), 0)
    if server_fd < 0 {
        eprintln(@LOCATION)
        C.perror(c'Socket creation failed')
        return -1
    }

    set_blocking(server_fd, false)

    // Enable SO_REUSEADDR and SO_REUSEPORT
    opt := 1
    if C.setsockopt(server_fd, C.SOL_SOCKET, C.SO_REUSEADDR, &opt, sizeof(opt)) < 0 {
        eprintln(@LOCATION)
        C.perror(c'setsockopt SO_REUSEADDR failed')
        close_socket(server_fd)
        return -1
    }
    if C.setsockopt(server_fd, C.SOL_SOCKET, C.SO_REUSEPORT, &opt, sizeof(opt)) < 0 {
        eprintln(@LOCATION)
        C.perror(c'setsockopt SO_REUSEPORT failed')
        close_socket(server_fd)
        return -1
    }

    addr := if server.family == .ip6 {
        net.new_ip6(u16(server.port), [u8(0), 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]!)
    } else {
        net.new_ip(u16(server.port), [u8(0), 0, 0, 0]!)
    }
    alen := addr.len()
    if C.bind(server_fd, voidptr(&addr), alen) < 0 {
        eprintln(@LOCATION)
        C.perror(c'Bind failed')
        close_socket(server_fd)
        return -1
    }
    if C.listen(server_fd, max_connection_size) < 0 {
        eprintln(@LOCATION)
        C.perror(c'Listen failed')
        close_socket(server_fd)
        return -1
    }
    return server_fd
}

// add_fd_to_epoll adds a file descriptor to the epoll instance
fn add_fd_to_epoll(epoll_fd int, fd int, events u32) int {
    mut ev := C.epoll_event{
        events: events
    }
    ev.data.fd = fd
    if C.epoll_ctl(epoll_fd, C.EPOLL_CTL_ADD, fd, &ev) == -1 {
        eprintln(@LOCATION)
        C.perror(c'epoll_ctl')
        return -1
    }
    return 0
}

// remove_fd_from_epoll removes a file descriptor from the epoll instance
fn remove_fd_from_epoll(epoll_fd int, fd int) bool {
    ret := C.epoll_ctl(epoll_fd, C.EPOLL_CTL_DEL, fd, C.NULL)
    if ret == -1 {
        eprintln('ERROR: epoll_ctl(DEL, fd=${fd}) failed with errno=${C.errno}')
        return false
    }
    return true
}

fn handle_accept_loop(epoll_fd int, listen_fd int, mut client_fds map[int]bool) {
    for {
        client_fd := C.accept4(listen_fd, C.NULL, C.NULL, C.SOCK_NONBLOCK)
        if client_fd < 0 {
            if C.errno == C.EAGAIN || C.errno == C.EWOULDBLOCK {
                break // No more incoming connections; exit loop.
            }
            eprintln(@LOCATION)
            C.perror(c'Accept failed')
            break
        }
        // Enable TCP_NODELAY for lower latency
        opt := 1
        C.setsockopt(client_fd, C.IPPROTO_TCP, C.TCP_NODELAY, &opt, sizeof(opt))
        // Register client socket with epoll
        if add_fd_to_epoll(epoll_fd, client_fd, u32(C.EPOLLIN | C.EPOLLET)) == -1 {
            close_socket(client_fd)
            continue
        }
        client_fds[client_fd] = true
    }
}

fn handle_client_closure(epoll_fd int, client_fd int, mut client_fds map[int]bool, mut client_buffers map[int][]u8, mut client_read_starts map[int]i64, mut closing_client_fds map[int]bool) {
    // Never close the listening socket here
    if client_fd == 0 {
        return
    }
    if client_fd <= 0 {
        eprintln('ERROR: Invalid FD=${client_fd} for closure')
        return
    }
    client_fds.delete(client_fd)
    client_buffers.delete(client_fd)
    client_read_starts.delete(client_fd)
    closing_client_fds.delete(client_fd)
    remove_fd_from_epoll(epoll_fd, client_fd)
    close_socket(client_fd)
}

fn close_worker_clients(epoll_fd int, mut client_fds map[int]bool, mut client_buffers map[int][]u8, mut client_read_starts map[int]i64, mut closing_client_fds map[int]bool) {
    for client_fd in client_fds.keys() {
        handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
            client_read_starts, mut closing_client_fds)
    }
}

fn drain_closing_client(epoll_fd int, client_fd int, mut client_fds map[int]bool, mut client_buffers map[int][]u8, mut client_read_starts map[int]i64, mut closing_client_fds map[int]bool) {
    mut drain_buf := []u8{len: 4096}
    for {
        bytes_read := C.recv(client_fd, unsafe { &drain_buf[0] }, drain_buf.len, 0)
        if bytes_read < 0 {
            if C.errno == C.EAGAIN || C.errno == C.EWOULDBLOCK {
                return
            }
            handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
                client_read_starts, mut closing_client_fds)
            return
        }
        if bytes_read == 0 {
            handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
                client_read_starts, mut closing_client_fds)
            return
        }
    }
}

fn send_terminal_response_and_drain(client_fd int, response []u8, mut client_buffers map[int][]u8, mut client_read_starts map[int]i64, mut closing_client_fds map[int]bool) {
    C.send(client_fd, response.data, response.len, C.MSG_NOSIGNAL)
    net.shutdown(client_fd, how: .write)
    client_buffers.delete(client_fd)
    client_read_starts.delete(client_fd)
    closing_client_fds[client_fd] = true
}

fn process_request(server &Server, epoll_fd int, client_fd int, request_buffer []u8, mut client_fds map[int]bool, mut client_buffers map[int][]u8, mut client_read_starts map[int]i64, mut closing_client_fds map[int]bool) {
    mut request_arena := voidptr(unsafe { nil })
    $if prealloc {
        request_arena = unsafe { prealloc_scope_begin() }
    }
    client_read_starts.delete(client_fd)
    server.begin_request()
    defer {
        server.end_request()
    }
    mut decoded_http_request := decode_http_request(request_buffer) or {
        eprintln('Error decoding request ${err}')
        C.send(client_fd, tiny_bad_request_response.data, tiny_bad_request_response.len,
            C.MSG_NOSIGNAL)
        end_request_arena_current_thread(request_arena)
        handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
            client_read_starts, mut closing_client_fds)
        return
    }
    $if trace_prealloc ? {
        unsafe { prealloc_scope_checkpoint(c'fasthttp decoded request') }
    }
    decoded_http_request.client_conn_fd = client_fd
    decoded_http_request.user_data = server.user_data
    mut response := server.request_handler(decoded_http_request) or {
        eprintln('Error handling request ${err}')
        C.send(client_fd, tiny_bad_request_response.data, tiny_bad_request_response.len,
            C.MSG_NOSIGNAL)
        end_request_arena_current_thread(request_arena)
        handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
            client_read_starts, mut closing_client_fds)
        return
    }
    $if trace_prealloc ? {
        unsafe { prealloc_scope_checkpoint(c'fasthttp handler returned') }
    }
    response.attach_request_arena_if_empty(request_arena)
    defer {
        response.free_owned_content()
        response.end_request_arena_current_thread()
    }

    match response.takeover_mode {
        .manual {
            // The handler has taken ownership of the connection.
            // Remove from epoll and tracking, but do NOT close the fd.
            client_fds.delete(client_fd)
            client_buffers.delete(client_fd)
            client_read_starts.delete(client_fd)
            closing_client_fds.delete(client_fd)
            remove_fd_from_epoll(epoll_fd, client_fd)
            response.abandon_request_arena_current_thread()
            return
        }
        .reusable {
            set_blocking(client_fd, false)
            client_buffers.delete(client_fd)
            if server.is_shutting_down() || response.should_close {
                handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
                    client_read_starts, mut closing_client_fds)
            }
            return
        }
        .none {}
    }

    if response.content.len > 0 {
        mut send_error := false
        mut pos := 0
        for pos < response.content.len {
            sent := C.send(client_fd, unsafe { &response.content[pos] },
                response.content.len - pos, C.MSG_NOSIGNAL)
            if sent <= 0 {
                eprintln('ERROR: send() failed with errno=${C.errno}')
                send_error = true
                break
            }
            pos += sent
        }
        if send_error {
            handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
                client_read_starts, mut closing_client_fds)
            return
        }
    }

    if response.file_path != '' {
        mut fd := C.open(response.file_path.str, C.O_RDONLY, 0)
        if fd == -1 {
            eprintln('ERROR: open file failed')
            handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
                client_read_starts, mut closing_client_fds)
            return
        }
        defer {
            if fd != -1 {
                C.close(fd)
            }
        }
        mut st := C.stat{}
        if C.fstat(fd, &st) != 0 {
            eprintln('ERROR: fstat failed')
            handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
                client_read_starts, mut closing_client_fds)
            return
        }
        mut offset := i64(0)
        mut remaining := i64(st.st_size)
        mut sf_retries := 0
        for remaining > 0 {
            ssize := C.sendfile(client_fd, fd, &offset, usize(remaining))
            if ssize > 0 {
                remaining -= i64(ssize)
                sf_retries = 0
                continue
            }
            errno_val := C.errno
            match errno_val {
                C.EAGAIN, C.EWOULDBLOCK, C.EINTR {
                    if sf_retries < 3 {
                        sf_retries++
                        continue
                    }
                    eprintln('ERROR: sendfile() transient failure after ${sf_retries} retries (errno=${errno_val})')
                }
                C.EBADF {
                    eprintln('ERROR: sendfile() EBADF: input fd or socket not open for required access (errno=${errno_val})')
                }
                C.EFAULT {
                    eprintln('ERROR: sendfile() EFAULT: bad address for offset (errno=${errno_val})')
                }
                C.EINVAL {
                    eprintln('ERROR: sendfile() EINVAL: invalid descriptor state or non-seekable input (errno=${errno_val})')
                }
                C.EIO {
                    eprintln('ERROR: sendfile() EIO: I/O error while reading input file (errno=${errno_val})')
                }
                C.ENOMEM {
                    eprintln('ERROR: sendfile() ENOMEM: insufficient kernel memory (errno=${errno_val})')
                }
                C.EOVERFLOW {
                    eprintln('ERROR: sendfile() EOVERFLOW: count exceeds file/socket limits (errno=${errno_val})')
                }
                C.ESPIPE {
                    eprintln('ERROR: sendfile() ESPIPE: input file not seekable with offset (errno=${errno_val})')
                }
                else {
                    eprintln('ERROR: sendfile() failed with errno=${errno_val}')
                }
            }

            handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
                client_read_starts, mut closing_client_fds)
            return
        }
    }

    client_buffers.delete(client_fd)
    if server.is_shutting_down() || response.should_close {
        handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
            client_read_starts, mut closing_client_fds)
    }
}

fn process_events(server &Server, epoll_fd int, listen_fd int) {
    mut events := [max_connection_size]C.epoll_event{}
    mut request_buffer := []u8{len: server.max_request_buffer_size, cap: server.max_request_buffer_size}
    mut client_fds := map[int]bool{}
    mut client_buffers := map[int][]u8{}
    mut client_read_starts := map[int]i64{}
    mut closing_client_fds := map[int]bool{}
    unsafe {
        request_buffer.flags.set(.noslices | .nogrow | .noshrink)
    }
    for {
        if server.is_shutting_down() && server.active_request_count() == 0 {
            close_worker_clients(epoll_fd, mut client_fds, mut client_buffers, mut
                client_read_starts, mut closing_client_fds)
            return
        }
        num_events := C.epoll_wait(epoll_fd, &events[0], max_connection_size, epoll_wait_timeout_ms)
        if num_events < 0 {
            if C.errno == C.EINTR {
                continue
            }
            if server.is_shutting_down() {
                continue
            }
            eprintln('ERROR: epoll_wait() failed with errno=${C.errno}')
            continue
        }
        for i := 0; i < num_events; i++ {
            client_fd := unsafe { events[i].data.fd }
            // Accept new connections when the listening socket is readable
            if client_fd == listen_fd {
                if server.is_shutting_down() {
                    continue
                }
                handle_accept_loop(epoll_fd, listen_fd, mut client_fds)
                continue
            }

            if events[i].events & u32((C.EPOLLHUP | C.EPOLLERR)) != 0 {
                if client_fd == listen_fd {
                    eprintln('ERROR: listen fd had HUP/ERR')
                    continue
                }
                if client_fd > 0 {
                    // Try to send 444 No Response before closing abnormal connection
                    C.send(client_fd, status_444_response.data, status_444_response.len,
                        C.MSG_NOSIGNAL)
                    handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
                        client_read_starts, mut closing_client_fds)
                } else {
                    eprintln('ERROR: Invalid FD from epoll: ${client_fd}')
                }
                continue
            }
            if events[i].events & u32(C.EPOLLIN) != 0 {
                if server.is_shutting_down() {
                    handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
                        client_read_starts, mut closing_client_fds)
                    continue
                }
                if closing_client_fds[client_fd] or { false } {
                    drain_closing_client(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
                        client_read_starts, mut closing_client_fds)
                    continue
                }
                // Read all available data from the socket
                mut total_bytes_read := 0
                mut readed_request_buffer := client_buffers[client_fd] or {
                    []u8{cap: server.max_request_buffer_size}
                }
                mut headers_complete := false
                mut header_too_large := false
                mut header_end_pos := -1
                mut request_complete := false
                mut peer_closed := false
                mut recv_error := false

                for {
                    bytes_read := C.recv(client_fd, unsafe { &request_buffer[0] },
                        server.max_request_buffer_size - 1, 0)
                    if bytes_read < 0 {
                        if C.errno == C.EAGAIN || C.errno == C.EWOULDBLOCK {
                            // No more data available right now
                            break
                        }
                        // Error occurred
                        eprintln('ERROR: recv() failed with errno=${C.errno}')
                        recv_error = true
                        break
                    } else if bytes_read == 0 {
                        // Connection closed by client
                        peer_closed = true
                        break
                    }

                    unsafe {
                        readed_request_buffer.push_many(&request_buffer[0], bytes_read)
                    }
                    total_bytes_read += bytes_read
                    if client_fd !in client_read_starts {
                        client_read_starts[client_fd] = time.sys_mono_now()
                    }

                    // Enforce the configured limit on request headers, not on the whole body.
                    buffer_len := readed_request_buffer.len
                    if !headers_complete && buffer_len >= 4 {
                        header_end_pos = find_header_end_in_buf(readed_request_buffer.data,
                            buffer_len)
                        if header_end_pos == -1 {
                            if buffer_len >= server.max_request_buffer_size {
                                header_too_large = true
                                break
                            }
                        } else {
                            headers_complete = true
                            if header_end_pos > server.max_request_buffer_size {
                                header_too_large = true
                                break
                            }
                        }
                    }

                    if headers_complete && has_complete_body(readed_request_buffer.data, buffer_len) {
                        request_complete = true
                        break
                    }
                }

                if header_too_large {
                    send_terminal_response_and_drain(client_fd, status_413_response, mut
                        client_buffers, mut client_read_starts, mut closing_client_fds)
                    continue
                }
                if request_complete {
                    process_request(server, epoll_fd, client_fd, readed_request_buffer, mut
                        client_fds, mut client_buffers, mut client_read_starts, mut
                        closing_client_fds)
                } else if recv_error {
                    // Unexpected recv error - send 444 No Response
                    C.send(client_fd, status_444_response.data, status_444_response.len,
                        C.MSG_NOSIGNAL)
                    handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
                        client_read_starts, mut closing_client_fds)
                } else if peer_closed || (total_bytes_read == 0 && readed_request_buffer.len == 0) {
                    // Normal client closure (FIN received)
                    handle_client_closure(epoll_fd, client_fd, mut client_fds, mut client_buffers, mut
                        client_read_starts, mut closing_client_fds)
                } else if readed_request_buffer.len > 0 {
                    client_buffers[client_fd] = readed_request_buffer
                }
            }
        }
        if server.timeout_in_seconds > 0 {
            now := time.sys_mono_now()
            timeout_ns := i64(server.timeout_in_seconds) * 1_000_000_000
            for client_fd in client_read_starts.keys() {
                started := client_read_starts[client_fd] or { continue }
                if now - started >= timeout_ns {
                    send_terminal_response_and_drain(client_fd, status_408_response, mut
                        client_buffers, mut client_read_starts, mut closing_client_fds)
                }
            }
        }
    }
}

fn (mut server Server) stop_accepting() {
    for i := 0; i < max_thread_pool_size; i++ {
        if server.listen_fds[i] < 0 {
            continue
        }
        if server.epoll_fds[i] >= 0 {
            remove_fd_from_epoll(server.epoll_fds[i], server.listen_fds[i])
        }
        close_socket(server.listen_fds[i])
        server.listen_fds[i] = -1
    }
}

// run starts the server and begins listening for incoming connections.
pub fn (mut server Server) run() ! {
    $if windows {
        eprintln('Windows is not supported yet')
        return
    }
    for i := 0; i < max_thread_pool_size; i++ {
        server.listen_fds[i] = create_server_socket(server)
        if server.listen_fds[i] < 0 {
            return
        }

        server.epoll_fds[i] = C.epoll_create1(0)
        if server.epoll_fds[i] < 0 {
            C.perror(c'epoll_create1 failed')
            close_socket(server.listen_fds[i])
            return
        }

        // Register the listening socket with each worker epoll for distributed accepts (edge-triggered)
        if add_fd_to_epoll(server.epoll_fds[i], server.listen_fds[i], u32(C.EPOLLIN | C.EPOLLET)) == -1 {
            close_socket(server.listen_fds[i])
            close_socket(server.epoll_fds[i])
            return
        }

        server.threads[i] = spawn process_events(server, server.epoll_fds[i], server.listen_fds[i])
    }

    server.mark_running()
    println('listening on http://0.0.0.0:${server.port}/')
    // Main thread waits for workers; accepts are handled in worker epoll loops
    for i in 0 .. max_thread_pool_size {
        server.threads[i].wait()
        if server.epoll_fds[i] >= 0 {
            close_socket(server.epoll_fds[i])
            server.epoll_fds[i] = -1
        }
    }
    server.mark_stopped()
}