# Copyright (c) Microsoft. All rights reserved.

import ast
import asyncio
import contextlib
import json
import logging
import sys
from abc import abstractmethod
from collections.abc import MutableSequence, Sequence
from copy import copy
from enum import Enum
from typing import Any, ClassVar, Final, Generic, TypeVar

from pydantic import SecretStr, ValidationError
from redis.asyncio.client import Redis
from redis.commands.search.field import Field as RedisField
from redis.commands.search.field import NumericField, TagField, TextField, VectorField
from redis.commands.search.index_definition import IndexDefinition, IndexType
from redisvl.index.index import process_results
from redisvl.query.filter import FilterExpression, Num, Tag, Text
from redisvl.query.query import BaseQuery, VectorQuery
from redisvl.redis.utils import array_to_buffer, buffer_to_array, convert_bytes
from redisvl.schema import StorageType

from semantic_kernel.connectors.ai.embedding_generator_base import EmbeddingGeneratorBase
from semantic_kernel.data.vector import (
    DistanceFunction,
    FieldTypes,
    GetFilteredRecordOptions,
    IndexKind,
    KernelSearchResults,
    SearchType,
    TModel,
    VectorSearch,
    VectorSearchOptions,
    VectorSearchResult,
    VectorStore,
    VectorStoreCollection,
    VectorStoreCollectionDefinition,
    VectorStoreField,
)
from semantic_kernel.exceptions import (
    VectorSearchExecutionException,
    VectorSearchOptionsException,
    VectorStoreInitializationException,
    VectorStoreOperationException,
)
from semantic_kernel.kernel_pydantic import KernelBaseSettings
from semantic_kernel.utils.feature_stage_decorator import release_candidate
from semantic_kernel.utils.list_handler import desync_list

if sys.version_info >= (3, 12):
    from typing import override  # pragma: no cover
else:
    from typing_extensions import override  # pragma: no cover

logger = logging.getLogger(__name__)

TKey = TypeVar("TKey", bound=str)
TQuery = TypeVar("TQuery", bound=BaseQuery)


class RedisCollectionTypes(str, Enum):
    """Redis collection types."""

    JSON = "json"
    HASHSET = "hashset"


STORAGE_TYPE_MAP: Final[dict[RedisCollectionTypes, StorageType]] = {
    RedisCollectionTypes.JSON: StorageType.JSON,
    RedisCollectionTypes.HASHSET: StorageType.HASH,
}
DISTANCE_FUNCTION_MAP: Final[dict[DistanceFunction, str]] = {
    DistanceFunction.COSINE_SIMILARITY: "COSINE",
    DistanceFunction.DOT_PROD: "IP",
    DistanceFunction.EUCLIDEAN_DISTANCE: "L2",
    DistanceFunction.DEFAULT: "COSINE",
}
INDEX_KIND_MAP: Final[dict[IndexKind, str]] = {
    IndexKind.HNSW: "HNSW",
    IndexKind.FLAT: "FLAT",
    IndexKind.DEFAULT: "HNSW",
}
INDEX_TYPE_MAP: Final[dict[RedisCollectionTypes, IndexType]] = {
    RedisCollectionTypes.JSON: IndexType.JSON,
    RedisCollectionTypes.HASHSET: IndexType.HASH,
}
DATATYPE_MAP_VECTOR: Final[dict[str, str]] = {
    "float": "FLOAT32",
    "int": "FLOAT16",
    "binary": "FLOAT16",
    "ndarray": "FLOAT32",
    "default": "FLOAT32",
}


def _field_to_redis_field_hashset(name: str, field: VectorStoreField) -> RedisField:
    if field.field_type == FieldTypes.VECTOR:
        if field.distance_function not in DISTANCE_FUNCTION_MAP:
            raise VectorStoreOperationException(
                f"Distance function {field.distance_function} is not supported. "
                f"Supported functions are: {list(DISTANCE_FUNCTION_MAP.keys())}"
            )
        if field.index_kind not in INDEX_KIND_MAP:
            raise VectorStoreOperationException(
                f"Index kind {field.index_kind} is not supported. Supported kinds are: {list(INDEX_KIND_MAP.keys())}"
            )
        return VectorField(
            name=name,
            algorithm=INDEX_KIND_MAP[field.index_kind],
            attributes={
                "type": DATATYPE_MAP_VECTOR[field.type_ or "default"],
                "dim": field.dimensions,
                "distance_metric": DISTANCE_FUNCTION_MAP[field.distance_function],
            },
        )
    if field.type_ in ["int", "float"]:
        return NumericField(name=name)
    if field.is_full_text_indexed:
        return TextField(name=name)
    return TagField(name=name)


