# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

# distutils: language = c++
# cython: embedsignature = True

include "config.pxi"

from cython.operator cimport dereference as deref

from kudu.compat import tobytes, frombytes
from kudu.schema cimport *
from kudu.errors cimport check_status
from kudu.client cimport PartialRow
from kudu.util import get_decimal_scale, to_unixtime_micros, to_unscaled_decimal, \
    unix_epoch_days_to_date, date_to_unix_epoch_days
from errors import KuduException

import six

from . import util

BOOL = KUDU_BOOL
STRING = KUDU_STRING
VARCHAR = KUDU_VARCHAR

INT8 = KUDU_INT8
INT16 = KUDU_INT16
INT32 = KUDU_INT32
INT64 = KUDU_INT64

FLOAT = KUDU_FLOAT
DOUBLE = KUDU_DOUBLE

UNIXTIME_MICROS = KUDU_UNIXTIME_MICROS
DATE = KUDU_DATE
BINARY = KUDU_BINARY

DECIMAL = KUDU_DECIMAL

cdef dict _reverse_dict(d):
    return dict((v, k) for k, v in d.items())


# CompressionType enums
COMPRESSION_DEFAULT = CompressionType_DEFAULT
COMPRESSION_NONE = CompressionType_NONE
COMPRESSION_SNAPPY = CompressionType_SNAPPY
COMPRESSION_LZ4 = CompressionType_LZ4
COMPRESSION_ZLIB = CompressionType_ZLIB

cdef dict _compression_types = {
    'default': COMPRESSION_DEFAULT,
    'none': COMPRESSION_NONE,
    'snappy': COMPRESSION_SNAPPY,
    'lz4': COMPRESSION_LZ4,
    'zlib': COMPRESSION_ZLIB,
}

cdef dict _compression_type_to_name = _reverse_dict(_compression_types)


# EncodingType enums
ENCODING_AUTO = EncodingType_AUTO
ENCODING_PLAIN = EncodingType_PLAIN
ENCODING_PREFIX = EncodingType_PREFIX
ENCODING_BIT_SHUFFLE = EncodingType_BIT_SHUFFLE
ENCODING_RLE = EncodingType_RLE
ENCODING_DICT = EncodingType_DICT

cdef dict _encoding_types = {
    'auto': ENCODING_AUTO,
    'plain': ENCODING_PLAIN,
    'prefix': ENCODING_PREFIX,
    'bitshuffle': ENCODING_BIT_SHUFFLE,
    'rle': ENCODING_RLE,
    'dict': ENCODING_DICT,
}

cdef dict _encoding_type_to_name = _reverse_dict(_encoding_types)


cdef class KuduType(object):

    """
    Usability wrapper for Kudu data type enum
    """

    def __cinit__(self, DataType type):
        self.type = type

    property name:

        def __get__(self):
            return _type_names[self.type]

    def __repr__(self):
        return 'KuduType({0})'.format(self.name)

    def new_value(self, value):
        return KuduValue(self, value)


int8 = KuduType(KUDU_INT8)
int16 = KuduType(KUDU_INT16)
int32 = KuduType(KUDU_INT32)
int64 = KuduType(KUDU_INT64)
string_ = KuduType(KUDU_STRING)
bool_ = KuduType(KUDU_BOOL)
float_ = KuduType(KUDU_FLOAT)
double_ = KuduType(KUDU_DOUBLE)
binary = KuduType(KUDU_BINARY)
unixtime_micros = KuduType(KUDU_UNIXTIME_MICROS)
decimal = KuduType(KUDU_DECIMAL)
varchar = KuduType(KUDU_VARCHAR)
date = KuduType(KUDU_DATE)


cdef dict _type_names = {
    INT8: 'int8',
    INT16: 'int16',
    INT32: 'int32',
    INT64: 'int64',
    STRING: 'string',
    BOOL: 'bool',
    FLOAT: 'float',
    DOUBLE: 'double',
    BINARY: 'binary',
    UNIXTIME_MICROS: 'unixtime_micros',
    DECIMAL: 'decimal',
    VARCHAR: 'varchar',
    DATE: 'date'
}


