# Copyright (c) 2022, California Institute of Technology and contributors
#
# You should have received a copy of the licensing terms for this
# software included in the file "LICENSE" located in the top-level
# directory of this package. If you did not, you can view a copy at
# https://git.ligo.org/ngdd/arrakis-server/-/raw/main/LICENSE

import itertools
from collections.abc import Generator, Iterable, Iterator
from typing import Any, TypeVar

import numpy
import pyarrow
from arrakis import Channel, schemas
from pyarrow import flight

T = TypeVar("T")


def batched(iterable: Iterable[T], n: int) -> Generator[Iterable[T], None, None]:
    """An implementation of python 3.12's itertools.batches.
    Taken from the python documentation for itertools recipies
    Given an Iterable object iterable, generate a series of
    Iterators that return chunks of 'n' items from iterable.
    The last batch may be smaller than n entries.
    """
    if n < 1:
        raise ValueError("n must be greater than or equal to 1")
    iterator = iter(iterable)
    while batch := tuple(itertools.islice(iterator, n)):
        yield batch


def create_batches(
    channels: Iterable[Channel],
    schema: pyarrow.Schema,
    batch_size: int = 1000,
) -> list[pyarrow.RecordBatch]:
    """Create record batches from channels using a schema-driven approach.

    Parameters
    ----------
    channels : Iterable[Channel]
        The channels to convert to batches.
    schema : pyarrow.Schema
        The Arrow schema defining the structure and field types.
    batch_size : int, optional
        Number of channels per batch (default: 1000).

    Returns
    -------
    list[pyarrow.RecordBatch]
        List of record batches conforming to the schema.

    """
    batches = []
    field_names = schema.names

    for channel_batch in batched(channels, batch_size):
        metadata = tuple(
            tuple(_extract_field_value(channel, field) for field in field_names)
            for channel in channel_batch
        )
        if metadata:
            field_data = [list(values) for values in zip(*metadata)]
        else:
            field_data = [[] for _ in field_names]

        arrays = [
            pyarrow.array(data, type=schema.field(field_name).type)
            for field_name, data in zip(field_names, field_data)
        ]
        batch = pyarrow.RecordBatch.from_arrays(arrays, schema=schema)
        batches.append(batch)

    return batches


def create_metadata_stream(channels: Iterable[Channel]) -> flight.RecordBatchStream:
    """Create a record batch stream from channel metadata."""
    schema = schemas.metadata()
    batches = create_batches(channels, schema)
    return flight.RecordBatchStream(
        pyarrow.RecordBatchReader.from_batches(schema, batches)
    )


def read_all_chunks(
    reader: flight.MetadataRecordBatchReader,
) -> Iterator[pyarrow.RecordBatch]:
    while True:
        try:
            batch, _ = reader.read_chunk()
            yield batch
        except StopIteration:
            return


def _extract_field_value(channel: Channel, field_name: str) -> Any:
    """Extract a value from a Channel object for a given schema field name."""
    if field_name == "channel":
        return channel.name
    elif field_name == "data_type":
        return numpy.dtype(channel.data_type).name
    else:
        return getattr(channel, field_name)