def _field_to_redis_field_json(name: str, field: VectorStoreField) -> RedisField:
    if field.field_type == FieldTypes.VECTOR:
        if field.distance_function not in DISTANCE_FUNCTION_MAP:
            raise VectorStoreOperationException(
                f"Distance function {field.distance_function} is not supported. "
                f"Supported functions are: {list(DISTANCE_FUNCTION_MAP.keys())}"
            )
        if field.index_kind not in INDEX_KIND_MAP:
            raise VectorStoreOperationException(
                f"Index kind {field.index_kind} is not supported. Supported kinds are: {list(INDEX_KIND_MAP.keys())}"
            )
        return VectorField(
            name=f"$.{name}",
            algorithm=INDEX_KIND_MAP[field.index_kind],
            attributes={
                "type": DATATYPE_MAP_VECTOR[field.type_ or "default"],
                "dim": field.dimensions,
                "distance_metric": DISTANCE_FUNCTION_MAP[field.distance_function],
            },
            as_name=name,
        )
    if field.type_ in ["int", "float"]:
        return NumericField(name=f"$.{name}", as_name=name)
    if field.is_full_text_indexed:
        return TextField(name=f"$.{name}", as_name=name)
    return TagField(name=f"$.{name}", as_name=name)


def _definition_to_redis_fields(
    definition: VectorStoreCollectionDefinition, collection_type: RedisCollectionTypes
) -> list[RedisField]:
    """Create a list of fields for Redis from a definition."""
    fields: list[RedisField] = []
    for field in definition.fields:
        if field.field_type == FieldTypes.KEY:
            continue
        if collection_type == RedisCollectionTypes.HASHSET:
            fields.append(_field_to_redis_field_hashset(field.storage_name or field.name, field))  # type: ignore
        elif collection_type == RedisCollectionTypes.JSON:
            fields.append(_field_to_redis_field_json(field.storage_name or field.name, field))  # type: ignore
    return fields


@release_candidate
class RedisSettings(KernelBaseSettings):
    """Redis model settings.

    Args:
    - connection_string (str | None):
        Redis connection string (Env var REDIS_CONNECTION_STRING)
    """

    env_prefix: ClassVar[str] = "REDIS_"

    connection_string: SecretStr