cdef dict _type_name_to_number = _reverse_dict(_type_names)

cdef dict _type_to_obj = {
    INT8: int8,
    INT16: int16,
    INT32: int32,
    INT64: int64,
    STRING: string_,
    BOOL: bool_,
    FLOAT: float_,
    DOUBLE: double_,
    BINARY: binary,
    UNIXTIME_MICROS: unixtime_micros,
    DECIMAL: decimal,
    VARCHAR: varchar,
    DATE: date
}


cdef KuduType to_data_type(object obj):
    if isinstance(obj, KuduType):
        return obj
    elif isinstance(obj, six.string_types):
        return _type_to_obj[_type_name_to_number[obj]]
    elif obj in _type_to_obj:
        return _type_to_obj[obj]
    else:
        raise ValueError('Invalid type: {0}'.format(obj))

cdef cppclass KuduColumnTypeAttributes:
        KuduColumnTypeAttributes()
        KuduColumnTypeAttributes(const KuduColumnTypeAttributes& other)
        KuduColumnTypeAttributes(int8_t precision, int8_t scale)
        KuduColumnTypeAttributes(uint16_t length)

        int8_t precision()
        int8_t scale()
        uint16_t length()

        c_bool Equals(KuduColumnTypeAttributes& other)
        void CopyFrom(KuduColumnTypeAttributes& other)

cdef class ColumnTypeAttributes:
    """
    Wraps a Kudu client ColumnTypeAttributes object.
    """
    def __cinit__(self):
        self.type_attributes = NULL

    def __dealloc__(self):
        if self.type_attributes is not NULL:
            del self.type_attributes

    property precision:
        def __get__(self):
            return self.type_attributes.precision()

    property scale:
        def __get__(self):
            return self.type_attributes.scale()

    property length:
        def __get__(self):
            return self.type_attributes.length()

    def __repr__(self):
        return ('ColumnTypeAttributes(precision=%s, scale=%s, length=%s)'
                % (self.type_attributes.precision(),
                   self.type_attributes.scale(),
                   self.type_attributes.length()))

cdef class ColumnSchema:
    """
    Wraps a Kudu client ColumnSchema object. Use schema.at(i) or schema[i] to
    construct one.
    """

    def __cinit__(self):
        self.schema = NULL
        self._type = None

    def __dealloc__(self):
        if self.schema is not NULL:
            del self.schema

    property name:
        def __get__(self):
            return frombytes(self.schema.name())

    property type:
        def __get__(self):
            if self._type is None:
                self._type = _type_to_obj[self.schema.type()]
            return self._type

    property nullable:
        def __get__(self):
            return self.schema.is_nullable()

    property type_attributes:
        def __get__(self):
            cdef ColumnTypeAttributes result = ColumnTypeAttributes()
            result.type_attributes = new KuduColumnTypeAttributes(self.schema.type_attributes())
            return result

    def equals(self, other):
        if not isinstance(other, ColumnSchema):
            return False
        return self.schema.Equals(deref((<ColumnSchema> other).schema))

    def __repr__(self):
        return ('ColumnSchema(name=%s, type=%s, nullable=%s)'
                % (self.name, self.type.name,
                   self.nullable))


#----------------------------------------------------------------------

