"""
Utility methods for parsing data returned from MapD
"""
import datetime
from collections import namedtuple
from sqlalchemy import text
import mapd.ttypes as T

from ._utils import seconds_to_time


Description = namedtuple("Description", ["name", "type_code", "display_size",
                                         "internal_size", "precision", "scale",
                                         "null_ok"])
ColumnDetails = namedtuple("ColumnDetails", ["name", "type", "nullable",
                                             "precision", "scale",
                                             "comp_param"])

_typeattr = {
    'SMALLINT': 'int',
    'INT': 'int',
    'BIGINT': 'int',
    'TIME': 'int',
    'TIMESTAMP': 'int',
    'DATE': 'int',
    'BOOL': 'int',
    'FLOAT': 'real',
    'DECIMAL': 'real',
    'DOUBLE': 'real',
    'STR': 'str',
}
_thrift_types_to_values = T.TDatumType._NAMES_TO_VALUES
_thrift_values_to_types = T.TDatumType._VALUES_TO_NAMES


def _extract_row_val(desc, val):
    # type: (T.TColumnType, T.TDatum) -> Any
    typename = T.TDatumType._VALUES_TO_NAMES[desc.col_type.type]
    if val.is_null:
        return None
    val = getattr(val.val, _typeattr[typename] + '_val')
    base = datetime.datetime(1970, 1, 1)
    if typename == 'TIMESTAMP':
        val = (base + datetime.timedelta(seconds=val))
    elif typename == 'DATE':
        val = (base + datetime.timedelta(seconds=val)).date()
    elif typename == 'TIME':
        val = seconds_to_time(val)
    return val


def _extract_col_vals(desc, val):
    # type: (T.TColumnType, T.TColumn) -> Any
    typename = T.TDatumType._VALUES_TO_NAMES[desc.col_type.type]
    nulls = val.nulls

    vals = getattr(val.data, _typeattr[typename] + '_col')
    vals = [None if null else v
            for null, v in zip(nulls, vals)]

    base = datetime.datetime(1970, 1, 1)
    if typename == 'TIMESTAMP':
        vals = [None if v is None else base + datetime.timedelta(seconds=v)
                for v in vals]
    elif typename == 'DATE':
        vals = [None if v is None else (base +
                                        datetime.timedelta(seconds=v)).date()
                for v in vals]
    elif typename == 'TIME':
        vals = [None if v is None else seconds_to_time(v) for v in vals]

    return vals


def _extract_description(row_desc):
    # type: (List[T.TColumnType]) -> List[Description]
    """
    Return a tuple of (name, type_code, display_size, internal_size,
                       precision, scale, null_ok)

    https://www.python.org/dev/peps/pep-0249/#description
    """
    return [Description(col.col_name, col.col_type.type,
                        None, None, None, None,
                        col.col_type.nullable)
            for col in row_desc]


def _extract_column_details(row_desc):
    # For Connection.get_table_details
    return [
        ColumnDetails(x.col_name, _thrift_values_to_types[x.col_type.type],
                      x.col_type.nullable, x.col_type.precision,
                      x.col_type.scale, x.col_type.comp_param)
        for x in row_desc
    ]


def _is_columnar(data):
    # type: (T.TQueryResult) -> bool
    return data.row_set.is_columnar


def _load_schema(buf):
    """
    Load a `pyarrow.Schema` from a buffer written to shared memory

    Parameters
    ----------
    buf : pyarrow.Buffer

    Returns
    -------
    schema : pyarrow.Schema
    """
    import pyarrow as pa

    reader = pa.RecordBatchStreamReader(buf)
    return reader.schema


def _load_data(buf, schema):
    """
    Load a `pandas.DataFrame` from a buffer written to shared memory

    Parameters
    ----------
    buf : pyarrow.Buffer
    shcema : pyarrow.Schema

    Returns
    -------
    df : pandas.DataFrame
    """
    import pyarrow as pa

    message = pa.read_message(buf)
    rb = pa.read_record_batch(message, schema)
    return rb.to_pandas()


def _parse_tdf_gpu(tdf):
    """
    Parse the results of a select ipc_gpu into a GpuDataFrame

    Parameters
    ----------
    tdf : TDataFrame

    Returns
    -------
    gdf : GpuDataFrame
    """
    import numpy as np
    from pygdf.gpuarrow import GpuArrowReader
    from pygdf.dataframe import DataFrame
    from numba import cuda
    from numba.cuda.cudadrv import drvapi

    from .shm import load_buffer

    ipc_handle = drvapi.cu_ipc_mem_handle(*tdf.df_handle)
    ipch = cuda.driver.IpcHandle(None, ipc_handle, size=tdf.df_size)
    ctx = cuda.current_context()
    dptr = ipch.open(ctx)

    schema_buffer = load_buffer(tdf.sm_handle, tdf.sm_size)
    # TODO: extra copy.
    schema_buffer = np.frombuffer(schema_buffer.to_pybytes(), dtype=np.uint8)

    dtype = np.dtype(np.byte)
    darr = cuda.devicearray.DeviceNDArray(shape=dptr.size,
                                          strides=dtype.itemsize,
                                          dtype=dtype,
                                          gpu_data=dptr)
    reader = GpuArrowReader(schema_buffer, darr)
    df = DataFrame()
    for k, v in reader.to_dict().items():
        df[k] = v

    return df


