from asyncio import Lock
from base64 import urlsafe_b64decode
from collections.abc import Callable, Iterable, MutableMapping, MutableSequence, Sequence
from typing import Any, cast

from haiway import AttributePath, AttributeRequirement, as_tuple

from draive.embedding import (
    Embedded,
    ImageEmbedding,
    TextEmbedding,
    mmr_vector_similarity_search,
    vector_similarity_search,
)
from draive.multimodal import TextContent
from draive.parameters import DataModel
from draive.resources import ResourceContent, ResourceReference
from draive.utils import VectorIndex

__all__ = ("VolatileVectorIndex",)


def VolatileVectorIndex() -> VectorIndex:  # noqa: C901, PLR0915
    lock: Lock = Lock()
    storage: MutableMapping[type[Any], MutableSequence[Embedded[Any]]] = {}

    async def index[Model: DataModel, Value: ResourceContent | TextContent | str](  # noqa: C901, PLR0912
        model: type[Model],
        /,
        *,
        attribute: Callable[[Model], Value] | AttributePath[Model, Value] | Value,
        values: Iterable[Model],
        **extra: Any,
    ) -> None:
        values_sequence: Sequence[Model] = as_tuple(values)
        async with lock:
            value_selector: Callable[[Model], Value]
            match attribute:
                case Callable() as selector:
                    value_selector = selector

                case path:
                    assert isinstance(  # nosec: B101
                        path, AttributePath
                    ), "Prepare parameter path by using Self._.path.to.property"
                    value_selector = cast(AttributePath[Model, Value], path).__call__

            selected_values: list[str | bytes] = []
            for value in values_sequence:
                match value_selector(value):
                    case str() as text:
                        selected_values.append(text)

                    case TextContent() as text_content:
                        selected_values.append(text_content.text)

                    case ResourceContent() as data:
                        if not data.mime_type.startswith("image"):
                            raise ValueError(f"{data.mime_type} embedding is not supported")

                        selected_values.append(urlsafe_b64decode(data.data.encode("utf-8")))

                    case ResourceReference():
                        raise ValueError("Resource references are not supported")

            embedded_values: Sequence[Embedded[str] | Embedded[bytes]]
            if all(isinstance(value, str) for value in selected_values):
                embedded_values = await TextEmbedding.embed(
                    cast(list[str], selected_values),
                    **extra,
                )

            elif all(isinstance(value, bytes | bytearray) for value in selected_values):
                embedded_values = await ImageEmbedding.embed(
                    cast(list[bytes], selected_values),
                    **extra,
                )

            else:
                raise ValueError("Selected attribute values have to be the same type")

            embedded_models: list[Embedded[Model]] = [
                Embedded(
                    value=value,
                    vector=embedded.vector,
                )
                for value, embedded in zip(
                    values_sequence,
                    embedded_values,
                    strict=True,
                )
            ]

            if model in storage:
                storage[model].extend(embedded_models)

            else:
                storage[model] = embedded_models

    async def search[Model: DataModel](  # noqa: C901, PLR0912
        model: type[Model],
        /,
        *,
        query: Sequence[float] | ResourceContent | TextContent | str | None = None,
        score_threshold: float | None = None,
        requirements: AttributeRequirement[Model] | None = None,
        limit: int | None = None,
        **extra: Any,
    ) -> Sequence[Model]:
        assert query is not None or (query is None and score_threshold is None)  # nosec: B101
        stored: Sequence[Embedded[Model]] | None
        async with lock:
            stored = storage.get(model)

        if not stored:
            return []

        filtered: Sequence[Embedded[Model]]
        if requirements:
            filtered = [
                element
                for element in stored
                if requirements.check(
                    element.value,
                    raise_exception=False,
                )
            ]

        else:
            filtered = stored

        if not filtered:
            return []

        query_vector: Sequence[float]
        match query:
            case None:
                if limit:
                    return [embedded.value for embedded in filtered[:limit]]

                else:
                    return [embedded.value for embedded in filtered]

            case str() as text:
                embedded_text: Embedded[str] = await TextEmbedding.embed(
                    text,
                    **extra,
                )
                query_vector = embedded_text.vector

            case TextContent() as text_content:
                embedded_text: Embedded[str] = await TextEmbedding.embed(
                    text_content.text,
                    **extra,
                )
                query_vector = embedded_text.vector

            case ResourceContent() as data:
                if not data.mime_type.startswith("image"):
                    raise ValueError(f"{data.mime_type} embedding is not supported")

                embedded_image: Embedded[bytes] = await ImageEmbedding.embed(
                    urlsafe_b64decode(data.data.encode("utf-8")),
                    **extra,
                )
                query_vector = embedded_image.vector

            case ResourceReference():
                raise ValueError("Resource references are not supported")

            case vector:
                query_vector = vector

        matching: Sequence[Embedded[Model]] = [
            filtered[index]
            for index in vector_similarity_search(
                query_vector=query_vector,
                values_vectors=[element.vector for element in filtered],
                score_threshold=score_threshold,
                limit=limit * 8  # feed MMR with more results
                if limit is not None
                else None,
            )
        ]

        if not matching:
            return []

        return [
            matching[index].value
            for index in mmr_vector_similarity_search(
                query_vector=query_vector,
                values_vectors=[element.vector for element in matching],
                limit=limit,
            )
        ]

    async def delete[Model: DataModel](
        model: type[Model],
        /,
        *,
        requirements: AttributeRequirement[Model] | None = None,
        **extra: Any,
    ) -> None:
        async with lock:
            if model not in storage:
                return

            if requirements is not None:
                storage[model] = [
                    element
                    for element in storage[model]
                    if not requirements.check(
                        element.value,
                        raise_exception=False,
                    )
                ]

            else:
                del storage[model]

    return VectorIndex(
        indexing=index,
        searching=search,
        deleting=delete,
    )
