# Copyright (C) 2016-present the asyncpg authors and contributors
# <see AUTHORS file>
#
# This module is part of asyncpg and is released under
# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0


from asyncpg import types as apg_types

from collections.abc import Sequence as SequenceABC

# defined in postgresql/src/include/utils/rangetypes.h
DEF RANGE_EMPTY  = 0x01  # range is empty
DEF RANGE_LB_INC = 0x02  # lower bound is inclusive
DEF RANGE_UB_INC = 0x04  # upper bound is inclusive
DEF RANGE_LB_INF = 0x08  # lower bound is -infinity
DEF RANGE_UB_INF = 0x10  # upper bound is +infinity


cdef enum _RangeArgumentType:
    _RANGE_ARGUMENT_INVALID = 0
    _RANGE_ARGUMENT_TUPLE = 1
    _RANGE_ARGUMENT_RANGE = 2


cdef inline bint _range_has_lbound(uint8_t flags):
    return not (flags & (RANGE_EMPTY | RANGE_LB_INF))


cdef inline bint _range_has_ubound(uint8_t flags):
    return not (flags & (RANGE_EMPTY | RANGE_UB_INF))


cdef inline _RangeArgumentType _range_type(object obj):
    if cpython.PyTuple_Check(obj) or cpython.PyList_Check(obj):
        return _RANGE_ARGUMENT_TUPLE
    elif isinstance(obj, apg_types.Range):
        return _RANGE_ARGUMENT_RANGE
    else:
        return _RANGE_ARGUMENT_INVALID


cdef range_encode(ConnectionSettings settings, WriteBuffer buf,
                  object obj, uint32_t elem_oid,
                  encode_func_ex encoder, const void *encoder_arg):
    cdef:
        ssize_t obj_len
        uint8_t flags = 0
        object lower = None
        object upper = None
        WriteBuffer bounds_data = WriteBuffer.new()
        _RangeArgumentType arg_type = _range_type(obj)

    if arg_type == _RANGE_ARGUMENT_INVALID:
        raise TypeError(
            'list, tuple or Range object expected (got type {})'.format(
                type(obj)))

    elif arg_type == _RANGE_ARGUMENT_TUPLE:
        obj_len = len(obj)
        if obj_len == 2:
            lower = obj[0]
            upper = obj[1]

            if lower is None:
                flags |= RANGE_LB_INF

            if upper is None:
                flags |= RANGE_UB_INF

            flags |= RANGE_LB_INC | RANGE_UB_INC

        elif obj_len == 1:
            lower = obj[0]
            flags |= RANGE_LB_INC | RANGE_UB_INF

        elif obj_len == 0:
            flags |= RANGE_EMPTY

        else:
            raise ValueError(
                'expected 0, 1 or 2 elements in range (got {})'.format(
                    obj_len))

    else:
        if obj.isempty:
            flags |= RANGE_EMPTY
        else:
            lower = obj.lower
            upper = obj.upper

            if obj.lower_inc:
                flags |= RANGE_LB_INC
            elif lower is None:
                flags |= RANGE_LB_INF

            if obj.upper_inc:
                flags |= RANGE_UB_INC
            elif upper is None:
                flags |= RANGE_UB_INF

    if _range_has_lbound(flags):
        encoder(settings, bounds_data, lower, encoder_arg)

    if _range_has_ubound(flags):
        encoder(settings, bounds_data, upper, encoder_arg)

    buf.write_int32(1 + bounds_data.len())
    buf.write_byte(<int8_t>flags)
    buf.write_buffer(bounds_data)


cdef range_decode(ConnectionSettings settings, FRBuffer *buf,
                  decode_func_ex decoder, const void *decoder_arg):
    cdef:
        uint8_t flags = <uint8_t>frb_read(buf, 1)[0]
        int32_t bound_len
        object lower = None
        object upper = None
        FRBuffer bound_buf

    if _range_has_lbound(flags):
        bound_len = hton.unpack_int32(frb_read(buf, 4))
        if bound_len == -1:
            lower = None
        else:
            frb_slice_from(&bound_buf, buf, bound_len)
            lower = decoder(settings, &bound_buf, decoder_arg)

    if _range_has_ubound(flags):
        bound_len = hton.unpack_int32(frb_read(buf, 4))
        if bound_len == -1:
            upper = None
        else:
            frb_slice_from(&bound_buf, buf, bound_len)
            upper = decoder(settings, &bound_buf, decoder_arg)

    return apg_types.Range(lower=lower, upper=upper,
                           lower_inc=(flags & RANGE_LB_INC) != 0,
                           upper_inc=(flags & RANGE_UB_INC) != 0,
                           empty=(flags & RANGE_EMPTY) != 0)


cdef multirange_encode(ConnectionSettings settings, WriteBuffer buf,
                       object obj, uint32_t elem_oid,
                       encode_func_ex encoder, const void *encoder_arg):
    cdef:
        WriteBuffer elem_data
        ssize_t elem_data_len
        ssize_t elem_count

    if not isinstance(obj, SequenceABC):
        raise TypeError(
            'expected a sequence (got type {!r})'.format(type(obj).__name__)
        )

    elem_data = WriteBuffer.new()

    for elem in obj:
        range_encode(settings, elem_data, elem, elem_oid, encoder, encoder_arg)

    elem_count = len(obj)
    if elem_count > INT32_MAX:
        raise OverflowError(f'too many elements in multirange value')

    elem_data_len = elem_data.len()
    if elem_data_len > INT32_MAX - 4:
        raise OverflowError(
            f'size of encoded multirange datum exceeds the maximum allowed'
            f' {INT32_MAX - 4} bytes')

    # Datum length
    buf.write_int32(4 + <int32_t>elem_data_len)
    # Number of elements in multirange
    buf.write_int32(<int32_t>elem_count)
    buf.write_buffer(elem_data)


cdef multirange_decode(ConnectionSettings settings, FRBuffer *buf,
                       decode_func_ex decoder, const void *decoder_arg):
    cdef:
        int32_t nelems = hton.unpack_int32(frb_read(buf, 4))
        FRBuffer elem_buf
        int32_t elem_len
        int i
        list result

    if nelems == 0:
        return []

    if nelems < 0:
        raise exceptions.ProtocolError(
            'unexpected multirange size value: {}'.format(nelems))

    result = cpython.PyList_New(nelems)
    for i in range(nelems):
        elem_len = hton.unpack_int32(frb_read(buf, 4))
        if elem_len == -1:
            raise exceptions.ProtocolError(
                'unexpected NULL element in multirange value')
        else:
            frb_slice_from(&elem_buf, buf, elem_len)
        elem = range_decode(settings, &elem_buf, decoder, decoder_arg)
        cpython.Py_INCREF(elem)
        cpython.PyList_SET_ITEM(result, i, elem)

    return result