# 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

"""Publisher API."""

from __future__ import annotations

import contextlib
import logging
from dataclasses import replace
from typing import TYPE_CHECKING, Literal

import numpy
import pyarrow
from pyarrow.flight import connect

from . import constants, schemas
from .client import Client
from .flight import (
    MultiFlightReader,
    RequestType,
    RequestValidator,
    create_descriptor,
    parse_url,
)

try:
    from confluent_kafka import Producer
except ImportError:
    HAS_KAFKA = False
else:
    HAS_KAFKA = True

if TYPE_CHECKING:
    from collections.abc import Generator, Iterable
    from datetime import timedelta

    from .block import SeriesBlock
    from .channel import Channel


logger = logging.getLogger("arrakis")


def serialize_batch(batch: pyarrow.RecordBatch):
    """Serialize a record batch to bytes.

    Parameters
    ----------
    batch : pyarrow.RecordBatch
        The batch to serialize.

    Returns
    -------
    bytes
        The serialized buffer.

    """
    sink = pyarrow.BufferOutputStream()
    with pyarrow.ipc.new_stream(sink, batch.schema) as writer:
        writer.write_batch(batch)
    return sink.getvalue()


def channel_to_dtype_name(channel: Channel) -> str:
    """Given a channel, return the data type's name."""
    assert isinstance(channel.data_type, numpy.dtype)
    return channel.data_type.name


class Publisher:
    """Publish timeseries data to Arrakis."""

    def __init__(self, publisher_id: str, url: str | None = None):
        """Initialize Publisher.

        Parameters
        ----------
        id : str
            Publisher ID string.
        url : str | None
            Initial Flight URL to connect to.  Will be automatically
            determined if not specified.

        """
        if not HAS_KAFKA:
            msg = (
                "Publishing requires confluent-kafka to be installed."
                "This is provided by the 'publish' extra or it can be "
                "installed manually through pip or conda."
            )
            raise ImportError(msg)

        self.publisher_id = publisher_id
        self.initial_url = parse_url(url)

        self.channels: dict[str, Channel] = {}
        self.stride: int | None = None

        self._producer: Producer
        self._validator = RequestValidator()

    def register(self, channels: Iterable[Channel] | None = None):
        """register channels for publication

        For most publishers, channels are not specified when
        registering and this method will query the server for the
        allowable channels for this publisher and register them
        internally.

        For publishers allowed to register their own channels, all
        channels they expect to publish should be provided as
        argument, and they will be registered with the server, who
        will in turn provide the necessarily partition information for
        the channels to be published into kafka.

        """

        if channels:
            msg = "Channel self-register is not currently supported."
            raise NotImplementedError(msg)

            for name, id_, index in self._partition_channels(channels):
                self.channels[name] = replace(
                    self.channels[name],
                    partition_id=id_,
                    partition_index=index,
                )

            # check that all channels requested channels have been
            # registered
            for channel in channels:
                if channel.name not in self.channels:
                    msg = f"channel {channel.name} was not properly registered."
                    raise ValueError(msg)

        # This find call should raise an error if the publisher_id
        # is not known to the server.
        for channel in Client(self.initial_url).find(publisher=self.publisher_id):
            if self.stride:
                assert channel.stride == self.stride, (
                    "channels specify inconsistent stride"
                )
            else:
                self.stride = channel.stride
            self.channels[channel.name] = channel

        if not self.channels:
            # FIXME: more informative error message here
            msg = f"unknown publisher ID '{self.publisher_id}'."
            raise ValueError(msg)

        for channel in self.channels.values():
            assert channel.partition_id is not None, (
                f"Channel {channel} is missing partition_id."
            )
            assert channel.partition_index is not None, (
                f"Channel {channel} is missing partition_index."
            )

        return self

    def enter(self) -> None:
        """Enter publication context manager"""
        # get connection properties
        producer_info: dict[str, str] = {}
        descriptor = create_descriptor(
            RequestType.Publish,
            publisher_id=self.publisher_id,
            validator=self._validator,
        )
        with connect(self.initial_url) as client:
            flight_info = client.get_flight_info(descriptor)
            with MultiFlightReader(flight_info.endpoints, client) as stream:
                for data in stream.unpack():
                    kv_pairs = data["properties"]
                    producer_info.update(dict(kv_pairs))
        logger.info("producer info: %s", producer_info)

        # set up producer
        self._producer = Producer(
            {
                "message.max.bytes": 10_000_000,  # 10 MB
                "enable.idempotence": True,
                **producer_info,
            }
        )

    def __enter__(self) -> Publisher:
        self.enter()
        return self

    def publish(
        self,
        block: SeriesBlock,
        timeout: timedelta = constants.DEFAULT_TIMEOUT,
    ) -> None:
        """Publish timeseries data

        Parameters
        ----------
        block : SeriesBlock
            A data block with all channels to publish.
        timeout : timedelta, optional
            The maximum time to wait to publish before timing out.
            Default is 2 seconds.

        """
        if not hasattr(self, "_producer") or not self._producer:
            msg = (
                "publication interface not initialized, "
                "please use context manager when publishing."
            )
            raise RuntimeError(msg)

        # check for attempt to publish invalid channels
        for name, channel in block.channels.items():
            if channel != self.channels[name]:
                msg = f"invalid channel for this publisher: {channel}"
                raise ValueError(msg)
        # FIXME: check for invalid block duration
        # FIXME: check for block time in the future
        # FIXME: warning for missing channels

        # publish data for each data type
        for partition_id, batch in block.to_row_batches(self.channels):
            topic = f"arrakis-{partition_id}"
            logger.debug("publishing to topic %s: %s", topic, batch)
            self._producer.produce(topic=topic, value=serialize_batch(batch))
            self._producer.flush()

    def _partition_channels(
        self, channels: Iterable[Channel]
    ) -> Generator[tuple[str, str, int]]:
        descriptor = create_descriptor(
            RequestType.Partition,
            publisher_id=self.publisher_id,
            validator=self._validator,
        )
        # FIXME: should we not get FlightInfo first?
        with connect(self.initial_url) as client:
            writer, reader = client.do_exchange(descriptor)

        # send over list of channels to map new/updated partitions for
        dtypes = [channel_to_dtype_name(channel) for channel in channels]

        schema = schemas.metadata(scope="partition")
        batch = pyarrow.RecordBatch.from_arrays(
            [
                pyarrow.array(
                    [str(channel) for channel in channels],
                    type=schema.field("channel").type,
                ),
                pyarrow.array(dtypes, type=schema.field("data_type").type),
                pyarrow.array(
                    [channel.sample_rate for channel in channels],
                    type=schema.field("sample_rate").type,
                ),
                pyarrow.array(
                    [None for _ in channels],
                    type=schema.field("partition_id").type,
                ),
                pyarrow.array(
                    [None for _ in channels],
                    type=schema.field("partition_index").type,
                ),
            ],
            schema=schema,
        )

        # send over the partitions
        writer.begin(schema)
        writer.write_batch(batch)
        writer.done_writing()

        # get back the partition IDs and update
        partitions = reader.read_all().to_pydict()
        for channel, id_, index in zip(
            partitions["channel"],
            partitions["partition_id"],
            partitions["partition_index"],
        ):
            yield channel, id_, index

    def close(self) -> None:
        """Exit publication context manager."""
        logger.info("closing kafka producer...")
        with contextlib.suppress(Exception):
            self._producer.flush()

    def __exit__(self, *exc) -> Literal[False]:
        self.close()
        return False
