
header "sys" as sys
header "openssl" as ossl

use core
use coro
use io
use utils

class Connection {
    ~ fd: FD
    ~ ssl: ?SSL (null)
    ~ ssl_enabled: bool (false)
    - closed: bool (false)

    + static fn new(fd: FD) Connection {
        io:set_non_block(fd, true)
        return Connection {
            fd: fd
        }
    }

    + fn close() {
        if this.closed : return
        this.closed = true
        Socket.close(this.fd)
    }

    fn gc_free() {
        this.close()
    }

    + fn ssl_connect(host: String, ca_cert_path: ?String (null)) !ssl_error {

        let fd = this.fd
        io:set_non_block(fd, false)

        if this.ssl_enabled : return

        let ssl = SSL.new(fd, host, ca_cert_path)

        let err = ossl:SSL_connect(ssl.ssl)
        if err != 1 {
            let ssl_err = ossl:SSL_get_error(ssl.ssl, err);
            let msg : ?String = null
            let e : uint = ossl:ERR_get_error()
            if e != 0 {
                msg = ossl:ERR_error_string(e, null)
            }
            throw ssl_error, msg ?? ("SSL error code: " + ssl_err)
        }
        this.ssl = ssl
        this.ssl_enabled = true

        io:set_non_block(fd, true)
    }

    + fn send(data: String) !connection {
        this.send_bytes(data.data, data.bytes, true) ! throw connection
    }
    + fn send_buffer(data: utils:ByteBuffer, skip_bytes: uint, send_all: bool) uint !connection {
        return this.send_bytes(data.data.offset(skip_bytes), data.length - skip_bytes, send_all) ! throw connection
    }

    + fn send_bytes(data: ptr, bytes: uint, send_all: bool) uint !connection !closed {

        if this.closed : throw closed

        if bytes == 0 : return 0
        let ssl = this.ssl
        let bytes_sent : uint = 0
        let bytes_to_send : uint = bytes

        while true {
            //////////////
            // SSL
            //////////////
            if isset(ssl) {
                let wbytes = ossl:SSL_write(ssl.ssl, data, bytes_to_send);
                if wbytes < 0 {
                    // SSL_ERROR_WANT_WRITE: 3
                    let err = ossl:SSL_get_error(ssl.ssl, wbytes);
                    if err == 3 {
                        // EAGAIN
                        let ev = coro:await_fd(this.fd, false, true)
                        if ev.is_closed() || !ev.is_writable() : throw closed
                        continue
                    }
                    throw connection;
                }

                let new_bytes = wbytes.@cast(uint)
                bytes_sent += new_bytes
                if bytes_sent == bytes : break
                if bytes_sent < bytes {
                    data += new_bytes
                    bytes_to_send -= new_bytes
                    if send_all : continue
                    break
                }

                throw connection
            }

            //////////////
            // NO SSL
            //////////////
            let ofd = io:os_fd(this.fd)
            #if OS == linux
            let wbytes = sys:send(ofd, data, bytes_to_send, sys:MSG_NOSIGNAL)
            #elif OS == win
            let wbytes = sys:send(ofd, data, bytes_to_send.@cast(i32), 0)
            #else
            let wbytes = sys:send(ofd, data, bytes_to_send, 0)
            #end

            if wbytes < 0 {
                #if OS == win
                let err = core:WSA.error()
                if err == sys:WSAEWOULDBLOCK {
                    // EAGAIN
                    let ev = coro:await_fd(this.fd, false, true)
                    if ev.is_closed() || !ev.is_writable() : throw closed
                    continue
                }
                #else
                let err = core:libc_errno()
                if err == sys:EAGAIN {
                    // EAGAIN
                    let ev = coro:await_fd(this.fd, false, true)
                    if ev.is_closed() || !ev.is_writable() : throw connection
                    continue
                }
                #end
                throw connection
            }

            let new_bytes = wbytes.@cast(uint)
            bytes_sent += new_bytes
            if bytes_sent == bytes : break
            if bytes_sent < bytes {
                data += new_bytes
                bytes_to_send -= new_bytes
                if send_all : continue
                break
            }

            throw connection
        }

        return bytes_sent
    }

    + fn recv(buffer: utils:ByteBuffer, bytes: uint) uint !connection !closed {

        if this.closed : throw closed

        buffer.minimum_free_space(bytes)
        let count : uint = 0
        let ssl = this.ssl;
        let err : i32 = 0

        while true {
            //////////////
            // SSL
            //////////////
            if isset(ssl) {
                let rcvd = ossl:SSL_read(ssl.ssl, buffer.data.offset(buffer.length), bytes);
                if rcvd <= 0 {
                    // SSL_ERROR_NONE: 0
                    // SSL_ERROR_WANT_READ: 2
                    // SSL_ERROR_ZERO_RETURN: 6
                    err = ossl:SSL_get_error(ssl.ssl, rcvd);
                    if err == 2 {
                        let ev = coro:await_fd(this.fd, true, false)
                        if ev.is_closed() || !ev.is_readable() : throw closed
                        continue
                    }
                    throw connection, "SSL Error #" + err;
                }
                count = rcvd.@cast(uint)
                break
            }

            //////////////
            // NO SSL
            //////////////
            #if OS == win
            let rcvd = sys:recv(this.fd.@cast(uint), buffer.data.offset(buffer.length), bytes.@cast(i32), 0)
            #else
            let rcvd = sys:recv(this.fd.@cast(i32), buffer.data.offset(buffer.length), bytes, 0)
            #end
            if (rcvd <= 0) {

                err = core:socket_errno()

                #if OS == win
                if err == sys:WSAEWOULDBLOCK {
                    // EAGAIN
                    let ev = coro:await_fd(this.fd, true, false)
                    if ev.is_closed() || !ev.is_readable() : throw closed
                    continue
                }
                if err == sys:WSAEINTR : continue // Ignore interrupts, retry recv
                if err == sys:WSAECONNRESET : throw closed
                #else
                if err == sys:EAGAIN {
                    // EAGAIN
                    let ev = coro:await_fd(this.fd, true, false)
                    if ev.is_closed() || !ev.is_readable() {
                        throw closed
                    }
                    if !ev.is_readable() {
                        throw closed
                    }
                    continue
                }
                if err == sys:EINTR : continue // Ignore interrupts, retry recv
                if err == sys:ECONNRESET : throw closed
                #end
                throw connection
            }
            count = rcvd.@cast(uint)
            break
        }
            
        buffer.length += count
        return count
    }
}