@release_candidate
class RedisCollection(
    VectorStoreCollection[TKey, TModel],
    VectorSearch[TKey, TModel],
    Generic[TKey, TModel],
):
    """A vector store record collection implementation using Redis."""

    redis_database: Redis
    prefix_collection_name_to_key_names: bool
    collection_type: RedisCollectionTypes
    supported_key_types: ClassVar[set[str] | None] = {"str"}
    supported_vector_types: ClassVar[set[str] | None] = {"float"}
    supported_search_types: ClassVar[set[SearchType]] = {SearchType.VECTOR}

    def __init__(
        self,
        record_type: type[TModel],
        definition: VectorStoreCollectionDefinition | None = None,
        collection_name: str | None = None,
        embedding_generator: EmbeddingGeneratorBase | None = None,
        redis_database: Redis | None = None,
        prefix_collection_name_to_key_names: bool = True,
        collection_type: RedisCollectionTypes = RedisCollectionTypes.HASHSET,
        connection_string: str | None = None,
        env_file_path: str | None = None,
        env_file_encoding: str | None = None,
        **kwargs: Any,
    ) -> None:
        """RedisMemoryStore is an abstracted interface to interact with a Redis node connection.

        See documentation about connections: https://redis-py.readthedocs.io/en/stable/connections.html
        See documentation about vector attributes: https://redis.io/docs/stack/search/reference/vectors.

        """
        if redis_database:
            super().__init__(
                record_type=record_type,
                definition=definition,
                collection_name=collection_name,
                embedding_generator=embedding_generator,
                redis_database=redis_database,
                prefix_collection_name_to_key_names=prefix_collection_name_to_key_names,
                collection_type=collection_type,
                managed_client=False,
                **kwargs,
            )
            return
        try:
            redis_settings = RedisSettings(
                connection_string=connection_string,
                env_file_path=env_file_path,
                env_file_encoding=env_file_encoding,
            )
        except ValidationError as ex:
            raise VectorStoreInitializationException("Failed to create Redis settings.", ex) from ex
        super().__init__(
            record_type=record_type,
            definition=definition,
            collection_name=collection_name,
            embedding_generator=embedding_generator,
            redis_database=Redis.from_url(redis_settings.connection_string.get_secret_value()),
            prefix_collection_name_to_key_names=prefix_collection_name_to_key_names,
            collection_type=collection_type,
            **kwargs,
        )

    def _get_redis_key(self, key: TKey) -> TKey:
        if self.prefix_collection_name_to_key_names:
            return f"{self.collection_name}:{key}"  # type: ignore
        return key

    def _unget_redis_key(self, key: TKey) -> TKey:
        if self.prefix_collection_name_to_key_names and ":" in key:
            return key[len(self.collection_name) + 1 :]  # type: ignore
        return key

    @override
    async def ensure_collection_exists(self, **kwargs) -> None:
        """Create a new index in Redis.

        Args:
            **kwargs: Additional keyword arguments.
                fields (list[Fields]): The fields to create the index with, when not supplied,
                    these are created from the definition.
                index_definition (IndexDefinition): The search index to create, if this is supplied
                    this is used instead of a index created based on the definition.
                other kwargs are passed to the create_index method.
        """
        if (index_definition := kwargs.pop("index_definition", None)) and (fields := kwargs.pop("fields", None)):
            if isinstance(index_definition, IndexDefinition):
                await self.redis_database.ft(self.collection_name).create_index(
                    fields, definition=index_definition, **kwargs
                )
                return
            raise VectorStoreOperationException("Invalid index type supplied.")
        fields = _definition_to_redis_fields(self.definition, self.collection_type)
        index_definition = IndexDefinition(
            prefix=f"{self.collection_name}:", index_type=INDEX_TYPE_MAP[self.collection_type]
        )
        await self.redis_database.ft(self.collection_name).create_index(fields, definition=index_definition, **kwargs)

    @override
    async def collection_exists(self, **kwargs) -> bool:
        try:
            await self.redis_database.ft(self.collection_name).info()
            return True
        except Exception:
            return False

    @override
    async def ensure_collection_deleted(self, **kwargs) -> None:
        exists = await self.collection_exists()
        if exists:
            await self.redis_database.ft(self.collection_name).dropindex(**kwargs)
        else:
            logger.debug("Collection does not exist, skipping deletion.")

    @override
    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
        """Exit the context manager."""
        if self.managed_client:
            await self.redis_database.aclose()  # type: ignore

    @override
    async def _inner_search(
        self,
        search_type: SearchType,
        options: VectorSearchOptions,
        values: Any | None = None,
        vector: Sequence[float | int] | None = None,
        **kwargs: Any,
    ) -> KernelSearchResults[VectorSearchResult[TModel]]:
        if not vector:
            vector = await self._generate_vector_from_values(values, options)
        if not vector:
            raise VectorSearchExecutionException("No vector found.")
        query = self._construct_vector_query(vector, options, **kwargs)
        results = await self.redis_database.ft(self.collection_name).search(  # type: ignore
            query=query.query, query_params=query.params
        )
        processed = process_results(results, query, STORAGE_TYPE_MAP[self.collection_type])
        return KernelSearchResults(
            results=self._get_vector_search_results_from_results(desync_list(processed)),
            total_count=results.total,
        )

    def _construct_vector_query(
        self, vector: Sequence[float | int], options: VectorSearchOptions, **kwargs: Any
    ) -> VectorQuery:
        vector_field = self.definition.try_get_vector_field(options.vector_property_name)
        if not vector_field:
            raise VectorSearchOptionsException("Vector field not found.")

        query = VectorQuery(
            vector=vector,  # type: ignore
            vector_field_name=vector_field.storage_name or vector_field.name,  # type: ignore
            num_results=options.top + options.skip,
            dialect=2,
            return_score=True,
        )
        if filter := self._build_filter(options.filter):  # type: ignore
            if isinstance(filter, list):
                expr = filter[0]
                for v in filter[1:]:
                    expr = expr & v

                query.set_filter(expr)
            else:
                query.set_filter(filter)
        query.paging(offset=options.skip, num=options.top + options.skip)
        query.sort_by(
            query.DISTANCE_ID,
            asc=(vector_field.distance_function)
            in [
                DistanceFunction.COSINE_SIMILARITY,
                DistanceFunction.DOT_PROD,
                DistanceFunction.DEFAULT,
            ],
        )
        return self._add_return_fields(query, options.include_vectors)

    @override
    def _lambda_parser(self, node: ast.AST) -> FilterExpression:
        """Parse the lambda AST and return a RedisVL FilterExpression."""

        def get_field_expr(field_name):
            # Find the field in the data model
            field = next(
                (f for f in self.definition.fields if (f.storage_name or f.name) == field_name),
                None,
            )
            if field is None:
                raise VectorStoreOperationException(f"Field '{field_name}' not found in data model.")
            if field.field_type == FieldTypes.DATA:
                if field.is_full_text_indexed:
                    return lambda: Text(field_name)
                if field.type_ in ("int", "float"):
                    return lambda: Num(field_name)
                return lambda: Tag(field_name)
            if field.field_type == FieldTypes.VECTOR:
                raise VectorStoreOperationException(f"Cannot filter on vector field '{field_name}'.")
            return lambda: Tag(field_name)

        match node:
            case ast.Compare():
                if len(node.ops) > 1:
                    # Chain comparisons (e.g., 1 < x < 3) become & of each comparison
                    expr = None
                    for idx in range(len(node.ops)):
                        left = node.left if idx == 0 else node.comparators[idx - 1]
                        right = node.comparators[idx]
                        op = node.ops[idx]
                        sub = self._lambda_parser(ast.Compare(left=left, ops=[op], comparators=[right]))
                        expr = expr & sub if expr else sub
                    return expr
                left = node.left
                right = node.comparators[0]
                op = node.ops[0]
                # Only support field op value or value op field
                if isinstance(left, (ast.Attribute, ast.Name)):
                    field_name = left.attr if isinstance(left, ast.Attribute) else left.id
                    field_expr = get_field_expr(field_name)()
                    value = self._lambda_parser(right)
                    match op:
                        case ast.Eq():
                            return field_expr == value
                        case ast.NotEq():
                            return field_expr != value
                        case ast.Gt():
                            return field_expr > value
                        case ast.GtE():
                            return field_expr >= value
                        case ast.Lt():
                            return field_expr < value
                        case ast.LtE():
                            return field_expr <= value
                        case ast.In():
                            return field_expr == value  # Tag/Text/Num support list equality
                        case ast.NotIn():
                            return ~(field_expr == value)
                    raise NotImplementedError(f"Unsupported operator: {type(op)}")
                if isinstance(right, (ast.Attribute, ast.Name)):
                    # Reverse: value op field
                    field_name = right.attr if isinstance(right, ast.Attribute) else right.id
                    field_expr = get_field_expr(field_name)()
                    value = self._lambda_parser(left)
                    match op:
                        case ast.Eq():
                            return field_expr == value
                        case ast.NotEq():
                            return field_expr != value
                        case ast.Gt():
                            return field_expr < value
                        case ast.GtE():
                            return field_expr <= value
                        case ast.Lt():
                            return field_expr > value
                        case ast.LtE():
                            return field_expr >= value
                        case ast.In():
                            return field_expr == value
                        case ast.NotIn():
                            return ~(field_expr == value)
                    raise NotImplementedError(f"Unsupported operator: {type(op)}")
                raise NotImplementedError("Comparison must be between a field and a value.")
            case ast.BoolOp():
                op = node.op  # type: ignore
                values = [self._lambda_parser(v) for v in node.values]
                if isinstance(op, ast.And):
                    expr = values[0]
                    for v in values[1:]:
                        expr = expr & v
                    return expr
                if isinstance(op, ast.Or):
                    expr = values[0]
                    for v in values[1:]:
                        expr = expr | v
                    return expr
                raise NotImplementedError(f"Unsupported BoolOp: {type(op)}")
            case ast.UnaryOp():
                match node.op:
                    case ast.Not():
                        operand = self._lambda_parser(node.operand)
                        return ~operand
                    case ast.UAdd() | ast.USub() | ast.Invert():
                        raise NotImplementedError("Unary +, -, ~ are not supported in RedisVL filters.")
            case ast.Attribute():
                # Only allow attributes that are in the data model
                if node.attr not in self.definition.storage_names:
                    raise VectorStoreOperationException(
                        f"Field '{node.attr}' not in data model (storage property names are used)."
                    )
                return node.attr
            case ast.Name():
                # Only allow names that are in the data model
                if node.id not in self.definition.storage_names:
                    raise VectorStoreOperationException(
                        f"Field '{node.id}' not in data model (storage property names are used)."
                    )
                return node.id
            case ast.Constant():
                return node.value
        raise NotImplementedError(f"Unsupported AST node: {type(node)}")

    @abstractmethod
    def _add_return_fields(self, query: TQuery, include_vectors: bool) -> TQuery:
        """Add the return fields to the query.

        There is a difference between the JSON and Hashset collections,
        this method should be overridden by the subclasses.

        """
        pass

    @override
    def _get_record_from_result(self, result: dict[str, Any]) -> Any:
        """Get a record from a result."""
        ret = result.copy()
        ret.pop("vector_distance", None)
        for key, value in ret.items():
            with contextlib.suppress(json.JSONDecodeError):
                ret[key] = json.loads(value) if isinstance(value, str) else value
        return ret

    @override
    def _get_score_from_result(self, result: dict[str, Any]) -> float | None:
        return result.get("vector_distance")


