from typing import Optional, cast
from uuid import uuid4

import faiss
from langchain_community.docstore import InMemoryDocstore
from langchain_community.vectorstores import FAISS
from langchain_community.vectorstores.utils import DistanceStrategy
from langchain_core.documents import Document
from langchain_huggingface import HuggingFaceEmbeddings
from openkosmos_ai.common.model import AIBaseConfig, SearchResult
from openkosmos_ai.nodes.base_node import BaseStateNode
from pydantic import BaseModel


class EmbeddingConfig(BaseModel):
    model_name: str
    model_kwargs: Optional[dict] = {"device": "cpu"}


class VectorStoreConfig(AIBaseConfig):
    embedding: EmbeddingConfig
    index_dir: Optional[str] = None
    index_name: Optional[str] = None
    save_index: bool = False
    load_index: bool = False

    def enable_save(self, enable=True):
        self.save_index = enable
        return self

    def enable_load(self, enable=True):
        self.load_index = enable
        return self


class VectorStoreNode(BaseStateNode[VectorStoreConfig]):
    def __init__(self, config: VectorStoreConfig):
        super().__init__(config)
        self.vector_store = None
        self.embedding_model = HuggingFaceEmbeddings(**config.embedding.model_dump())
        if config.load_index:
            self.load_index()

    def store(self):
        return self.vector_store

    def load_index(self):
        self.vector_store = FAISS.load_local(folder_path=self.config().index_dir,
                                             index_name=self.config().index_name,
                                             embeddings=self.embedding_model,
                                             normalize_L2=True,
                                             distance_strategy=DistanceStrategy.MAX_INNER_PRODUCT,
                                             allow_dangerous_deserialization=True)

    def build_index(self, documents: list[Document], document_ids: list[str] = None):
        self.vector_store = FAISS(
            embedding_function=self.embedding_model,
            index=faiss.IndexFlatIP(len(self.embedding_model.embed_query("AI"))),
            docstore=InMemoryDocstore(),
            index_to_docstore_id={},
            normalize_L2=True,
            distance_strategy=DistanceStrategy.MAX_INNER_PRODUCT
        )

        doc_ids = document_ids
        if document_ids is None:
            doc_ids = [str(uuid4()) for _ in range(len(documents))]

        self.vector_store.add_documents(documents=documents, ids=doc_ids)

        if self.config().save_index:
            self.vector_store.save_local(folder_path=self.config().index_dir,
                                         index_name=self.config().index_name)
        return doc_ids

    def search(self, query: str, threshold=0.6, filter_meta={}, top_k: int = 3) -> list[SearchResult]:
        results = self.vector_store.similarity_search_with_score(query=query,
                                                                 k=top_k,
                                                                 score_threshold=threshold,
                                                                 filter=filter_meta)
        return [
            SearchResult(id=doc.id, similarity_score=score, content=doc.page_content, metadata=doc.metadata)
            for doc, score in cast(list[tuple[Document, float]], results)]
