import pyarrow as pa
import numpy as np
from itertools import repeat
from typing import Union, Dict, Any

from ray.experimental.data.deltacat.storage.model.types import DeltaType
from ray.experimental.data.deltacat.compute.compactor.model import delta_file_envelope as dfe

_SYS_COL_UUID = "4000f124-dfbd-48c6-885b-7b22621a6d41"


def _get_sys_col_name(suffix):
    return f"{_SYS_COL_UUID}_{suffix}"


_PK_HASH_DIGEST_BYTE_WIDTH = 20
_PK_HASH_COLUMN_NAME = _get_sys_col_name("hash")
_PK_HASH_COLUMN_TYPE = pa.binary(_PK_HASH_DIGEST_BYTE_WIDTH)
_PK_HASH_COLUMN_FIELD = pa.field(
    _PK_HASH_COLUMN_NAME,
    _PK_HASH_COLUMN_TYPE,
)

_DEDUPE_TASK_IDX_COLUMN_NAME = _get_sys_col_name("dedupe_task_idx")
_DEDUPE_TASK_IDX_COLUMN_TYPE = pa.int32()
_DEDUPE_TASK_IDX_COLUMN_FIELD = pa.field(
    _DEDUPE_TASK_IDX_COLUMN_NAME,
    _DEDUPE_TASK_IDX_COLUMN_TYPE,
)

_PARTITION_STREAM_POSITION_COLUMN_NAME = _get_sys_col_name("stream_position")
_PARTITION_STREAM_POSITION_COLUMN_TYPE = pa.int64()
_PARTITION_STREAM_POSITION_COLUMN_FIELD = pa.field(
    _PARTITION_STREAM_POSITION_COLUMN_NAME,
    _PARTITION_STREAM_POSITION_COLUMN_TYPE,
)

_ORDERED_FILE_IDX_COLUMN_NAME = _get_sys_col_name("file_index")
_ORDERED_FILE_IDX_COLUMN_TYPE = pa.int32()
_ORDERED_FILE_IDX_COLUMN_FIELD = pa.field(
    _ORDERED_FILE_IDX_COLUMN_NAME,
    _ORDERED_FILE_IDX_COLUMN_TYPE,
)

_ORDERED_RECORD_IDX_COLUMN_NAME = _get_sys_col_name("record_index")
_ORDERED_RECORD_IDX_COLUMN_TYPE = pa.int64()
_ORDERED_RECORD_IDX_COLUMN_FIELD = pa.field(
    _ORDERED_RECORD_IDX_COLUMN_NAME,
    _ORDERED_RECORD_IDX_COLUMN_TYPE,
)

_DELTA_TYPE_COLUMN_NAME = _get_sys_col_name("delta_type")
_DELTA_TYPE_COLUMN_TYPE = pa.bool_()
_DELTA_TYPE_COLUMN_FIELD = pa.field(
    _DELTA_TYPE_COLUMN_NAME,
    _DELTA_TYPE_COLUMN_TYPE,
)

_IS_SOURCE_COLUMN_NAME = _get_sys_col_name("is_source")
_IS_SOURCE_COLUMN_TYPE = pa.bool_()
_IS_SOURCE_COLUMN_FIELD = pa.field(
    _IS_SOURCE_COLUMN_NAME,
    _IS_SOURCE_COLUMN_TYPE,
)


def get_pk_hash_column_array(obj) -> Union[pa.Array, pa.ChunkedArray]:
    return pa.array(
        obj,
        _PK_HASH_COLUMN_TYPE
    )


def pk_hash_column_np(table: pa.Table) -> np.ndarray:
    return table[_PK_HASH_COLUMN_NAME].to_numpy()


def pk_hash_column(table: pa.Table) -> pa.ChunkedArray:
    return table[_PK_HASH_COLUMN_NAME]


def get_dedupe_task_idx_column_array(obj) -> Union[pa.Array, pa.ChunkedArray]:
    return pa.array(
        obj,
        _DEDUPE_TASK_IDX_COLUMN_TYPE,
    )


def get_stream_position_column_array(obj) -> Union[pa.Array, pa.ChunkedArray]:
    return pa.array(
        obj,
        _PARTITION_STREAM_POSITION_COLUMN_TYPE,
    )


def stream_position_column_np(table: pa.Table) -> np.ndarray:
    return table[_PARTITION_STREAM_POSITION_COLUMN_NAME].to_numpy()