def _bind_parameters(operation, parameters):
    return (text(operation)
            .bindparams(**parameters)
            .compile(compile_kwargs={"literal_binds": True}))

[((12, 14, 14, 52), 'collections.namedtuple', 'namedtuple', ({(12, 25, 12, 38): '"""Description"""', (12, 40, 14, 51): "['name', 'type_code', 'display_size', 'internal_size', 'precision', 'scale',\n    'null_ok']"}, {}), "('Description', ['name', 'type_code', 'display_size',\n    'internal_size', 'precision', 'scale', 'null_ok'])", False, 'from collections import namedtuple\n'), ((15, 16, 17, 59), 'collections.namedtuple', 'namedtuple', ({(15, 27, 15, 42): '"""ColumnDetails"""', (15, 44, 17, 58): "['name', 'type', 'nullable', 'precision', 'scale', 'comp_param']"}, {}), "('ColumnDetails', ['name', 'type', 'nullable', 'precision',\n    'scale', 'comp_param'])", False, 'from collections import namedtuple\n'), ((42, 11, 42, 40), 'datetime.datetime', 'datetime.datetime', ({(42, 29, 42, 33): '1970', (42, 35, 42, 36): '1', (42, 38, 42, 39): '1'}, {}), '(1970, 1, 1)', False, 'import datetime\n'), ((61, 11, 61, 40), 'datetime.datetime', 'datetime.datetime', ({(61, 29, 61, 33): '1970', (61, 35, 61, 36): '1', (61, 38, 61, 39): '1'}, {}), '(1970, 1, 1)', False, 'import datetime\n'), ((118, 13, 118, 44), 'pyarrow.RecordBatchStreamReader', 'pa.RecordBatchStreamReader', ({(118, 40, 118, 43): 'buf'}, {}), '(buf)', True, 'import pyarrow as pa\n'), ((137, 14, 137, 34), 'pyarrow.read_message', 'pa.read_message', ({(137, 30, 137, 33): 'buf'}, {}), '(buf)', True, 'import pyarrow as pa\n'), ((138, 9, 138, 46), 'pyarrow.read_record_batch', 'pa.read_record_batch', ({(138, 30, 138, 37): 'message', (138, 39, 138, 45): 'schema'}, {}), '(message, schema)', True, 'import pyarrow as pa\n'), ((162, 17, 162, 57), 'numba.cuda.cudadrv.drvapi.cu_ipc_mem_handle', 'drvapi.cu_ipc_mem_handle', ({(162, 42, 162, 56): '*tdf.df_handle'}, {}), '(*tdf.df_handle)', False, 'from numba.cuda.cudadrv import drvapi\n'), ((163, 11, 163, 68), 'numba.cuda.driver.IpcHandle', 'cuda.driver.IpcHandle', (), '', False, 'from numba import cuda\n'), ((164, 10, 164, 32), 'numba.cuda.current_context', 'cuda.current_context', ({}, {}), '()', False, 'from numba import cuda\n'), ((171, 12, 171, 29), 'numpy.dtype', 'np.dtype', ({(171, 21, 171, 28): 'np.byte'}, {}), '(np.byte)', True, 'import numpy as np\n'), ((172, 11, 175, 56), 'numba.cuda.devicearray.DeviceNDArray', 'cuda.devicearray.DeviceNDArray', (), '', False, 'from numba import cuda\n'), ((176, 13, 176, 48), 'pygdf.gpuarrow.GpuArrowReader', 'GpuArrowReader', ({(176, 28, 176, 41): 'schema_buffer', (176, 43, 176, 47): 'darr'}, {}), '(schema_buffer, darr)', False, 'from pygdf.gpuarrow import GpuArrowReader\n'), ((177, 9, 177, 20), 'pygdf.dataframe.DataFrame', 'DataFrame', ({}, {}), '()', False, 'from pygdf.dataframe import DataFrame\n'), ((44, 22, 44, 53), 'datetime.timedelta', 'datetime.timedelta', (), '', False, 'import datetime\n'), ((63, 46, 63, 75), 'datetime.timedelta', 'datetime.timedelta', (), '', False, 'import datetime\n'), ((185, 12, 185, 27), 'sqlalchemy.text', 'text', ({(185, 17, 185, 26): 'operation'}, {}), '(operation)', False, 'from sqlalchemy import text\n'), ((46, 22, 46, 53), 'datetime.timedelta', 'datetime.timedelta', (), '', False, 'import datetime\n'), ((67, 40, 67, 69), 'datetime.timedelta', 'datetime.timedelta', (), '', False, 'import datetime\n')]