# 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

"""Arrow Flight schema definitions."""

from collections.abc import Iterable

import pyarrow

from . import Channel

METADATA_FIELDS = {
    "name": pyarrow.field("channel", pyarrow.string(), nullable=False),
    "data_type": pyarrow.field("data_type", pyarrow.string(), nullable=False),
    "sample_rate": pyarrow.field("sample_rate", pyarrow.int32(), nullable=False),
    "publisher": pyarrow.field("publisher", pyarrow.string()),
    "partition_id": pyarrow.field("partition_id", pyarrow.string()),
    "partition_index": pyarrow.field("partition_index", pyarrow.uint32()),
    "stride": pyarrow.field("stride", pyarrow.uint64()),
    "max_latency": pyarrow.field("max_latency", pyarrow.uint64()),
}


def metadata(scope: str | None = None) -> pyarrow.Schema:
    """Create an Arrow Flight metadata schema.

    Parameters
    ----------
    scope : None | str
        Metadata scope, currently accepts only None or "partition".

    Returns
    -------
    pyarrow.Schema
        The metadata schema.

    """
    fields = [
        "name",
        "data_type",
        "sample_rate",
        "partition_id",
        "partition_index",
    ]
    if scope != "partition":
        fields.extend(
            [
                "publisher",
                "stride",
                "max_latency",
            ]
        )
    return pyarrow.schema([METADATA_FIELDS[field] for field in fields])


def count() -> pyarrow.Schema:
    """Create an Arrow Flight schema for `count`.

    Returns
    -------
    pyarrow.Schema
        The count schema.

    """
    return pyarrow.schema([pyarrow.field("count", pyarrow.int64())])


def stream(channels: Iterable[Channel]) -> pyarrow.Schema:
    """Create an Arrow Flight schema for `stream`.

    Parameters
    ----------
    channels : Iterable[Channel]
        The list of channels for the stream request.

    Returns
    -------
    pyarrow.Schema
        The stream schema.

    """
    columns = [pyarrow.field("time", pyarrow.int64(), nullable=False)]
    for channel in channels:
        dtype = pyarrow.from_numpy_dtype(channel.data_type)
        field = pyarrow.field(channel.name, pyarrow.list_(dtype)).with_metadata(
            {"rate": str(channel.sample_rate)}
        )
        columns.append(field)
    return pyarrow.schema(columns)


def publish() -> pyarrow.Schema:
    """Create an Arrow Flight schema for `publish`.

    Returns
    -------
    pyarrow.Schema
        The publish schema.

    """
    dtype = pyarrow.map_(pyarrow.string(), pyarrow.string())
    return pyarrow.schema([pyarrow.field("properties", dtype, nullable=False)])
