import bisect
import itertools
import time
from collections import Counter
from dataclasses import dataclass
from operator import itemgetter
from typing import List, Union, Tuple, Optional, NamedTuple, Dict, Any

from fakeredis._commands import BeforeAny, AfterAny
from fakeredis._helpers import current_time


class StreamEntryKey(NamedTuple):
    ts: int
    seq: int

    def encode(self) -> bytes:
        return f'{self.ts}-{self.seq}'.encode()

    @staticmethod
    def parse_str(entry_key_str: Union[bytes, str]) -> 'StreamEntryKey':
        if isinstance(entry_key_str, bytes):
            entry_key_str = entry_key_str.decode()
        s = entry_key_str.split('-')
        (timestamp, sequence) = (int(s[0]), 0) if len(s) == 1 else (int(s[0]), int(s[1]))
        return StreamEntryKey(timestamp, sequence)


class StreamRangeTest:
    """Argument converter for sorted set LEX endpoints."""

    def __init__(self, value: Union[StreamEntryKey, BeforeAny, AfterAny], exclusive: bool):
        self.value = value
        self.exclusive = exclusive

    @staticmethod
    def valid_key(entry_key: Union[bytes, str]) -> bool:
        try:
            StreamEntryKey.parse_str(entry_key)
            return True
        except ValueError:
            return False

    @classmethod
    def decode(cls, value: bytes, exclusive=False):
        if value == b'-':
            return cls(BeforeAny(), True)
        elif value == b'+':
            return cls(AfterAny(), True)
        elif value[:1] == b'(':
            return cls(StreamEntryKey.parse_str(value[1:]), True)
        return cls(StreamEntryKey.parse_str(value), exclusive)


@dataclass
class StreamConsumerInfo(object):
    name: bytes
    pending: int
    last_attempt: int  # Impacted by XREADGROUP, XCLAIM, XAUTOCLAIM
    last_success: int  # Impacted by XREADGROUP, XCLAIM, XAUTOCLAIM

    def __init__(self, name):
        self.name = name
        self.pending = 0
        _time = current_time()
        self.last_attempt = _time
        self.last_success = _time

    def info(self, curr_time: int) -> List[bytes]:
        return [
            b'name', self.name,
            b'pending', self.pending,
            b'idle', curr_time - self.last_attempt,
            b'inactive', curr_time - self.last_success,
        ]


