# 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 heapq
import logging
import math
import os
import time
import warnings
from collections import defaultdict, deque
from collections.abc import Iterable, Iterator, Mapping
from dataclasses import dataclass
from datetime import timedelta
from enum import Enum, auto
from threading import RLock
from typing import TYPE_CHECKING, Any, Generator, Generic, TypeVar

import numpy
from gpstime import gpsnow

from .block import SeriesBlock, Time, combine_blocks, concatenate_blocks, time_as_ns
from .stream import StreamReader

if TYPE_CHECKING:
    from .channel import Channel


logger = logging.getLogger("arrakis")

T = TypeVar("T")

DEFAULT_TIMEOUT = timedelta(seconds=1)


class OnDrop(Enum):
    IGNORE = auto()
    RAISE = auto()
    WARN = auto()


ONDROP_DEFAULT = os.getenv("ARRAKIS_ONDROP", "warn")


@dataclass
class MuxedData(Mapping, Generic[T]):
    """Container that holds timestamped data.

    Parameters
    ----------
    time : int
        The timestamp associated with this data, in nanoseconds.
    data : dict[str, T]
        The keyed data.

    """

    time: int
    data: dict[str, T]

    def __getitem__(self, index: str) -> T:
        return self.data[index]

    def __iter__(self) -> Iterator[str]:
        return iter(self.data)

    def __len__(self) -> int:
        return len(self.data)


class Muxer(Generic[T]):
    """A data structure that multiplexes items from multiple named streams.

    Given items from multiple named streams with monotonically increasing
    integer timestamps, this data structure can be used to pull out sets of
    synchronized items (items all with the same timestamp).

    The oldest items will be held until either all named streams are
    available or until the timeout has been reached. If a start time has been
    set, any items with an older timestamp will be rejected.

    Parameters
    ----------
    keys : Iterable[str]
        Identifiers for the named streams to expect when adding items.
    start : int, optional
        The GPS time to start muxing items for.
        If not set, accept items from any time.
    timeout : timedelta or None, optional
        The maximum time to wait for messages from named streams, in seconds,
        before multiplexing. If None is specified, wait indefinitely. Default
        is 1 second.

    """

    def __init__(
        self,
        keys: Iterable[str],
        start: int | None = None,
        timeout: timedelta | None = DEFAULT_TIMEOUT,
    ) -> None:
        self._keys = set(keys)
        self._items: dict[int, dict[str, T]] = defaultdict(lambda: defaultdict())
        self._times: list[int] = []
        self._last_time = (
            (start - 1) if start is not None else numpy.iinfo(numpy.int64).min
        )
        self._start = start
        self._timeout = timeout

        # track when processing started to handle lookback properly
        self._processing_start_time = int(gpsnow() * Time.SECONDS)

    def push(self, time: int, key: str, item: T, on_drop: str = ONDROP_DEFAULT) -> None:
        """Push an item into the muxer.

        Parameters
        ----------
        time : int
            The timestamp associated with this item.
        key : str
            The key stream associated with this item. Must match a key provided
            at initialization.
        item : T
            The item to add.
        on_drop : str, optional
            Specifies behavior when the item would be dropped from the muxer,
            in the case that it was not provided to the muxer before the
            specified timeout. Options are 'ignore', 'raise', or 'warn'.
            Default is 'warn'.

        """
        if key not in self._keys:
            msg = f"{key} doesn't match keys provided at initialization"
            raise KeyError(msg)

        # skip over items that have already been pulled
        if time <= self._last_time:
            if self._start is not None and time < self._start:
                return
            msg = f"item's timestamp is too old: ({time} <= {self._last_time})"
            match OnDrop[on_drop.upper()]:
                case OnDrop.IGNORE:
                    return
                case OnDrop.RAISE:
                    raise ValueError(msg)
                case OnDrop.WARN:
                    logger.warning(msg)
                    warnings.warn(msg, stacklevel=2)
                    return

        # add item
        if time in self._items:
            if key not in self._items[time]:
                self._items[time][key] = item
        else:
            heapq.heappush(self._times, time)
            self._items[time][key] = item

    def pull(self) -> Iterator[MuxedData[T]]:
        """Pull monotonically increasing synchronized items from the muxer.

        Yields
        ------
        MuxedData[T]
            Synchronized items with a common timestamp, keyed by stream keys.

        """
        if not self._times:
            return

        # yield items in monotonically increasing order as long
        # as conditions are met
        time = self._times[0]
        while (
            self._has_all_items(time)
            or self._are_items_stale(time)
            or self._has_complete_newer_timestamp(time)
        ):
            yield MuxedData(time, self._items.pop(time))
            self._last_time = heapq.heappop(self._times)
            if not self._times:
                break
            time = self._times[0]

    def _has_all_items(self, time: int):
        """Check if a timestamp has all items requested."""
        return len(self._items[time]) == len(self._keys)

    def _has_complete_newer_timestamp(self, time: int):
        """Check if there's a newer complete timestamp, making this one safe to yield.

        Based on monotonic timestamp assumption: if we have complete data for a newer
        timestamp, no more data will arrive for older timestamps, so we can safely
        yield partial older data.
        """
        for newer_time in self._times:
            if newer_time > time and self._has_all_items(newer_time):
                return True
        return False

    def _are_items_stale(self, time):
        """Check if a timestamp is older than the latency cutoff."""
        if self._timeout is None:
            return False

        time_now = gpsnow()
        dt_timeout = self._timeout.total_seconds()

        if time < self._processing_start_time:
            # historical data: give extra time based on how far back in history
            dt_lookback = (self._processing_start_time - time) / float(Time.SECONDS)
            oldest_time_allowed = time_now - dt_timeout + dt_lookback
        else:
            # live data: use normal timeout
            oldest_time_allowed = time_now - dt_timeout

        return time <= int(oldest_time_allowed * Time.SECONDS)