@release_candidate
class RedisHashsetCollection(RedisCollection[TKey, TModel], Generic[TKey, TModel]):
    """A vector store record collection implementation using Redis Hashsets."""

    def __init__(
        self,
        record_type: type[TModel],
        definition: VectorStoreCollectionDefinition | None = None,
        collection_name: str | None = None,
        embedding_generator: EmbeddingGeneratorBase | None = None,
        redis_database: Redis | None = None,
        prefix_collection_name_to_key_names: bool = False,
        connection_string: str | None = None,
        env_file_path: str | None = None,
        env_file_encoding: str | None = None,
        **kwargs: Any,
    ) -> None:
        """RedisMemoryStore is an abstracted interface to interact with a Redis node connection.

        See documentation about connections: https://redis-py.readthedocs.io/en/stable/connections.html
        See documentation about vector attributes: https://redis.io/docs/stack/search/reference/vectors.

        """
        super().__init__(
            record_type=record_type,
            definition=definition,
            collection_name=collection_name,
            embedding_generator=embedding_generator,
            redis_database=redis_database,
            prefix_collection_name_to_key_names=prefix_collection_name_to_key_names,
            collection_type=RedisCollectionTypes.HASHSET,
            connection_string=connection_string,
            env_file_path=env_file_path,
            env_file_encoding=env_file_encoding,
            **kwargs,
        )

    @override
    async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence[TKey]:
        return await asyncio.gather(*[self._single_upsert(record) for record in records])

    async def _single_upsert(self, upsert_record: Any) -> TKey:
        await self.redis_database.hset(**upsert_record)
        return self._unget_redis_key(upsert_record["name"])

    @override
    async def _inner_get(
        self,
        keys: Sequence[TKey] | None = None,
        options: GetFilteredRecordOptions | None = None,
        **kwargs,
    ) -> Sequence[dict[str, Any]] | None:
        if not keys:
            if options is not None:
                raise NotImplementedError("Get without keys is not yet implemented.")
            return None
        results = await asyncio.gather(*[self._single_get(self._get_redis_key(key)) for key in keys])
        return [result for result in results if result]

    async def _single_get(self, key: str) -> dict[str, Any] | None:
        result = await self.redis_database.hgetall(key)
        if result:
            result = convert_bytes(result)
            result[self.definition.key_name] = key
        return result

    @override
    async def _inner_delete(self, keys: Sequence[TKey], **kwargs: Any) -> None:
        await self.redis_database.delete(*[self._get_redis_key(key) for key in keys])

    @override
    def _serialize_dicts_to_store_models(
        self,
        records: Sequence[dict[str, Any]],
        **kwargs: Any,
    ) -> Sequence[dict[str, Any]]:
        """Serialize the dict to a Redis store model."""
        results: MutableSequence[dict[str, Any]] = []
        for record in records:
            result: dict[str, Any] = {"mapping": {}}
            for field in self.definition.fields:
                if field.field_type == FieldTypes.VECTOR:
                    dtype = DATATYPE_MAP_VECTOR[field.type_ or "default"].lower()
                    result["mapping"][field.storage_name or field.name] = array_to_buffer(record[field.name], dtype)
                    continue
                if field.field_type == FieldTypes.KEY:
                    result["name"] = self._get_redis_key(record[field.name])
                    continue
                result["mapping"][field.storage_name or field.name] = record[field.name]
            results.append(result)
        return results

    @override
    def _deserialize_store_models_to_dicts(
        self,
        records: Sequence[dict[str, Any]],
        **kwargs: Any,
    ) -> Sequence[dict[str, Any]]:
        results = []
        for record in records:
            rec = record.copy()
            for field in self.definition.fields:
                match field.field_type:
                    case FieldTypes.KEY:
                        rec[field.name] = self._unget_redis_key(rec[field.name])
                    case "vector":
                        dtype = DATATYPE_MAP_VECTOR[field.type_ or "default"]
                        rec[field.name] = buffer_to_array(rec[field.name], dtype)
            results.append(rec)
        return results

    def _add_return_fields(self, query: TQuery, include_vectors: bool) -> TQuery:
        """Add the return fields to the query.

        For a Hashset index this should not be decoded, that is the only difference
        between this and the JSON collection.

        """
        for field in self.definition.fields:
            match field.field_type:
                case "vector":
                    if include_vectors:
                        query.return_field(field.name, decode_field=False)
                case _:
                    query.return_field(field.name)
        return query