class StreamGroup(object):
    def __init__(self, stream: 'XStream', name: bytes, start_key: StreamEntryKey, entries_read: int = None):
        self.stream = stream
        self.name = name
        self.start_key = start_key
        self.entries_read = entries_read
        # consumer_name -> #pending_messages
        self.consumers: Dict[bytes, StreamConsumerInfo] = dict()
        self.last_delivered_key = start_key
        self.last_ack_key = start_key
        # Pending Entries List, see https://redis.io/commands/xreadgroup/
        # msg_id -> consumer_name, time read
        self.pel = dict()

    def set_id(self, last_delivered_str: bytes, entries_read: Union[int, None]) -> None:
        """Set last_delivered_id for group
        """
        self.start_key = self.stream.parse_ts_seq(last_delivered_str)
        start_index, _ = self.stream.find_index(self.start_key)
        self.entries_read = entries_read
        self.last_delivered_key = self.stream[min(start_index + (entries_read or 0), len(self.stream) - 1)]

    def add_consumer(self, consumer_name: bytes) -> int:
        if consumer_name in self.consumers:
            return 0
        self.consumers[consumer_name] = StreamConsumerInfo(consumer_name)
        return 1

    def del_consumer(self, consumer_name: bytes) -> int:
        if consumer_name not in self.consumers:
            return 0
        res = self.consumers[consumer_name].pending
        del self.consumers[consumer_name]
        return res

    def consumers_info(self):
        return [self.consumers[k].info(current_time()) for k in self.consumers]

    def group_info(self) -> List[bytes]:
        start_index, _ = self.stream.find_index(self.start_key)
        last_delivered_index, _ = self.stream.find_index(self.last_delivered_key)
        last_ack_index, _ = self.stream.find_index(self.last_ack_key)
        if start_index + (self.entries_read or 0) > len(self.stream):
            lag = len(self.stream) - start_index - self.entries_read
        else:
            lag = len(self.stream) - 1 - last_delivered_index
        res = {
            b'name': self.name,
            b'consumers': len(self.consumers),
            b'pending': last_delivered_index - last_ack_index,
            b'last-delivered-id': self.last_delivered_key.encode(),
            b'entries-read': self.entries_read,
            b'lag': lag,
        }
        return list(itertools.chain(*res.items()))

    def group_read(self, consumer_name: bytes, start_id: bytes, count: int, noack: bool) -> List:
        _time = current_time()
        if consumer_name not in self.consumers:
            self.consumers[consumer_name] = StreamConsumerInfo(consumer_name)

        self.consumers[consumer_name].last_attempt = _time
        if start_id == b'>':
            start_key = self.last_delivered_key
        else:
            start_key = max(StreamEntryKey.parse_str(start_id), self.last_delivered_key)
        ids_read = self.stream.stream_read(start_key, count)
        if not noack:
            for k in ids_read:
                self.pel[k] = (consumer_name, _time)
        if len(ids_read) > 0:
            self.last_delivered_key = max(self.last_delivered_key, ids_read[-1])
            self.entries_read = (self.entries_read or 0) + len(ids_read)
        self.consumers[consumer_name].last_success = _time
        self.consumers[consumer_name].pending += len(ids_read)
        return [self.stream.format_record(x) for x in ids_read]

    def _calc_consumer_last_time(self):
        new_last_success_map = {k: min(v)[1] for k, v in itertools.groupby(self.pel.values(), key=itemgetter(0))}
        for consumer in new_last_success_map:
            if consumer not in self.consumers:
                self.consumers[consumer] = StreamConsumerInfo(consumer)
            self.consumers[consumer].last_attempt = new_last_success_map[consumer]
            self.consumers[consumer].last_success = new_last_success_map[consumer]

    def ack(self, args: Tuple[bytes]) -> int:
        res = 0
        for k in args:
            try:
                parsed = StreamEntryKey.parse_str(k)
            except Exception:
                continue
            if parsed in self.pel:
                consumer_name = self.pel[parsed][0]
                self.consumers[consumer_name].pending -= 1
                del self.pel[parsed]
                res += 1
        self._calc_consumer_last_time()
        return res

    def pending(self, idle: Optional[int],
                start: Optional[StreamRangeTest],
                end: Optional[StreamRangeTest],
                count: Optional[int],
                consumer: Optional[str]):
        _time = current_time()
        relevent_ids = list(self.pel.keys())
        if consumer is not None:
            relevent_ids = [k for k in relevent_ids
                            if self.pel[k][0] == consumer]
        if idle is not None:
            relevent_ids = [
                k for k in relevent_ids
                if self.pel[k][1] + idle < _time
            ]
        if start is not None and end is not None:
            relevent_ids = [
                k for k in relevent_ids
                if (((start.value < k) or (start.value == k and not start.exclusive))
                    and ((end.value > k) or (end.value == k and not end.exclusive)))
            ]
        if count is not None:
            relevent_ids = sorted(relevent_ids)[:count]

        return [[k.encode(), self.pel[k][0]] for k in relevent_ids]

    def pending_summary(self) -> List[bytes]:
        counter = Counter([self.pel[k][0] for k in self.pel])
        data = [
            len(self.pel),
            min(self.pel).encode() if len(self.pel) > 0 else None,
            max(self.pel).encode() if len(self.pel) > 0 else None,
            [[i, counter[i]] for i in counter],
        ]
        return data

    def claim(self,
              min_idle_ms: int,
              msgs: List[bytes],
              consumer_name: bytes,
              _time: Optional[int],
              force: bool) -> Tuple[List, List]:
        curr_time = current_time()
        if _time is None:
            _time = curr_time
        self.consumers.get(consumer_name, StreamConsumerInfo(consumer_name)).last_attempt = curr_time
        claimed_msgs, deleted_msgs = [], []
        for msg in msgs:
            try:
                key = StreamEntryKey.parse_str(msg) if isinstance(msg, bytes) else msg
            except Exception:
                continue
            if key not in self.pel:
                if force:
                    self.pel[key] = (consumer_name, _time)  # Force claim msg
                    if key in self.stream:
                        claimed_msgs.append(key)
                    else:
                        deleted_msgs.append(key)
                        del self.pel[key]
                continue
            if curr_time - self.pel[key][1] < min_idle_ms:
                continue  # Not idle enough time to be claimed
            self.pel[key] = (consumer_name, _time)
            if key in self.stream:
                claimed_msgs.append(key)
            else:
                deleted_msgs.append(key)
                del self.pel[key]
        self._calc_consumer_last_time()
        return sorted(claimed_msgs), sorted(deleted_msgs)

    def read_pel_msgs(self, min_idle_ms: int, start: bytes, count: int):
        start_key = StreamEntryKey.parse_str(start)
        curr_time = current_time()
        msgs = sorted([k for k in self.pel
                       if (curr_time - self.pel[k][1] >= min_idle_ms)
                       and k >= start_key])
        count = min(count, len(msgs))
        return msgs[:count]