class TimedQueue(Generic[T]):
    """A sequential, time-stamped queue handling gaps and timeouts

    Parameters
    ----------
    start : int
        GPS start time of queue, in nanoseconds.
    stride : int
        Time step for elements in the queue, in nanoseconds.
    timeout : int
        Timeout for elements in the queue, in nanoseconds.
    """

    def __init__(self, stride: int, timeout: int, start: int | None = None):
        self.stride = stride
        self.timeout = timeout
        self.start = start
        # Track if queue has received data yet (only relevant when start is None)
        self._initialized = start is not None
        # set the initial last_time attribute to be the first multiple
        # of stride before the requested start time
        if start is None:
            start = time_as_ns(gpsnow())
        self.last_time = ((start // self.stride) - 1) * self.stride
        self._queue: deque = deque()
        # FIXME: what's the right value here?
        # self._max_queue = int(2 * self.timeout / self.stride)
        self._max_queue = 1_000_000
        self._lock = RLock()

    def __repr__(self) -> str:
        return (
            f"<{self.__class__.__name__} stride={self.stride:_} "
            f"timeout={self.timeout:_} last_time={self.last_time:_} "
            f"length={len(self._queue)}>"
        )

    def __len__(self):
        return len(self._queue)

    def _append(self, *args):
        assert len(self) <= self._max_queue, (
            f"queue is too long: {len(self)} > {self._max_queue}"
        )
        self._queue.append(args)

    def push(self, time: int, element: Any, on_drop: str = ONDROP_DEFAULT) -> None:
        """Push an element into the queue.

        The time being pushed into the queue must be a multiple of the
        time stride specified at initialization of the queue.

        If the time associated with the pushed element is less than
        what has already been processed by the queue, the batch will
        be dropped and this operation will be a no-op.

        Parameters
        ----------
        time : int
            GPS time associated with the element, in nanoseconds
        element : Any
            element being pushed into the queue.
        on_drop : str, optional
            Specifies behavior when the item would be dropped from the muxer,
            in the case that it was not provided to the muxer before the
            specified timeout. Options are 'ignore', 'raise', or 'warn'.
            Default is 'warn'.

        """
        assert time % self.stride == 0, (
            f"time {time} is not a multiple of queue stride {self.stride:_}"
        )
        # If this is the first real block (not a gap), reset last_time to align
        # with it if start time was not set initially. This mitigates the
        # amount of data to be dropped while still allowing gaps to be returned
        # in a timely manner if no data has ever been received.
        if not self._initialized and element is not None:
            self.last_time = ((time // self.stride) - 1) * self.stride
            self._initialized = True

        # if time is older that what's already in the queue, drop it
        if time <= self.last_time:
            msg = f"item's timestamp is too old: ({time:_} <= {self.last_time:_})"
            match OnDrop[on_drop.upper()]:
                case OnDrop.IGNORE:
                    return
                case OnDrop.RAISE:
                    raise ValueError(msg)
                case OnDrop.WARN:
                    logger.warning(msg)
                    warnings.warn(msg, stacklevel=2)
                    return
        # use a lock so that the last_time attribute is always
        # synchronous with element append
        with self._lock:
            # backfill any skipped times with None elements (gaps)
            # FIXME: could be deque.extend?
            for t in range(self.last_time, time, self.stride)[1:]:
                self._append(t, None)
            # insert the new element
            self._append(time, element)
            self.last_time = time

    def _update_timeout(self):
        # only do timeouts for live streams (i.e. start is None)
        if self.start is not None:
            return
        # we handle timeouts by always attempting to push a gap at the
        # timeout time.  either the timeout time is before the last
        # queued time and the gap will be dropped and ignored, or it's
        # ahead of the last queued time, indicating that the queue did
        # not receive elements before the timeout and everything
        # between the last queued time and the timeout time should be
        # filled with gaps.
        timeout = time_as_ns(gpsnow()) - self.timeout
        # round down to the previous stride
        timeout = (timeout // self.stride) * self.stride
        self.push(timeout, None, on_drop="ignore")

    def ready(self, duration: int) -> int | None:
        """Check if queue holds duration worth of elements

        Returns
        -------
        int or None
            Returns either the number of elements that span duration,
            or None if the queue does not have duration worth of
            elements,

        """
        self._update_timeout()
        assert duration % self.stride == 0, (
            f"duration {duration:_} is not a multiple of queue stride {self.stride:_}"
        )
        n_elements = int(duration / self.stride)
        if len(self) >= n_elements:
            return n_elements
        return None

    def pull(self, duration: int | None = None) -> Iterator[tuple[int, Any]]:
        """Drain the queue.

        Gaps are represented by None elements.

        Parameters
        ----------
        duration : int
            Duration to extract, in nanoseconds.  If the specified
            duration is not available, no elements will be returned.
            If not specified, the queue will be drained.

        Yields
        ------
        tuple[time, element]
            The element from the queue and it's associated timestamp.

        """
        self._update_timeout()

        with self._lock:
            # if duration specified, pull the requested number of elements
            if duration:
                if n_elements := self.ready(duration):
                    for _ in range(n_elements):
                        yield self._queue.popleft()

            # else drain the queue
            else:
                yield from self._queue
                self._queue.clear()

    def front_time(self) -> int | None:
        """Return the timestamp of the front element, or None if empty."""
        with self._lock:
            if len(self._queue) > 0:
                return self._queue[0][0]
            return None

    def drain_until(self, target_time: int) -> None:
        """Remove elements until the front element is at or after target_time."""
        with self._lock:
            while len(self._queue) > 0 and self._queue[0][0] < target_time:
                self._queue.popleft()


class BlockMuxStream(Generic[T]):
    """A time-aware, gap-handling multiplexer for SeriesBlock streams.

    Given SeriesBlocks from multiple named streams with monotonically
    increasing integer timestamps, this data structure can be used to
    pull out sets of synchronized blocks, all with the same
    timestamps.  If data on the streams is not available before
    timeouts are reached, gap blocks will be returned.

    The oldest items will be held until either all named streams are
    available or until the timeout has been reached. If a start time
    has been set, any items with an older timestamp will be rejected.

    Parameters
    ----------
    reader : StreamReader
        StreamReader object producing multiple stream to multiplex.
    start_ns : int
        GPS start time of stream, in nanoseconds.
    timeout : int = 0
        Overall timeout for the muxer, in nanoseconds.  Overrides
        individual queue timeouts.  If not specified the mux
        timeout will be the max of the individual queue timeouts.

    """

    def __init__(
        self,
        reader: StreamReader,
        start: int | None = None,
        timeout: int = 0,
    ):
        self.reader = reader
        self.start = start
        # extract stream info
        self._queues: dict[str, TimedQueue] = {}
        self.channels: list[Channel] = []
        for stream_name, channels in self.reader.streams.items():
            # gather stride and latency
            # done this way to satisfy type checking consistency
            strides = []
            timeouts = []
            for channel in channels:
                assert channel.stride
                strides.append(channel.stride)
                assert channel.max_latency
                timeouts.append(channel.max_latency)
            # the stride for a stream is the least common multiple of
            # the stride of the individual channels
            qstride = math.lcm(*strides)
            # timeout is max of all expected latencies
            qtimeout = max(timeouts)
            self._queues[stream_name] = TimedQueue(
                stride=qstride,
                # mux timeout overrides individual timesouts
                timeout=timeout or qtimeout,
                start=start,
            )
            self.channels.extend(channels)
        # the overall stride for multiple streams is the least common
        # multiple of the stride of the individual streams
        self.stride = math.lcm(*(q.stride for q in self._queues.values()))
        self.timeout = max(q.timeout for q in self._queues.values())

    def __getitem__(self, key: str) -> TimedQueue:
        """Access an individual queue."""
        return self._queues[key]

    def push(self, stream_name: str, block: SeriesBlock):
        """Push an element for time into a particular queue."""
        time_ns = block.time_ns
        q = self._queues[stream_name]
        q.push(time_ns, block)

    def ready(self) -> bool:
        """True if all queues have the expected number of elements

        covering the overall muxer stride.

        """
        # Align queues first, then check if all have enough elements
        self._align_queues()
        return all(q.ready(self.stride) for q in self._queues.values())

    def _align_queues(self) -> None:
        """Align all queues to the same starting timestamp.

        If queues have different front timestamps (due to different
        initialization times when start=None), drain the behind queues
        until they catch up to the maximum timestamp.
        """
        front_times = [q.front_time() for q in self._queues.values()]
        valid_times = [t for t in front_times if t is not None]
        if not valid_times:
            return
        target_time = max(valid_times)
        for q in self._queues.values():
            q.drain_until(target_time)

    def pull(self) -> SeriesBlock:
        """Pull synchronized, concatenated, combined blocks from all
        streams covering the overall specified stride.

        Parameters
        ----------
        channels : Iterable[Channel]
            Channel list for creating gap blocks

        Returns
        -------
        SeriesBlock

        """
        blocks = []
        for queue in self._queues.values():
            q_blocks = []
            for time_ns, block in queue.pull(self.stride):
                if block is None:
                    block = SeriesBlock.full_gap(
                        time_ns,
                        self.stride,
                        self.channels,
                    )
                q_blocks.append(block)
            q_block = concatenate_blocks(*q_blocks)
            blocks.append(q_block)
        return combine_blocks(*blocks)

    def stream(self, end_ns: int | None = None) -> Generator[SeriesBlock, None, None]:
        with self.reader as reader:
            while True:
                if not reader.done():
                    for stream_id, block in reader.read(convert_blocks=True):
                        self.push(stream_id, block)

                    if not self.ready():
                        # FIXME: what's the optimal sleep time here, needs
                        # to be less than the smallest stride
                        time.sleep(1 / 64)
                        continue

                block = self.pull()

                if end_ns is not None and block.time_ns > end_ns:
                    return

                yield block

                if end_ns is not None and block.end_ns >= end_ns:
                    return
