# 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-python/-/raw/main/LICENSE

"""Client-based API access."""

import json
import logging
from collections.abc import Generator
from typing import TYPE_CHECKING, TypeAlias
from urllib.parse import urlparse

import numpy
from pyarrow import flight
from pyarrow.flight import connect

from . import constants
from .block import SeriesBlock, concatenate_blocks, time_as_ns
from .channel import Channel
from .flight import (
    MultiFlightReader,
    RequestType,
    RequestValidator,
    create_descriptor,
    parse_url,
)
from .kafka import KafkaReader
from .mux import BlockMuxStream

if TYPE_CHECKING:
    from .stream import StreamReader

logger = logging.getLogger("arrakis")

DataTypeLike: TypeAlias = (
    str | list[str] | type | list[type] | numpy.dtype | list[numpy.dtype]
)


def get_flight_info(
    client: flight.FlightClient, descriptor: flight.FlightDescriptor
) -> flight.FlightInfo:
    logger.debug(
        "retrieving '%s' flight info...", json.loads(descriptor.command)["request"]
    )
    flight_info = client.get_flight_info(descriptor)
    n_endpoints = len(flight_info.endpoints)
    logger.debug("%g endpoint%s identified", n_endpoints, "s"[: n_endpoints ^ 1])
    return flight_info