cdef class ColumnSpec:

    """
    Helper class for configuring a column's settings while using the
    SchemaBuilder.
    """

    def type(self, type_):
        self._type = to_data_type(type_)
        self.spec.Type(self._type.type)
        return self

    def default(self, value):
        """
        Set a default value for the column
        """
        cdef:
            KuduValue kval

        if not self._type:
            raise ValueError("You must set the Column type before setting " +
                             "the default value.")
        else:
            kval = self._type.new_value(value)
            self.spec.Default(kval._value)
        return self

    def remove_default(self):
        """
        Remove a default value set.
        """
        self.spec.RemoveDefault()
        return self

    def compression(self, compression):
        """
        Set the compression type

        Parameters
        ----------
        compression : string or int
          One of {'default', 'none', 'snappy', 'lz4', 'zlib'}
          Or see kudu.COMPRESSION_* constants

        Returns
        -------
        self
        """
        cdef CompressionType type
        if isinstance(compression, int):
            # todo: validation
            type = <CompressionType> compression
        else:
            if compression is None:
                type = CompressionType_NONE
            else:
                try:
                    type = _compression_types[compression.lower()]
                except KeyError:
                    raise ValueError('Invalid compression type: {0}'
                                     .format(compression))

        self.spec.Compression(type)
        return self

    def encoding(self, encoding):
        """
        Set the encoding type

        Parameters
        ----------
        encoding : string or int
          One of {'auto', 'plain', 'prefix', 'bitshuffle', 'rle', 'dict'}
          Or see kudu.ENCODING_* constants

        Returns
        -------
        self
        """
        cdef EncodingType type
        if isinstance(encoding, six.string_types):
            try:
                type = _encoding_types[encoding.lower()]
            except KeyError:
                raise ValueError('Invalid encoding type: {0}'
                                 .format(encoding))
        else:
            # todo: validation
            type = <EncodingType> encoding

        self.spec.Encoding(type)
        return self

    def precision(self, precision):
        """
        Set the precision for the column.

        Clients must specify a precision for decimal columns. Precision is the total
        number of digits that can be represented by the column, regardless of the
        location of the decimal point. For example, representing integer values up to 9999,
        and fractional values up to 99.99, both require a precision of 4. You can also
        represent corresponding negative values, without any change in the precision.
        For example, the range -9999 to 9999 still only requires a precision of 4.

        The precision must be between 1 and 38.

        Returns
        -------
        self
        """
        self.spec.Precision(precision)
        return self

    def scale(self, scale):
        """
        Set the scale for the column.

        Clients can specify a scale for decimal columns. Scale represents the number
        of fractional digits. This value must be less than or equal to precision.
        A scale of 0 produces integral values, with no fractional part. If precision
        and scale are equal, all the digits come after the decimal point, making all
        the values between 0.9999 and -0.9999.

        The scale must be greater than 0 and less than the column's precision.
        If no scale is provided a default scale of 0 is used.

        Returns
        -------
        self
        """
        self.spec.Scale(scale)
        return self

    def length(self, length):
        """
        Set the length for the column.

        Clients can specify a length for varchar columns. Length is the maximum
        length in characters (UTF-8) of the string that the varchar can hold.

        The length must be between 1 and 65,535 (inclusive).

        Returns
        -------
        self
        """
        self.spec.Length(length)
        return self

    def primary_key(self):
        """
        Make this column a primary key. If you use this method, it will be the
        only primary key. Otherwise see set_primary_keys method on
        SchemaBuilder.

        Returns
        -------
        self
        """
        self.spec.PrimaryKey()
        return self

    def nullable(self, bint is_nullable=True):
        """
        Set nullable (True) or not nullable (False)

        Parameters
        ----------
        is_nullable : boolean, default True

        Returns
        -------
        self
        """
        if is_nullable:
            self.spec.Nullable()
        else:
            self.spec.NotNull()
        return self

    def rename(self, new_name):
        """
        Change the column name.
        """
        self.spec.RenameTo(tobytes(new_name))
        return self

    def block_size(self, block_size):
        """
        Set the target block size for the column.

        This is the number of bytes of user data packed per block on disk, and
        represents the unit of IO when reading the column. Larger values may
        improve scan performance, particularly on spinning media. Smaller
        values may improve random access performance, particularly for
        workloads that have high cache hit rates or operate on fast storage
        such as SSD.

        Parameters
        ----------
        block_size : int
          Block size (in bytes) to use.

        Returns
        -------
        self : ColumnSpec
        """
        self.spec.BlockSize(block_size)
        return self


