from abc import ABC, abstractmethod
from concurrent import futures
from dataclasses import dataclass
from typing import Optional, Any
from attr import define, field, Factory
from swarms.utils.futures import execute_futures_dict
from griptape.artifacts import TextArtifact


@define
class BaseVectorStore(ABC):
    """ """

    DEFAULT_QUERY_COUNT = 5

    @dataclass
    class QueryResult:
        id: str
        vector: list[float]
        score: float
        meta: Optional[dict] = None
        namespace: Optional[str] = None

    @dataclass
    class Entry:
        id: str
        vector: list[float]
        meta: Optional[dict] = None
        namespace: Optional[str] = None

    embedding_driver: Any
    futures_executor: futures.Executor = field(
        default=Factory(lambda: futures.ThreadPoolExecutor()), kw_only=True)

    def upsert_text_artifacts(self,
                              artifacts: dict[str, list[TextArtifact]],
                              meta: Optional[dict] = None,
                              **kwargs) -> None:
        execute_futures_dict({
            namespace:
            self.futures_executor.submit(self.upsert_text_artifact, a,
                                         namespace, meta, **kwargs)
            for namespace, artifact_list in artifacts.items()
            for a in artifact_list
        })

    def upsert_text_artifact(self,
                             artifact: TextArtifact,
                             namespace: Optional[str] = None,
                             meta: Optional[dict] = None,
                             **kwargs) -> str:
        if not meta:
            meta = {}

        meta["artifact"] = artifact.to_json()

        if artifact.embedding:
            vector = artifact.embedding
        else:
            vector = artifact.generate_embedding(self.embedding_driver)

        return self.upsert_vector(vector,
                                  vector_id=artifact.id,
                                  namespace=namespace,
                                  meta=meta,
                                  **kwargs)

    def upsert_text(self,
                    string: str,
                    vector_id: Optional[str] = None,
                    namespace: Optional[str] = None,
                    meta: Optional[dict] = None,
                    **kwargs) -> str:
        return self.upsert_vector(self.embedding_driver.embed_string(string),
                                  vector_id=vector_id,
                                  namespace=namespace,
                                  meta=meta if meta else {},
                                  **kwargs)

    @abstractmethod
    def upsert_vector(self,
                      vector: list[float],
                      vector_id: Optional[str] = None,
                      namespace: Optional[str] = None,
                      meta: Optional[dict] = None,
                      **kwargs) -> str:
        ...

    @abstractmethod
    def load_entry(self,
                   vector_id: str,
                   namespace: Optional[str] = None) -> Entry:
        ...

    @abstractmethod
    def load_entries(self, namespace: Optional[str] = None) -> list[Entry]:
        ...

    @abstractmethod
    def query(self,
              query: str,
              count: Optional[int] = None,
              namespace: Optional[str] = None,
              include_vectors: bool = False,
              **kwargs) -> list[QueryResult]:
        ...