@release_candidate
class RedisJsonCollection(RedisCollection[TKey, TModel], Generic[TKey, TModel]):
    """A vector store record collection implementation using Redis Json."""

    def __init__(
        self,
        record_type: type[TModel],
        definition: VectorStoreCollectionDefinition | None = None,
        collection_name: str | None = None,
        embedding_generator: EmbeddingGeneratorBase | None = None,
        redis_database: Redis | None = None,
        prefix_collection_name_to_key_names: bool = False,
        connection_string: str | None = None,
        env_file_path: str | None = None,
        env_file_encoding: str | None = None,
        **kwargs: Any,
    ) -> None:
        """RedisMemoryStore is an abstracted interface to interact with a Redis node connection.

        See documentation about connections: https://redis-py.readthedocs.io/en/stable/connections.html
        See documentation about vector attributes: https://redis.io/docs/stack/search/reference/vectors.

        """
        super().__init__(
            record_type=record_type,
            definition=definition,
            collection_name=collection_name,
            redis_database=redis_database,
            prefix_collection_name_to_key_names=prefix_collection_name_to_key_names,
            collection_type=RedisCollectionTypes.JSON,
            connection_string=connection_string,
            env_file_path=env_file_path,
            env_file_encoding=env_file_encoding,
            embedding_generator=embedding_generator,
            **kwargs,
        )

    @override
    async def _inner_upsert(self, records: Sequence[Any], **kwargs: Any) -> Sequence[TKey]:
        return await asyncio.gather(*[self._single_upsert(record) for record in records])

    async def _single_upsert(self, upsert_record: Any) -> TKey:
        await self.redis_database.json().set(upsert_record["name"], "$", upsert_record["value"])
        return self._unget_redis_key(upsert_record["name"])

    @override
    async def _inner_get(
        self,
        keys: Sequence[TKey] | None = None,
        options: GetFilteredRecordOptions | None = None,
        **kwargs,
    ) -> Sequence[dict[str, Any]] | None:
        if not keys:
            if options is not None:
                raise NotImplementedError("Get without keys is not yet implemented.")
            return None
        kwargs_copy = copy(kwargs)
        kwargs_copy.pop("include_vectors", None)
        redis_keys = [self._get_redis_key(key) for key in keys]
        results = await self.redis_database.json().mget(redis_keys, "$", **kwargs_copy)
        return [self._add_key(key, result[0]) for key, result in zip(redis_keys, results) if result]

    def _add_key(self, key: TKey, record: dict[str, Any]) -> dict[str, Any]:
        record[self.definition.key_name] = key
        return record

    @override
    async def _inner_delete(self, keys: Sequence[str], **kwargs: Any) -> None:
        await asyncio.gather(*[self.redis_database.json().delete(key, **kwargs) for key in keys])

    @override
    def _serialize_dicts_to_store_models(
        self,
        records: Sequence[dict[str, Any]],
        **kwargs: Any,
    ) -> Sequence[dict[str, Any]]:
        """Serialize the dict to a Redis store model."""
        results: MutableSequence[dict[str, Any]] = []
        for record in records:
            result: dict[str, Any] = {"value": {}}
            for field in self.definition.fields:
                if field.field_type == FieldTypes.KEY:
                    result["name"] = self._get_redis_key(record[field.name])
                    continue
                if field.field_type == "vector":
                    result["value"][field.storage_name or field.name] = record[field.name]
                result["value"][field.storage_name or field.name] = record[field.name]
            results.append(result)
        return results

    @override
    def _deserialize_store_models_to_dicts(
        self,
        records: Sequence[dict[str, Any]],
        **kwargs: Any,
    ) -> Sequence[dict[str, Any]]:
        results = []
        key_field_name = self.definition.key_name
        for record in records:
            rec = record.copy()
            rec[key_field_name] = self._unget_redis_key(record[key_field_name])
            results.append(rec)
        return results

    def _add_return_fields(self, query: TQuery, include_vectors: bool) -> TQuery:
        """Add the return fields to the query."""
        for field in self.definition.fields:
            match field.field_type:
                case FieldTypes.VECTOR:
                    if include_vectors:
                        query.return_field(field.name)
                case _:
                    query.return_field(field.name)
        return query