def get_file_index_column_array(obj) \
        -> Union[pa.Array, pa.ChunkedArray]:
    return pa.array(
        obj,
        _ORDERED_FILE_IDX_COLUMN_TYPE,
    )


def file_index_column_np(table: pa.Table) -> np.ndarray:
    return table[_ORDERED_FILE_IDX_COLUMN_NAME].to_numpy()


def get_record_index_column_array(obj) -> \
        Union[pa.Array, pa.ChunkedArray]:
    return pa.array(
        obj,
        _ORDERED_RECORD_IDX_COLUMN_TYPE,
    )


def record_index_column_np(table: pa.Table) -> np.ndarray:
    return table[_ORDERED_RECORD_IDX_COLUMN_NAME].to_numpy()


def is_source_column_np(table: pa.Table) -> np.ndarray:
    return table[_IS_SOURCE_COLUMN_NAME].to_numpy()


def get_delta_type_column_array(obj) -> Union[pa.Array, pa.ChunkedArray]:
    return pa.array(
        obj,
        _DELTA_TYPE_COLUMN_TYPE,
    )


def get_is_source_column_array(obj) -> Union[pa.Array, pa.ChunkedArray]:
    return pa.array(
        obj,
        _IS_SOURCE_COLUMN_TYPE,
    )


def project_delta_file_metadata_on_table(delta_file_envelope: Dict[str, Any]) \
        -> pa.Table:

    table = dfe.get_table(delta_file_envelope)

    # append event timestamp column
    stream_position = dfe.get_stream_position(delta_file_envelope)
    stream_position_iterator = repeat(
        int(stream_position),
        len(table),
    )
    table = append_stream_position_column(table, stream_position_iterator)

    # append ordered file number column
    ordered_file_number = dfe.get_file_index(delta_file_envelope)
    ordered_file_number_iterator = repeat(
        int(ordered_file_number),
        len(table),
    )
    table = append_file_idx_column(table, ordered_file_number_iterator)

    # append delta type column
    delta_type = dfe.get_delta_type(delta_file_envelope)
    delta_type_iterator = repeat(
        delta_type_to_field(delta_type),
        len(table),
    )
    table = append_delta_type_col(table, delta_type_iterator)

    # append is source column
    is_source_iterator = repeat(
        True,
        len(table),
    )
    table = append_is_source_col(table, is_source_iterator)
    return table


def append_stream_position_column(
        table: pa.Table,
        stream_positions):

    table = table.append_column(
        _PARTITION_STREAM_POSITION_COLUMN_FIELD,
        get_stream_position_column_array(stream_positions)
    )
    return table


def append_file_idx_column(
        table: pa.Table,
        ordered_file_indices):

    table = table.append_column(
        _ORDERED_FILE_IDX_COLUMN_FIELD,
        get_file_index_column_array(ordered_file_indices)
    )
    return table


def append_pk_hash_column(
        table: pa.Table,
        pk_hashes) -> pa.Table:

    table = table.append_column(
        _PK_HASH_COLUMN_FIELD,
        get_pk_hash_column_array(pk_hashes)
    )
    return table


def append_record_idx_col(
        table: pa.Table,
        ordered_record_indices) -> pa.Table:

    table = table.append_column(
        _ORDERED_RECORD_IDX_COLUMN_FIELD,
        get_record_index_column_array(ordered_record_indices),
    )
    return table


def append_dedupe_task_idx_col(
        table: pa.Table,
        dedupe_task_indices) -> pa.Table:

    table = table.append_column(
        _DEDUPE_TASK_IDX_COLUMN_FIELD,
        get_dedupe_task_idx_column_array(dedupe_task_indices),
    )
    return table


def delta_type_to_field(delta_type: DeltaType) -> bool:
    return True if delta_type is DeltaType.UPSERT else False


def delta_type_from_field(delta_type_field: bool) -> DeltaType:
    return DeltaType.UPSERT if delta_type_field else DeltaType.DELETE


def append_delta_type_col(
        table: pa.Table,
        delta_types) -> pa.Table:

    table = table.append_column(
        _DELTA_TYPE_COLUMN_FIELD,
        get_delta_type_column_array(delta_types),
    )
    return table


def append_is_source_col(
        table: pa.Table,
        booleans) -> pa.Table:

    table = table.append_column(
        _IS_SOURCE_COLUMN_FIELD,
        get_is_source_column_array(booleans),
    )
    return table
