from __future__ import annotations

from typing import Generator, Dict, Type, NamedTuple, List

from pyarrow import RecordBatch

from deltacat.experimental.storage.rivulet.reader.data_reader import (
    DataReader,
    MEMORY_FORMAT,
)
import pyarrow as pa


class RecordBatchRowIndex(NamedTuple):
    """
    Named tuple for a record batch with an index into a specific row
    Note that record batches store data by column, so the row index should be
    used to index into each column array
    """

    batch: RecordBatch
    row_index: int


class ArrowDataReader(DataReader[RecordBatchRowIndex]):
    """
    Parquet reader to iteratively load records from parquet files
    """

    def deserialize_records(
        self, record: RecordBatchRowIndex, output_type: Type[MEMORY_FORMAT]
    ) -> Generator[MEMORY_FORMAT, None, None]:
        """
        Deserialize records into the specified format.

        Note that output_type gets set based on what a DataScan converts results to,
            e.g. to_arrow, to_dict

        :param record: Input data (generated by generate_records method)
        :param output_type: Type to deserialize into
        :returns: A generator yielding records of the specified type.
        """
        batch, row_idx = record[0].batch, record[0].row_index

        if output_type == Dict:
            yield {
                column: batch.column(column_idx)[row_idx].as_py()
                for column_idx, column in enumerate(batch.column_names)
            }

        elif output_type == RecordBatch:
            # only yield full record batch if row_idx is 0.
            # TODO this logic will need to change in zipper use case across data formats
            if row_idx == 0:
                yield batch

    def join_deserialize_records(
        self,
        records: List[RecordBatchRowIndex],
        output_type: Type[MEMORY_FORMAT],
        join_key: str,
    ) -> Generator[MEMORY_FORMAT, None, None]:
        """
        Deserialize records into the specified format.

        Note that output_type gets set based on what a DataScan converts results to,
            e.g. to_arrow, to_dict

        :param records: Input data (generated by generate_records method)
        :param output_type: Type to deserialize into
        :returns: A generator yielding records of the specified type.
        """

        if output_type == Dict:
            yield self.__join_records_as_dict(records)
        elif output_type == RecordBatch:
            yield self.__join_records_as_record_batch(records, join_key)

    @staticmethod
    def __join_records_as_dict(records: List[RecordBatchRowIndex]) -> Dict[str, any]:
        """
        Deserialize records into a PyDict

        :param records: input record data
        :returns: A PyDict that's joined the given records around the primary key.
        """
        batch: RecordBatch
        row_idx: int
        out = {}
        for r in records:
            batch, row_idx = r
            # Note this stomps over join key but that's OK
            for column_idx, column in enumerate(batch.schema.names):
                col = batch.column(column_idx)
                if len(col) <= row_idx:
                    raise IndexError(
                        f"row index {row_idx} out of bounds for column {column} with length {len(col)}"
                    )

                out.update({column: col[row_idx].as_py()})
        return out

    @staticmethod
    def __join_records_as_record_batch(
        records: List[RecordBatchRowIndex], join_key: str
    ) -> RecordBatch:
        """
        Deserialize records into a RecordBatch

        :param records: input record data
        :returns: RecordBatch that's inner-joined the given records around the primary key.
        """
        batch: RecordBatch
        row_idx: int
        out: pa.Table | None = None
        for record in records:
            batch, row_idx = record
            batch_slice: RecordBatch = batch.slice(row_idx, 1)
            if not out:
                out = pa.Table.from_batches([batch_slice])
            else:
                table2 = pa.Table.from_batches([batch_slice])
                out = out.join(table2, keys=join_key, join_type="inner")
        return out.to_batches()[0]