@release_candidate
class RedisStore(VectorStore):
    """Create a Redis Vector Store."""

    redis_database: Redis

    def __init__(
        self,
        connection_string: str | None = None,
        embedding_generator: EmbeddingGeneratorBase | None = None,
        env_file_path: str | None = None,
        env_file_encoding: str | None = None,
        redis_database: Redis | None = None,
        **kwargs: Any,
    ) -> None:
        """RedisMemoryStore is an abstracted interface to interact with a Redis instance."""
        if redis_database:
            super().__init__(
                redis_database=redis_database,
                embedding_generator=embedding_generator,
                **kwargs,
            )
            return
        try:
            redis_settings = RedisSettings(
                connection_string=connection_string,
                env_file_path=env_file_path,
                env_file_encoding=env_file_encoding,
            )
        except ValidationError as ex:
            raise VectorStoreInitializationException("Failed to create Redis settings.", ex) from ex
        super().__init__(
            redis_database=Redis.from_url(redis_settings.connection_string.get_secret_value()),
            embedding_generator=embedding_generator,
            **kwargs,
        )

    @override
    async def list_collection_names(self, **kwargs) -> Sequence[str]:
        return [name.decode() for name in await self.redis_database.execute_command("FT._LIST")]

    @override
    def get_collection(
        self,
        record_type: type[TModel],
        *,
        definition: VectorStoreCollectionDefinition | None = None,
        collection_name: str | None = None,
        embedding_generator: EmbeddingGeneratorBase | None = None,
        collection_type: RedisCollectionTypes = RedisCollectionTypes.HASHSET,
        **kwargs: Any,
    ) -> RedisCollection:
        """Get a RedisCollection instance.

        Args:
            record_type: The type of the data model.
            definition: The model fields, optional.
            collection_name: The name of the collection.
            embedding_generator: The embedding generator to use.
            collection_type: The type of the collection, can be JSON or HASHSET.
            **kwargs: Additional keyword arguments, passed to the collection constructor.
        """
        if collection_type == RedisCollectionTypes.HASHSET:
            return RedisHashsetCollection(
                record_type=record_type,
                definition=definition,
                collection_name=collection_name,
                redis_database=self.redis_database,
                embedding_generator=embedding_generator or self.embedding_generator,
                **kwargs,
            )
        if collection_type == RedisCollectionTypes.JSON:
            return RedisJsonCollection(
                record_type=record_type,
                definition=definition,
                collection_name=collection_name,
                redis_database=self.redis_database,
                embedding_generator=embedding_generator or self.embedding_generator,
                **kwargs,
            )
        raise VectorStoreOperationException(
            f"Collection type {collection_type} is not supported. Supported types are: {RedisCollectionTypes}"
        )

    @override
    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
        """Exit the context manager."""
        if self.managed_client:
            await self.redis_database.aclose()  # type: ignore