cdef class SchemaBuilder:

    def copy_column(self, colschema):
        """
        Add a new column to the schema by copying it from an existing one.
        Returns a ColumnSpec object for further configuration and use in a
        fluid programming style. This method allows the SchemaBuilder to be
        more easily used to build a new Schema from an existing one.

        Parameters
        ----------
        colschema : ColumnSchema

        Examples
        --------
        for col in scanner.get_projection_schema():
            builder.copy_column(col).compression('lz4')
        builder.set_primary_keys(['key'])

        Returns
        -------
        spec : ColumnSpec
        """
        return self.add_column(colschema.name,
                               colschema.type,
                               colschema.nullable)

    def add_column(self, name, type_=None, nullable=None, compression=None,
                   encoding=None, primary_key=False, block_size=None,
                   default=None, precision=None, scale=None, length=None):
        """
        Add a new column to the schema. Returns a ColumnSpec object for further
        configuration and use in a fluid programming style.

        Parameters
        ----------
        name : string
        type_ : string or KuduType
          Data type e.g. 'int32' or kudu.int32
        nullable : boolean, default None
          New columns are nullable by default. Set boolean value for explicit
          nullable / not-nullable
        compression : string or int
          One of {'default', 'none', 'snappy', 'lz4', 'zlib'}
          Or see kudu.COMPRESSION_* constants
        encoding : string or int
          One of {'auto', 'plain', 'prefix', 'bitshuffle', 'rle', 'dict'}
          Or see kudu.ENCODING_* constants
        primary_key : boolean, default False
          Use this column as the table primary key
        block_size : int, optional
          Block size (in bytes) to use for the target column.
        default : obj
          Use this to set the column default value
        precision : int
          Use this precision for the decimal column
        scale : int
          Use this scale for the decimal column
        length : int
          Use this length for the varchar column

        Examples
        --------
        (builder.add_column('foo')
         .nullable(True)
         .compression('lz4'))

        Returns
        -------
        spec : ColumnSpec
        """
        cdef:
            ColumnSpec result = ColumnSpec()
            string c_name = tobytes(name)

        result.spec = self.builder.AddColumn(c_name)

        if type_ is not None:
            result.type(type_)

        if nullable is not None:
            result.nullable(nullable)

        if compression is not None:
            result.compression(compression)

        if encoding is not None:
            result.encoding(encoding)

        if precision is not None:
            result.precision(precision)

        if scale is not None:
            result.scale(scale)

        if length is not None:
            result.length(length)

        if primary_key:
            result.primary_key()

        if block_size:
            result.block_size(block_size)

        if default:
            result.default(default)

        return result

    def set_primary_keys(self, key_names):
        """
        Set indicated columns (by name) to be the primary keys of the table
        schema

        Parameters
        ----------
        key_names : list of Python strings

        Returns
        -------
        None
        """
        cdef:
            vector[string] key_col_names

        for name in key_names:
            key_col_names.push_back(tobytes(name))

        self.builder.SetPrimaryKey(key_col_names)

    def build(self):
        """
        Creates an immutable Schema object after the user has finished adding
        and onfiguring columns

        Returns
        -------
        schema : Schema
        """
        cdef Schema result = Schema()
        cdef KuduSchema* schema = new KuduSchema()
        check_status(self.builder.Build(schema))

        result.schema = schema
        return result