class XStream:
    """Class representing stream.

    The stream contains entries with keys (timestamp, sequence) and field->value pairs.
    This implementation has them as a sorted list of tuples, the first value in the tuple
    is the key (timestamp, sequence).

    Structure of _values list:
    [
       ((timestamp,sequence), [field1, value1, field2, value2, ...])
       ((timestamp,sequence), [field1, value1, field2, value2, ...])
    ]
    """

    def __init__(self):
        self._ids: List[StreamEntryKey] = list()
        self._values_dict: Dict[StreamEntryKey, List] = dict()
        self._groups: Dict[bytes, StreamGroup] = dict()
        self._max_deleted_id = StreamEntryKey(0, 0)
        self._entries_added = 0

    def group_get(self, group_name: bytes) -> Optional[StreamGroup]:
        return self._groups.get(group_name, None)

    def group_add(self, name: bytes, start_key_str: bytes, entries_read: Union[int, None]) -> None:
        """Add a group listening to stream

        :param name: group name
        :param start_key_str: start_key in `timestamp-sequence` format, or $ listen from last.
        :param entries_read: number of entries read.
        """
        if start_key_str == b'$':
            start_key = self._ids[-1] if len(self._ids) > 0 else StreamEntryKey(0, 0)
        else:
            start_key = StreamEntryKey.parse_str(start_key_str)
        self._groups[name] = StreamGroup(self, name, start_key, entries_read)

    def group_delete(self, group_name: bytes) -> int:
        if group_name in self._groups:
            del self._groups[group_name]
            return 1
        return 0

    def groups_info(self) -> List[List[bytes]]:
        res = []
        for group in self._groups.values():
            group_res = group.group_info()
            res.append(group_res)
        return res

    def stream_info(self, full: bool) -> List[bytes]:
        res = {
            b'length': len(self._ids),
            b'groups': len(self._groups),
            b'first-entry': self.format_record(self._ids[0]) if len(self._ids) > 0 else None,
            b'last-entry': self.format_record(self._ids[-1]) if len(self._ids) > 0 else None,
            b'max-deleted-entry-id': self._max_deleted_id.encode(),
            b'entries-added': self._entries_added,
            b'recorded-first-entry-id': self._ids[0].encode() if len(self._ids) > 0 else b'0-0',
        }
        if full:
            res[b'entries'] = [self.format_record(i) for i in self._ids]
            res[b'groups'] = [g.group_info() for g in self._groups.values()]
        return list(itertools.chain(*res.items()))

    def delete(self, lst: List[Union[str, bytes]]) -> int:
        """Delete items from stream

        :param lst: list of IDs to delete, in the form of `timestamp-sequence`.
        :returns: Number of items deleted
        """
        res = 0
        for item in lst:
            ind, found = self.find_index_key_as_str(item)
            if found:
                self._max_deleted_id = max(self._ids[ind], self._max_deleted_id)
                del self._values_dict[self._ids[ind]]
                del self._ids[ind]
                res += 1
        return res

    def add(self, fields: List, entry_key: str = '*') -> Union[None, bytes]:
        """Add entry to a stream.

        If the entry_key can not be added (because its timestamp is before the last entry, etc.),
        nothing is added.

        :param fields: list of fields to add, must [key1, value1, key2, value2, ... ]
        :param entry_key:
            key for the entry, formatted as 'timestamp-sequence'
            If entry_key is '*', the timestamp will be calculated as current time and the sequence based
            on the last entry key of the stream.
            If entry_key is 'ts-*', and the timestamp is greater or equal than the last entry timestamp,
            then the sequence will be calculated accordingly.
        :returns:
            The key of the added entry.
            None if nothing was added.
        :raises AssertionError: if len(fields) is not even.
        """
        assert len(fields) % 2 == 0
        if isinstance(entry_key, bytes):
            entry_key = entry_key.decode()

        if entry_key is None or entry_key == '*':
            ts, seq = int(1000 * time.time()), 0
            if (len(self._ids) > 0
                    and self._ids[-1].ts == ts
                    and self._ids[-1].seq >= seq):
                seq = self._ids[-1].seq + 1
            ts_seq = StreamEntryKey(ts, seq)
        elif entry_key[-1] == '*':  # entry_key has `timestamp-*` structure
            split = entry_key.split('-')
            if len(split) != 2:
                return None
            ts, seq = int(split[0]), split[1]
            if len(self._ids) > 0 and ts == self._ids[-1].ts:
                seq = self._ids[-1].seq + 1
            else:
                seq = 0
            ts_seq = StreamEntryKey(ts, seq)
        else:
            ts_seq = StreamEntryKey.parse_str(entry_key)

        if len(self._ids) > 0 and self._ids[-1] > ts_seq:
            return None
        self._ids.append(ts_seq)
        self._values_dict[ts_seq] = list(fields)
        self._entries_added += 1
        return ts_seq.encode()

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

    def __iter__(self):
        def gen():
            for k in self._ids:
                yield self.format_record(k)

        return gen()

    def __getitem__(self, key):
        if isinstance(key, int):  # Index
            return self._ids[key]
        if isinstance(key, bytes):
            return self._values_dict[StreamEntryKey.parse_str(key)]
        return None

    def __contains__(self, key: StreamEntryKey):
        return key in self._values_dict

    def find_index(self, entry_key: StreamEntryKey, from_left=True) -> Tuple[int, bool]:
        """Find the closest index to entry_key_str in the stream
        :param entry_key: key for the entry.
        :param from_left: if not found exact match, return index of last smaller element
        :returns: A tuple of
            ( index of entry with the closest (from the left) key to entry_key_str,
              Whether the entry key is equal )
        """
        if len(self._ids) == 0:
            return 0, False
        if from_left:
            ind = bisect.bisect_left(self._ids, entry_key)
        else:
            ind = bisect.bisect_right(self._ids, entry_key)
        return ind, (ind < len(self._ids) and self._ids[ind] == entry_key)

    def find_index_key_as_str(self, entry_key_str: Union[str, bytes]) -> Tuple[int, bool]:
        """Find the closest index to entry_key_str in the stream
        :param entry_key_str: key for the entry, formatted as 'timestamp-sequence'.
        :returns: A tuple of
            ( index of entry with the closest (from the left) key to entry_key_str,
              Whether the entry key is equal )
        """
        if entry_key_str == b'$':
            return max(len(self._ids) - 1, 0), True
        ts_seq = StreamEntryKey.parse_str(entry_key_str)
        return self.find_index(ts_seq)

    @staticmethod
    def parse_ts_seq(ts_seq_str: Union[str, bytes]) -> StreamEntryKey:
        if ts_seq_str == b'$':
            return StreamEntryKey(0, 0)
        return StreamEntryKey.parse_str(ts_seq_str)

    def trim(self,
             max_length: Optional[int] = None,
             start_entry_key: Optional[str] = None,
             limit: Optional[int] = None) -> int:
        """Trim a stream

        :param max_length: max length of resulting stream after trimming (number of last values to keep)
        :param start_entry_key: min entry-key to keep, can not be given together with max_length.
        :param limit: number of entries to keep from minid.
        :returns: The resulting stream after trimming.
        :raises ValueError: When both max_length and start_entry_key are passed.
        """
        if max_length is not None and start_entry_key is not None:
            raise ValueError('Can not use both max_length and start_entry_key')
        start_ind = None
        if max_length is not None:
            start_ind = len(self._ids) - max_length
        elif start_entry_key is not None:
            ind, exact = self.find_index_key_as_str(start_entry_key)
            start_ind = ind
        res = max(start_ind, 0)
        if limit is not None:
            res = min(start_ind, limit)
        remove_keys, self._ids = self._ids[:res], self._ids[res:]
        for k in remove_keys:
            del self._values_dict[k]
        return res

    def irange(self, start: StreamRangeTest, stop: StreamRangeTest, reverse=False) -> List[Any]:
        """Returns a range of the stream from start to stop.

        :param start: start key
        :param stop: stop key
        :param reverse: Should the range be in reverse order?
        :returns: the range between start and stop
        """

        def _find_index(elem: StreamRangeTest, from_left=True) -> int:
            if isinstance(elem.value, BeforeAny):
                return 0
            if isinstance(elem.value, AfterAny):
                return len(self._ids)
            ind, found = self.find_index(elem.value, from_left)
            if found and elem.exclusive:
                ind += 1
            return ind

        start_ind = _find_index(start)
        stop_ind = _find_index(stop, from_left=False)
        matches = map(lambda x: self.format_record(self._ids[x]), range(start_ind, stop_ind))
        if reverse:
            return list(reversed(tuple(matches)))
        return list(matches)

    def last_item_key(self) -> bytes:
        return self._ids[-1].encode() if len(self._ids) > 0 else '0-0'.encode()

    def stream_read(self, start_key: StreamEntryKey, count: Union[int, None]) -> List[StreamEntryKey]:
        start_ind, found = self.find_index(start_key)
        if found:
            start_ind += 1
        if start_ind >= len(self):
            return []
        end_ind = len(self) if count is None or start_ind + count >= len(self) else start_ind + count
        return self._ids[start_ind:end_ind]

    def format_record(self, key: StreamEntryKey) -> List:
        results = self._values_dict[key]
        return [key.encode(), results]
