import time
from typing import List, Union, Optional, Any, Set, Dict

from fakeredis import _msgs as msgs
from fakeredis._command_args_parsing import extract_args
from fakeredis._commands import command, Key, CommandItem, Int, Float, Timestamp
from fakeredis._helpers import Database, SimpleString, OK, SimpleError, casematch
from fakeredis.model import TimeSeries, TimeSeriesRule, AGGREGATORS


class TimeSeriesCommandsMixin:  # TimeSeries commands
    _timeseries_keys: Set[bytes] = set()
    DUPLICATE_POLICIES = [b"BLOCK", b"FIRST", b"LAST", b"MIN", b"MAX", b"SUM"]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._db: Database

    @staticmethod
    def _filter_expression_check(ts: TimeSeries, filter_expression: bytes) -> bool:
        if not filter_expression:
            return True
        if filter_expression.find(b"!=") != -1:
            if len(filter_expression.split(b"!=")) != 2:
                raise SimpleError(msgs.TIMESERIES_BAD_FILTER_EXPRESSION)
            label, value = filter_expression.split(b"!=")
            if value == "-":
                return label in ts.labels

            if value[0] == b"(" and value[-1] == b")":
                values = set(value[1:-1].split(b","))
                return label in ts.labels and ts.labels[label] not in values
            return label not in ts.labels or ts.labels[label] != value
        if filter_expression.find(b"=") != -1:
            if len(filter_expression.split(b"=")) != 2:
                raise SimpleError(msgs.TIMESERIES_BAD_FILTER_EXPRESSION)
            label, value = filter_expression.split(b"=")
            if value == "-":
                return label not in ts.labels
            if value[0] == b"(" and value[-1] == b")":
                values = set(value[1:-1].split(b","))
                return label in ts.labels and ts.labels[label] in values
            return label in ts.labels and ts.labels[label] == value
        raise SimpleError(msgs.TIMESERIES_BAD_FILTER_EXPRESSION)

    def _get_timeseries(self, filter_expressions: List[bytes]) -> List["TimeSeries"]:
        res: List["TimeSeries"] = list()
        TimeSeriesCommandsMixin._timeseries_keys = {
            k for k in TimeSeriesCommandsMixin._timeseries_keys if k in self._db
        }
        for ts_key in TimeSeriesCommandsMixin._timeseries_keys:
            ts = self._db.get(ts_key).value
            if all([self._filter_expression_check(ts, expr) for expr in filter_expressions]):
                res.append(ts)
        return res

    @staticmethod
    def _validate_duplicate_policy(duplicate_policy: bytes) -> bool:
        return duplicate_policy is None or any(
            [casematch(duplicate_policy, item) for item in TimeSeriesCommandsMixin.DUPLICATE_POLICIES]
        )

    def _create_timeseries(self, name: bytes, *args) -> TimeSeries:
        (retention, encoding, chunk_size, duplicate_policy, (ignore_max_time_diff, ignore_max_val_diff)), left_args = (
            extract_args(
                args,
                ("+retention", "*encoding", "+chunk_size", "*duplicate_policy", "++ignore"),
                error_on_unexpected=False,
            )
        )
        retention = retention or 0
        encoding = encoding or b"COMPRESSED"
        if not (casematch(encoding, b"COMPRESSED") or casematch(encoding, b"UNCOMPRESSED")):
            raise SimpleError(msgs.BAD_SUBCOMMAND_MSG.format("TS.CREATE"))
        encoding = encoding.lower()
        chunk_size = chunk_size or 4096
        if chunk_size % 8 != 0:
            raise SimpleError(msgs.TIMESERIES_BAD_CHUNK_SIZE)
        if not self._validate_duplicate_policy(duplicate_policy):
            raise SimpleError(msgs.TIMESERIES_INVALID_DUPLICATE_POLICY)
        duplicate_policy = duplicate_policy.lower() if duplicate_policy else None
        if len(left_args) > 0 and (not casematch(left_args[0], b"LABELS") or len(left_args) % 2 != 1):
            raise SimpleError(msgs.BAD_SUBCOMMAND_MSG.format("TS.ADD"))
        labels = dict(zip(left_args[1::2], left_args[2::2])) if len(left_args) > 0 else {}

        if duplicate_policy is None and self.version >= (8,):
            # In Redis 8.0, the default duplicate policy is BLOCK
            duplicate_policy = b"block"
        res = TimeSeries(
            name=name,
            database=self._db,
            retention=retention,
            encoding=encoding,
            chunk_size=chunk_size,
            duplicate_policy=duplicate_policy,
            ignore_max_time_diff=ignore_max_time_diff,
            ignore_max_val_diff=ignore_max_val_diff,
            labels=labels,
        )
        self._timeseries_keys.add(name)
        return res

    @command(name="TS.INFO", fixed=(Key(TimeSeries),), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE)
    def ts_info(self, key: CommandItem, *args: bytes) -> List[Any]:
        if key.value is None:
            raise SimpleError(msgs.TIMESERIES_KEY_DOES_NOT_EXIST)
        return [
            b"totalSamples",
            len(key.value.sorted_list),
            b"memoryUsage",
            len(key.value.sorted_list) * 8 + len(key.value.encoding),
            b"firstTimestamp",
            key.value.sorted_list[0][0] if len(key.value.sorted_list) > 0 else 0,
            b"lastTimestamp",
            key.value.sorted_list[-1][0] if len(key.value.sorted_list) > 0 else 0,
            b"retentionTime",
            key.value.retention,
            b"chunkCount",
            len(key.value.sorted_list) * 8 // key.value.chunk_size,
            b"chunkSize",
            key.value.chunk_size,
            b"chunkType",
            key.value.encoding,
            b"duplicatePolicy",
            key.value.duplicate_policy,
            b"labels",
            [[k, v] for k, v in key.value.labels.items()],
            b"sourceKey",
            key.value.source_key,
            b"rules",
            [
                [rule.dest_key.name, rule.bucket_duration, rule.aggregator.upper(), rule.align_timestamp]
                for rule in key.value.rules
            ],
            b"keySelfName",
            key.value.name,
            b"Chunks",
            [],
        ]

    @command(name="TS.CREATE", fixed=(Key(TimeSeries),), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE)
    def ts_create(self, key: CommandItem, *args: bytes) -> SimpleString:
        if key.value is not None:
            raise SimpleError(msgs.TIMESERIES_KEY_EXISTS)
        key.value = self._create_timeseries(key.key, *args)
        return OK

    @command(name="TS.ADD", fixed=(Key(TimeSeries), Timestamp, Float), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE)
    def ts_add(self, key: CommandItem, timestamp: int, value: float, *args: bytes) -> int:
        (on_duplicate,), left_args = extract_args(args, ("*on_duplicate",), error_on_unexpected=False)
        if key.value is None:
            key.update(self._create_timeseries(key.key, *args))
        if not self._validate_duplicate_policy(on_duplicate):
            raise SimpleError(msgs.TIMESERIES_INVALID_DUPLICATE_POLICY)
        res = key.value.add(timestamp, value, on_duplicate)
        return res

    @command(name="TS.GET", fixed=(Key(TimeSeries),), repeat=(bytes,))
    def ts_get(self, key: CommandItem, *args: bytes) -> Optional[List[Union[int, float]]]:
        if key.value is None:
            raise SimpleError(msgs.TIMESERIES_KEY_DOES_NOT_EXIST)
        return key.value.get()

    @command(
        name="TS.MADD",
        fixed=(Key(TimeSeries), Timestamp, Float),
        repeat=(Key(TimeSeries), Timestamp, Float),
        flags=msgs.FLAG_DO_NOT_CREATE,
    )
    def ts_madd(self, *args: Any) -> List[int]:
        if len(args) % 3 != 0:
            raise SimpleError(msgs.WRONG_ARGS_MSG6)
        results: List[int] = list()
        for i in range(0, len(args), 3):
            key, timestamp, value = args[i : i + 3]
            if key.value is None:
                raise SimpleError(msgs.TIMESERIES_KEY_DOES_NOT_EXIST)
            results.append(key.value.add(timestamp, value))
        return results

    @command(name="TS.DEL", fixed=(Key(TimeSeries), Int, Int), repeat=(), flags=msgs.FLAG_DO_NOT_CREATE)
    def ts_del(self, key: CommandItem, from_ts: int, to_ts: int) -> bytes:
        if key.value is None:
            raise SimpleError(msgs.TIMESERIES_KEY_DOES_NOT_EXIST)
        return key.value.delete(from_ts, to_ts)

    @command(
        name="TS.CREATERULE",
        fixed=(Key(TimeSeries), Key(TimeSeries), bytes, bytes, Int),
        repeat=(bytes,),
        flags=msgs.FLAG_DO_NOT_CREATE,
    )
    def ts_createrule(
        self,
        source_key: CommandItem,
        dest_key: CommandItem,
        _: bytes,
        aggregator: bytes,
        bucket_duration: int,
        *args: bytes,
    ) -> SimpleString:
        if source_key.value is None:
            raise SimpleError(msgs.TIMESERIES_KEY_DOES_NOT_EXIST)
        if dest_key.value is None:
            raise SimpleError(msgs.TIMESERIES_KEY_DOES_NOT_EXIST)
        if len(args) > 1:
            raise SimpleError(msgs.WRONG_ARGS_MSG6.format("ts.createrule"))
        try:
            align_timestamp = int(args[0]) if len(args) == 1 else 0
        except ValueError:
            raise SimpleError(msgs.TIMESERIES_BAD_TIMESTAMP)
        existing_rule = source_key.value.get_rule(dest_key.key)
        if existing_rule is not None:
            raise SimpleError(msgs.TIMESERIES_RULE_EXISTS)
        if aggregator not in AGGREGATORS:
            raise SimpleError(msgs.TIMESERIES_BAD_AGGREGATION_TYPE)
        rule = TimeSeriesRule(source_key.value, dest_key.value, aggregator, bucket_duration, align_timestamp)
        source_key.value.add_rule(rule)
        return OK

    @command(
        name="TS.DELETERULE",
        fixed=(Key(TimeSeries), Key(TimeSeries)),
        repeat=(),
        flags=msgs.FLAG_DO_NOT_CREATE,
    )
    def ts_deleterule(self, source_key: CommandItem, dest_key: CommandItem) -> bytes:
        if source_key.value is None:
            raise SimpleError(msgs.TIMESERIES_KEY_DOES_NOT_EXIST)
        res: Optional[TimeSeriesRule] = source_key.value.get_rule(dest_key.key)
        if res is None:
            raise SimpleError(msgs.TIMESERIES_RULE_DOES_NOT_EXIST)
        source_key.value.delete_rule(res)
        return OK

    def _ts_inc_or_dec(self, key: CommandItem, addend: float, *args: bytes) -> bytes:
        (ts,), left_args = extract_args(
            args,
            ("+timestamp",),
            error_on_unexpected=False,
        )
        if key.value is None:
            key.update(self._create_timeseries(key.key, *left_args))
        timeseries = key.value
        if ts is None:
            if len(timeseries.sorted_list) == 0:
                ts = int(time.time())
            else:
                ts = timeseries.sorted_list[-1][0]
        if len(timeseries.sorted_list) > 0 and ts < timeseries.sorted_list[-1][0]:
            raise SimpleError(msgs.TIMESERIES_INVALID_TIMESTAMP)
        try:
            return key.value.incrby(ts, addend)
        except ValueError:
            msg = (
                msgs.TIMESERIES_TIMESTAMP_LOWER_THAN_MAX_V7
                if self.version >= (7,)
                else msgs.TIMESERIES_TIMESTAMP_LOWER_THAN_MAX_V6
            )
            raise SimpleError(msg)

    @command(name="TS.INCRBY", fixed=(Key(TimeSeries), Float), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE)
    def ts_incrby(self, key: CommandItem, addend: float, *args: bytes) -> bytes:
        return self._ts_inc_or_dec(key, addend, *args)

    @command(name="TS.DECRBY", fixed=(Key(TimeSeries), Float), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE)
    def ts_decrby(self, key: CommandItem, subtrahend: float, *args: bytes) -> bytes:
        return self._ts_inc_or_dec(key, -subtrahend, *args)

    @command(name="TS.ALTER", fixed=(Key(TimeSeries),), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE)
    def ts_alter(self, key: CommandItem, *args: bytes) -> bytes:
        if key.value is None:
            raise SimpleError(msgs.TIMESERIES_KEY_DOES_NOT_EXIST)

        ((retention, chunk_size, duplicate_policy, (ignore_max_time_diff, ignore_max_val_diff)), left_args) = (
            extract_args(
                args, ("+retention", "+chunk_size", "*duplicate_policy", "++ignore"), error_on_unexpected=False
            )
        )

        if chunk_size is not None and chunk_size % 8 != 0:
            raise SimpleError(msgs.TIMESERIES_BAD_CHUNK_SIZE)
        if not self._validate_duplicate_policy(duplicate_policy):
            raise SimpleError(msgs.TIMESERIES_INVALID_DUPLICATE_POLICY)
        duplicate_policy = duplicate_policy.lower() if duplicate_policy else None
        if len(left_args) > 0 and (not casematch(left_args[0], b"LABELS") or len(left_args) % 2 != 1):
            raise SimpleError(msgs.BAD_SUBCOMMAND_MSG.format("TS.ADD"))
        labels = dict(zip(left_args[1::2], left_args[2::2])) if len(left_args) > 0 else {}

        key.value.retention = retention or key.value.retention
        key.value.chunk_size = chunk_size or key.value.chunk_size
        key.value.duplicate_policy = duplicate_policy or key.value.duplicate_policy
        key.value.ignore_max_time_diff = ignore_max_time_diff or key.value.ignore_max_time_diff
        key.value.ignore_max_val_diff = ignore_max_val_diff or key.value.ignore_max_val_diff
        key.value.labels = labels or key.value.labels
        key.updated()
        return OK

    def _range(
        self, reverse: bool, ts: TimeSeries, from_ts: int, to_ts: int, *args: bytes
    ) -> List[List[Union[int, float]]]:
        RANGE_ARGS = ("latest", "++filter_by_value", "+count", "*align", "*+aggregation", "*buckettimestamp", "empty")
        (
            (
                latest,
                (value_min, value_max),
                count,
                align,
                (aggregator, bucket_duration),
                bucket_timestamp,
                empty,
            ),
            left_args,
        ) = extract_args(args, RANGE_ARGS, error_on_unexpected=False, left_from_first_unexpected=False)
        latest = True
        filter_ts: Optional[List[int]] = None
        if len(left_args) > 0:
            if not casematch(left_args[0], b"FILTER_BY_TS"):
                raise SimpleError(msgs.WRONG_ARGS_MSG6)
            left_args = left_args[1:]
            filter_ts = [int(x) for x in left_args]
        if aggregator is None and (align is not None or bucket_timestamp is not None or empty):
            raise SimpleError(msgs.WRONG_ARGS_MSG6)
        if bucket_timestamp is not None and bucket_timestamp not in (b"-", b"+", b"~"):
            raise SimpleError(msgs.WRONG_ARGS_MSG6)
        if align is not None:
            if align == b"+":
                align = to_ts
            elif align == b"-":
                align = from_ts
            else:
                align = int(align)
        if aggregator is not None and aggregator not in AGGREGATORS:
            raise SimpleError(msgs.TIMESERIES_BAD_AGGREGATION_TYPE)
        if aggregator is None:
            res = ts.range(from_ts, to_ts, value_min, value_max, count, filter_ts, reverse)
        else:
            res = ts.aggregate(
                from_ts,
                to_ts,
                latest,
                value_min,
                value_max,
                count,
                filter_ts,
                align,
                aggregator,
                bucket_duration,
                bucket_timestamp,
                empty,
                reverse,
            )

        res = [[x[0], x[1]] for x in res]
        return res

    @command(
        name="TS.RANGE", fixed=(Key(TimeSeries), Timestamp, Timestamp), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE
    )
    def ts_range(self, key: CommandItem, from_ts: int, to_ts: int, *args: bytes) -> List[List[Union[int, float]]]:
        if key.value is None:
            raise SimpleError(msgs.TIMESERIES_KEY_DOES_NOT_EXIST)
        return self._range(False, key.value, from_ts, to_ts, *args)

    @command(
        name="TS.REVRANGE",
        fixed=(Key(TimeSeries), Timestamp, Timestamp),
        repeat=(bytes,),
        flags=msgs.FLAG_DO_NOT_CREATE,
    )
    def ts_revrange(self, key: CommandItem, from_ts: int, to_ts: int, *args: bytes) -> List[List[Union[int, float]]]:
        if key.value is None:
            raise SimpleError(msgs.TIMESERIES_KEY_DOES_NOT_EXIST)
        res = self._range(True, key.value, from_ts, to_ts, *args)
        return res

    @command(name="TS.MGET", fixed=(bytes,), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE)
    def ts_mget(self, *args: bytes) -> List[List[Union[bytes, List[List[Union[int, float]]]]]]:
        latest, with_labels, selected_labels, filter_expression = False, False, None, None
        i = 0
        while i < len(args):
            if casematch(args[i], b"LATEST"):
                latest = True  # noqa: F841
                i += 1
            elif casematch(args[i], b"WITHLABELS"):
                with_labels = True
                i += 1
            elif casematch(args[i], b"SELECTED_LABELS"):
                selected_labels = list()
                i += 1
                while i < len(args) and casematch(args[i], b"FILTER"):
                    selected_labels.append(args[i])
            elif casematch(args[i], b"FILTER"):
                filter_expression = list()
                i += 1
                while i < len(args):
                    filter_expression.append(args[i])
                    i += 1

        if with_labels and selected_labels is not None:
            raise SimpleError(msgs.WRONG_ARGS_MSG6.format("ts.mget"))
        if filter_expression is None or len(filter_expression) == 0:
            raise SimpleError(msgs.WRONG_ARGS_MSG6.format("ts.mget"))

        timeseries = self._get_timeseries(filter_expression)
        if with_labels:
            return [[ts.name, [[k, v] for (k, v) in ts.labels.items()], ts.get()] for ts in timeseries]
        if selected_labels is not None:
            res = [
                [ts.name, [[label, ts.labels[label]] for label in selected_labels if label in ts.labels], ts.get()]
                for ts in timeseries
            ]
        else:
            res = [[ts.name, [], ts.get()] for ts in timeseries]
        return res

    @command(name="TS.QUERYINDEX", fixed=(bytes,), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE)
    def ts_queryindex(self, *args: bytes) -> List[bytes]:
        filter_expressions = list(args)
        timeseries = self._get_timeseries(filter_expressions)
        return [ts.name for ts in timeseries]

    def _group_by_label(self, reverse: bool, ts_list: List[Any], label: bytes, reducer: bytes) -> TimeSeries:
        # ts_list: [[name, labels, measurements], ...]
        reducer = reducer.lower()
        if reducer not in AGGREGATORS:
            raise SimpleError(msgs.TIMESERIES_BAD_AGGREGATION_TYPE)
        ts_map: Dict[bytes, Dict[int, List[float]]] = dict()  # label_value -> timestamp -> values
        for ts in ts_list:
            # Find label value
            labels, label_value = ts[1], None
            for label_name, current_value in labels:
                if label_name == label:
                    label_value = current_value
                    break
            if not label_value:
                raise SimpleError(msgs.TIMESERIES_BAD_FILTER_EXPRESSION)
            if label_value not in ts_map:
                ts_map[label_value] = dict()
            # Collect measurements
            for timestamp, value in ts[2]:
                if timestamp not in ts_map[label_value]:
                    ts_map[label_value][timestamp] = list()
                ts_map[label_value][timestamp].append(value)
        res = []
        for label_value, timestamp_values in ts_map.items():
            sorted_timestamps = sorted(timestamp_values.keys())
            name = f"{label.decode()}={label_value.decode()}"
            sources = (", ".join([ts[0].decode() for ts in ts_list])).encode("utf-8")
            labels = {label: label_value, b"__reducer__": reducer, b"__source__": sources}
            measurements: List[List[Union[int, float]]] = [
                [timestamp, float(AGGREGATORS[reducer](timestamp_values[timestamp]))] for timestamp in sorted_timestamps
            ]
            if reverse:
                measurements.reverse()
            res.append([name.encode("utf-8"), [[k, v] for (k, v) in labels.items()], measurements])
        return res

    def _mrange(self, reverse: bool, from_ts: int, to_ts: int, *args: bytes):
        args_lower = [arg.lower() for arg in args]
        arg_words = {
            b"latest",
            b"withlabels",
            b"selected_labels",
            b"filter",
            b"groupby",
            b"reduce",
            b"count",
            b"aggregation",
            b"filter_by_value",
            b"filter_by_ts",
            b"align",
            b"aggregation",
        }
        left_args = []
        latest, with_labels, selected_labels, filter_expression, group_by, reducer = (
            False,
            False,
            None,
            None,
            None,
            None,
        )
        i = 0
        while i < len(args_lower):
            if args_lower[i] == b"latest":
                latest = True  # noqa: F841
                i += 1
            elif args_lower[i] == b"withlabels":
                with_labels = True
                i += 1
            elif args_lower[i] == b"selected_labels":
                selected_labels = list()
                i += 1
                while i < len(args_lower) and args_lower[i] not in arg_words:
                    selected_labels.append(args_lower[i])
                    i += 1
            elif args_lower[i] == b"filter":
                filter_expression = list()
                i += 1
                while i < len(args_lower) and args_lower[i] not in arg_words:
                    filter_expression.append(args[i])
                    i += 1
            elif i + 3 < len(args_lower) and args_lower[i] == b"groupby" and args_lower[i + 2] == b"reduce":
                group_by = args[i + 1]
                reducer = args_lower[i + 3]
                i += 4
            else:
                left_args.append(args[i])
                i += 1

        if with_labels and selected_labels is not None:
            raise SimpleError(msgs.WRONG_ARGS_MSG6.format("ts.mrange"))
        if filter_expression is None or len(filter_expression) == 0:
            raise SimpleError(msgs.WRONG_ARGS_MSG6.format("ts.mrange"))

        timeseries = self._get_timeseries(filter_expression)
        if with_labels or (group_by is not None and reducer is not None):
            res = [
                [
                    ts.name,
                    [[k, v] for (k, v) in ts.labels.items()],
                    self._range(reverse, ts, from_ts, to_ts, *left_args),
                ]
                for ts in timeseries
            ]
        elif selected_labels is not None:
            res = [
                [
                    ts.name,
                    [[label, ts.labels[label]] for label in selected_labels if label in ts.labels],
                    self._range(reverse, ts, from_ts, to_ts, *left_args),
                ]
                for ts in timeseries
            ]
        else:
            res = [[ts.name, [], self._range(reverse, ts, from_ts, to_ts, *left_args)] for ts in timeseries]
        if group_by is not None and reducer is not None:
            return self._group_by_label(reverse, res, group_by, reducer)
        return res

    @command(name="TS.MRANGE", fixed=(Timestamp, Timestamp), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE)
    def ts_mrange(
        self, from_ts: int, to_ts: int, *args: bytes
    ) -> List[List[Union[bytes, List[List[Union[int, float]]]]]]:
        return self._mrange(False, from_ts, to_ts, *args)

    @command(name="TS.MREVRANGE", fixed=(Timestamp, Timestamp), repeat=(bytes,), flags=msgs.FLAG_DO_NOT_CREATE)
    def ts_mrevrange(
        self, from_ts: int, to_ts: int, *args: bytes
    ) -> List[List[Union[bytes, List[List[Union[int, float]]]]]]:
        return self._mrange(True, from_ts, to_ts, *args)