class Client:
    """Retrieve channel information and timeseries data from Arrakis.

    Parameters
    ----------
    url : str, optional
        The URL to connect to.
        If the URL is not set, connect to a default server
        or one set by ARRAKIS_SERVER.

    """

    def __init__(self, url: str | None = None):
        self.initial_url = parse_url(url)
        logger.debug("initial url: %s", self.initial_url)
        self._validator = RequestValidator()

    def find(
        self,
        pattern: str = constants.DEFAULT_MATCH,
        data_type: DataTypeLike | None = None,
        min_rate: int | None = constants.MIN_SAMPLE_RATE,
        max_rate: int | None = constants.MAX_SAMPLE_RATE,
        publisher: str | list[str] | None = None,
    ) -> Generator[Channel, None, None]:
        """Find channels matching a set of conditions

        Parameters
        ----------
        pattern : str, optional
            Channel pattern to match channels with, using regular expressions.
        data_type : numpy.dtype-like | list[numpy.dtype-like], optional
            If set, find all channels with these data types.
        min_rate : int, optional
            Minimum sampling rate for channels.
        max_rate : int, optional
            Maximum sampling rate for channels.
        publisher : str | list[str], optional
            If set, find all channels associated with these publishers.

        Yields
        -------
        Channel
            Channel objects for all channels matching query.

        """
        data_type = _parse_data_types(data_type)
        if min_rate is None:
            min_rate = constants.MIN_SAMPLE_RATE
        if max_rate is None:
            max_rate = constants.MAX_SAMPLE_RATE
        if publisher is None:
            publisher = []
        elif isinstance(publisher, str):
            publisher = [publisher]

        descriptor = create_descriptor(
            RequestType.Find,
            pattern=pattern,
            data_type=data_type,
            min_rate=min_rate,
            max_rate=max_rate,
            publisher=publisher,
            validator=self._validator,
        )
        with connect(self.initial_url) as client:
            yield from self._stream_channel_metadata(client, descriptor)

    def count(
        self,
        pattern: str = constants.DEFAULT_MATCH,
        data_type: DataTypeLike | None = None,
        min_rate: int | None = constants.MIN_SAMPLE_RATE,
        max_rate: int | None = constants.MAX_SAMPLE_RATE,
        publisher: str | list[str] | None = None,
    ) -> int:
        """Count channels matching a set of conditions

        Parameters
        ----------
        pattern : str, optional
            Channel pattern to match channels with, using regular expressions.
        data_type : numpy.dtype-like | list[numpy.dtype-like], optional
            If set, find all channels with these data types.
        min_rate : int, optional
            The minimum sampling rate for channels.
        max_rate : int, optional
            The maximum sampling rate for channels.
        publisher : str | list[str], optional
            If set, find all channels associated with these publishers.

        Returns
        -------
        int
            The number of channels matching query.

        """
        data_type = _parse_data_types(data_type)
        if min_rate is None:
            min_rate = constants.MIN_SAMPLE_RATE
        if max_rate is None:
            max_rate = constants.MAX_SAMPLE_RATE
        if publisher is None:
            publisher = []
        elif isinstance(publisher, str):
            publisher = [publisher]

        descriptor = create_descriptor(
            RequestType.Count,
            pattern=pattern,
            data_type=data_type,
            min_rate=min_rate,
            max_rate=max_rate,
            publisher=publisher,
            validator=self._validator,
        )
        count = 0
        with connect(self.initial_url) as client:
            flight_info = get_flight_info(client, descriptor)
            with MultiFlightReader(flight_info.endpoints, client) as stream:
                for data in stream.unpack():
                    count += data["count"]
        return count

    def _describe(
        self,
        client: flight.FlightClient,
        channels: list[str],
    ) -> dict[str, Channel]:
        descriptor = create_descriptor(
            RequestType.Describe, channels=channels, validator=self._validator
        )
        return {
            channel.name: channel
            for channel in self._stream_channel_metadata(client, descriptor)
        }

    def describe(self, channels: list[str]) -> dict[str, Channel]:
        """Get channel metadata for channels requested

        Parameters
        ----------
        channels : list[str]
            List of channels to request.

        Returns
        -------
        dict[str, Channel]
            Mapping of channel names to channel metadata.

        """
        with connect(self.initial_url) as client:
            return self._describe(client, channels)

    def stream(
        self,
        channels: list[str],
        start: float | None = None,
        end: float | None = None,
        kafka_url: str | None = None,
    ) -> Generator[SeriesBlock, None, None]:
        """Stream live or offline timeseries data

        Parameters
        ----------
        channels : list[str]
            List of channels to request.
        start : float, optional
            GPS start time, in seconds.
        end : float, optional
            GPS end time, in seconds.

        Yields
        ------
        SeriesBlock
            Dictionary-like object containing all requested channel data.

        Setting neither start nor end begins a live stream starting
        from now.

        """
        start_ns = time_as_ns(start) if start is not None else None
        end_ns = time_as_ns(end) if end is not None else None
        metadata: dict[str, Channel] = {}
        reader: StreamReader

        with connect(self.initial_url) as client:
            # initial describe request to get the full metadata needed
            # to construct the appropriate muxer
            metadata = self._describe(client, channels)
            for channel in metadata.values():
                assert channel.stride, "Channels do not include stride info."
                assert channel.max_latency, "Channels do not include latency info."

            # get the flight info for streams
            descriptor = create_descriptor(
                RequestType.Stream,
                channels=channels,
                start=start_ns,
                end=end_ns,
                validator=self._validator,
            )
            flight_info = get_flight_info(client, descriptor)

        # confirm that we don't have kafka endpoints mixed with
        # other endpoints
        if len(flight_info.endpoints) > 1:
            for endpoint in flight_info.endpoints:
                assert not _endpoint_is_kafka(endpoint), (
                    "Only a single kafka endpoint currently supported."
                )

        # construct reader
        if kafka_url is not None or _endpoint_is_kafka(flight_info.endpoints[0]):
            if kafka_url is None:
                kafka_url = flight_info.endpoint.locations[0]
            reader = KafkaReader(
                kafka_url,
                metadata,
                start_ns,
            )
        else:
            reader = MultiFlightReader(
                flight_info.endpoints,
                self.initial_url,
                metadata=metadata,
            )

        mux: BlockMuxStream = BlockMuxStream(reader, start=start_ns)

        for block in mux.stream(end_ns):
            yield block

    def fetch(
        self,
        channels: list[str],
        start: float,
        end: float,
    ) -> SeriesBlock:
        """Fetch timeseries data

        Parameters
        ----------
        channels : list[str]
            List of channels to request.
        start : float
            GPS start time, in seconds.
        end : float
            GPS end time, in seconds.

        Returns
        -------
        SeriesBlock
            Dictionary-like object containing all requested channel data.

        """
        return concatenate_blocks(*self.stream(channels, start, end))

    def _stream_channel_metadata(
        self,
        client: flight.FlightClient,
        descriptor: flight.FlightDescriptor,
    ) -> Generator[Channel, None, None]:
        """stream channel metadata."""
        flight_info = get_flight_info(client, descriptor)
        with MultiFlightReader(flight_info.endpoints, client) as stream:
            for channel_meta in stream.unpack():
                name = channel_meta.pop("channel")
                yield Channel(name, **channel_meta)


def _parse_data_types(
    data_types: DataTypeLike | None,
) -> list[str]:
    """Parse numpy-like data types to be JSON-serializable."""
    if data_types is None:
        return []
    if isinstance(data_types, (str, type, numpy.dtype)):
        return [numpy.dtype(data_types).name]
    return [numpy.dtype(dtype).name for dtype in data_types]


def _endpoint_is_kafka(endpoint):
    """True if endpoint is a kafka endpoint"""
    location = endpoint.locations[0]
    scheme = urlparse(location.uri.decode()).scheme
    return scheme == "kafka"