cdef class Schema:

    """
    Container for a Kudu table schema. Obtain from Table instances or create
    new ones using kudu.SchemaBuilder
    """

    def __cinit__(self):
        # Users should not call this directly
        self.schema = NULL
        self.own_schema = 1
        self._col_mapping.clear()
        self._mapping_initialized = 0

    def __dealloc__(self):
        if self.schema is not NULL and self.own_schema:
            del self.schema

    property names:

        def __get__(self):
            result = []
            for i in range(self.schema.num_columns()):
                name = frombytes(self.schema.Column(i).name())
                result.append(name)

            return result

    def __repr__(self):
        # Got to be careful with huge schemas, maybe some kind of summary repr
        # when more than 20-30 columns?
        buf = six.StringIO()

        col_names = self.names
        space = 2 + max(len(x) for x in col_names)

        for i in range(len(self)):
            col = self.at(i)
            not_null = '' if col.nullable else ' NOT NULL'

            buf.write('\n{0}{1}{2}'
                      .format(col.name.ljust(space),
                              col.type.name, not_null))

        pk_string = ', '.join(col_names[i] for i in self.primary_key_indices())
        buf.write('\nPRIMARY KEY ({0})'.format(pk_string))

        return "kudu.Schema {{{0}\n}}".format(util.indent(buf.getvalue(), 2))

    def __len__(self):
        return self.schema.num_columns()

    def __getitem__(self, key):
        if isinstance(key, six.string_types):
            key = self.get_loc(key)

        if key < 0:
            key += len(self)
        return self.at(key)

    def equals(self, Schema other):
        """
        Returns True if the table schemas are equal
        """
        return self.schema.Equals(deref(other.schema))

    cdef int get_loc(self, name) except -1:
        if not self._mapping_initialized:
            for i in range(self.schema.num_columns()):
                self._col_mapping[self.schema.Column(i).name()] = i
            self._mapping_initialized = 1

        name = tobytes(name)

        # TODO: std::map is slightly verbose and inefficient here (O(lg n)
        # lookups), may consider replacing with a better / different hash table
        # should it become a performance bottleneck
        cdef map[string, int].iterator it = self._col_mapping.find(name)
        if it == self._col_mapping.end():
            raise KeyError(name)
        return self._col_mapping[name]

    def at(self, size_t i):
        """
        Return the ColumnSchema for a column index. Analogous to schema[i].

        Returns
        -------
        col_schema : ColumnSchema
        """
        cdef ColumnSchema result = ColumnSchema()

        if i < 0 or i >= self.schema.num_columns():
            raise IndexError('Column index {0} is not in range'
                             .format(i))

        result.schema = new KuduColumnSchema(self.schema.Column(i))

        return result

    def new_row(self, record=None):
        """
        Create a new row corresponding to this schema. If a record is provided,
        a PartialRow will be initialized with values from the input record.
        The record can be in the form of a tuple, dict, or list. Dictionary
        keys can be either column names, indexes, or a mix of both names and
        indexes.

        Parameters
        ----------
        record : tuple/list/dict

        Returns
        -------
        row : PartialRow
        """
        result = PartialRow(self)
        result.row = self.schema.NewRow()

        if record:
            result.from_record(record)

        return result

    def primary_key_indices(self):
        """
        Return the indices of the columns used as primary keys

        Returns
        -------
        key_indices : list[int]
        """
        cdef:
            vector[int] indices
            size_t i

        self.schema.GetPrimaryKeyColumnIndexes(&indices)

        result = []
        for i in range(indices.size()):
            result.append(indices[i])
        return result

    def primary_keys(self):
        """
        Return the names of the columns used as primary keys

        Returns
        -------
        key_names : list[str]
        """
        indices = self.primary_key_indices()
        return [self.at(i).name for i in indices]


cdef class KuduValue:

    def __cinit__(self, KuduType type_, value):
        cdef:
            Slice slc

        if (type_.name[:3] == 'int'):
            self._value = C_KuduValue.FromInt(value)
        elif (type_.name in ['string', 'binary', 'varchar']):
            if isinstance(value, unicode):
                value = value.encode('utf8')

            slc = Slice(<char*> value, len(value))
            self._value = C_KuduValue.CopyString(slc)
        elif (type_.name == 'bool'):
            self._value = C_KuduValue.FromBool(value)
        elif (type_.name == 'float'):
            self._value = C_KuduValue.FromFloat(value)
        elif (type_.name == 'double'):
            self._value = C_KuduValue.FromDouble(value)
        elif (type_.name == 'unixtime_micros'):
            value = to_unixtime_micros(value)
            self._value = C_KuduValue.FromInt(value)
        elif (type_.name == 'date'):
            val = date_to_unix_epoch_days(value)
            self._value = C_KuduValue.FromInt(val)
        elif (type_.name == 'decimal'):
            IF PYKUDU_INT128_SUPPORTED == 1:
                scale = get_decimal_scale(value)
                value = to_unscaled_decimal(value)
                self._value = C_KuduValue.FromDecimal(value, scale)
            ELSE:
                raise KuduException("The decimal type is not supported when GCC version is < 4.6.0" % self)
        else:
            raise TypeError("Cannot initialize KuduValue for kudu type <{0}>"
                            .format(type_.name))

    def __dealloc__(self):
        # We don't own this.
        pass