import logging
import random
import string
from collections import defaultdict
from collections.abc import Iterator
from typing import Generator

import pyarrow
from confluent_kafka import Consumer, TopicPartition
from gpstime import gps2unix, gpsnow
from pyarrow import RecordBatch

from .block import SeriesBlock, Time, time_as_ns
from .channel import Channel
from .stream import StreamReader

logger = logging.getLogger("arrakis")


class KafkaReader(StreamReader):
    """A connection object to read data from Kafka.

    Parameters
    ----------
    url : str
        URL of Kafka broker to connect to.
    metadata : dict[str, Channel]]
        Dictionary of channel metadata for request.
    start : int | None
        GPS start time of stream in nanoseconds, defaults to "now".

    """

    # how frequently the kafka consumer polls the broker for new
    # messages, in seconds.  Needs to be long enough to give the
    # server time to respond (50-100ms).
    POLL_TIMEOUT = 0.1

    def __init__(
        self,
        url: str,
        metadata: dict[str, Channel],
        start: int | None = None,
    ):
        self.url = url
        self.metadata = metadata
        self.start = start

        # track stream -> index -> channel
        self._name_lookup: dict[str, dict[int, Channel]] = defaultdict(defaultdict)
        # filters for partitions
        self._partition_filter_id_sets = defaultdict(set)

        for channel in self.metadata.values():
            partition_id = channel.partition_id
            assert partition_id is not None
            partition_index = channel.partition_index
            assert partition_index is not None
            stream_id = self._id(partition_id)
            self._name_lookup[stream_id][partition_index] = channel
            self._partition_filter_id_sets[partition_id].add(partition_index)

        # pre-compute Arrow arrays for filtering to avoid repeated
        # array creation
        self._partition_filter_id_arrays = {
            partition_id: pyarrow.array(channel_id, type=pyarrow.int32())
            for partition_id, channel_id in self._partition_filter_id_sets.items()
        }

        # Optimize IPC stream reader creation/destruction overhead
        self._ipc_options = pyarrow.ipc.IpcReadOptions(use_threads=False)

        # create Kafka consumer
        consumer_settings = {
            "bootstrap.servers": url,
            "group.id": generate_groupid(),
            "message.max.bytes": 10_000_000,  # 10 MB
            "enable.auto.commit": False,
        }
        self._consumer = Consumer(consumer_settings)
        logger.debug("Kafka consumer: %s", consumer_settings)
        self._topics = [
            f"arrakis-{partition_id}" for partition_id in self._partition_filter_id_sets
        ]
        logger.debug("kafka topics: %s", self._topics)

    def _id(self, partition_id: str) -> str:
        """stream id for partition"""
        return f"{self.url}/{partition_id}"

    @property
    def streams(self) -> dict[str, list[Channel]]:
        return {
            stream_id: list(index_chan_map.values())
            for stream_id, index_chan_map in self._name_lookup.items()
        }

    def _smart_filter_batch(
        self, batch: pyarrow.RecordBatch, partition_id: str
    ) -> pyarrow.RecordBatch:
        """Adaptive filtering with bailout for cases where no filtering is needed."""
        id_column = batch.column("id")

        # get unique channels in this batch
        unique_batch_channels = pyarrow.compute.unique(id_column)
        unique_batch_set = set(unique_batch_channels.to_pylist())

        allowed_ids = self._partition_filter_id_sets[partition_id]

        # check if no unrelated channels are contained in batch
        if unique_batch_set.issubset(allowed_ids):
            return batch

        # keep only rows with requested channels
        mask = pyarrow.compute.is_in(
            id_column,
            self._partition_filter_id_arrays[partition_id],
        )
        return pyarrow.compute.filter(batch, mask)

    def enter(self):
        logger.debug("initiating kafka subscriptions...")
        now = time_as_ns(gpsnow())
        # if start time is specified, adjust the consumer's
        # offset to point at the data requested
        if self.start and self.start <= now:
            # convert to UNIX time in ms
            offset_time = int(gps2unix(self.start // Time.SECONDS) * 1000)
            # get offsets corresponding to times
            partitions = [
                TopicPartition(topic, partition=0, offset=offset_time)
                for topic in self._topics
            ]
            partitions = self._consumer.offsets_for_times(partitions)
            # reassign topic partitions to consumer
            self._consumer.unsubscribe()
            self._consumer.assign(partitions)

        # FIXME: start times in the future are being handled as if a
        # start time was not specified
        else:
            self._consumer.subscribe(self._topics)

    def read(
        self, *, convert_blocks=False
    ) -> Generator[tuple[str, RecordBatch | SeriesBlock], None, None]:
        while True:
            msg = self._consumer.poll(self.POLL_TIMEOUT)

            if not msg:
                break

            if msg.error():
                logger.error("KAFKA CONSUME ERROR: %s", msg.error())
                raise RuntimeError(msg.error())

            # extract partition ID from message
            partition_id = msg.topic().split("-", 1)[1]
            stream_id = self._id(partition_id)

            # deserialize streams
            with pyarrow.ipc.open_stream(
                msg.value(),
                options=self._ipc_options,
            ) as reader:
                for batch_input in read_all_batches(reader):
                    batch = self._smart_filter_batch(batch_input, partition_id)
                    if convert_blocks:
                        data = SeriesBlock.from_row_batch(
                            batch,
                            self._name_lookup[stream_id],
                        )
                    else:
                        data = batch
                    yield stream_id, data

    def close(self) -> None:
        logger.debug("closing kafka subscriptions...")
        self._consumer.unassign()
        self._consumer.unsubscribe()
        self._consumer.close()


def generate_groupid() -> str:
    """Generate a random Kafka group ID."""
    return random_alphanum(16)


def random_alphanum(n: int) -> str:
    """Generate a random alpha-numeric sequence of N characters."""
    alphanum = string.ascii_uppercase + string.digits
    return "".join(random.SystemRandom().choice(alphanum) for _ in range(n))


def read_all_batches(
    reader: pyarrow.ipc.RecordBatchStreamReader,
) -> Iterator[pyarrow.RecordBatch]:
    while True:
        try:
            yield reader.read_next_batch()
        except StopIteration:
            return
