"""
Knowledge base implementation using Ollama embeddings.
Provides tools for searching and adding documents.
"""

import numpy as np
from typing import List
from langchain.docstore.document import Document
from embeddings import get_embedding_model

class KnowledgeBase:
    """
    In-memory knowledge base that stores documents and their embeddings.
    """

    def __init__(self):
        self.embedding_model = get_embedding_model()
        self.documents: List[Document] = []
        self.embeddings: np.ndarray = np.empty((0, self.embedding_model.get_sentence_embedding_dimension()))

    def add_to_knowledge_base(self, content: str) -> str:
        """
        Adds a new document to the knowledge base.

        Args:
            content (str): The text content to add.

        Returns:
            str: Confirmation message.
        """
        doc = Document(page_content=content)
        embedding = self.embedding_model.embed_query(content)
        embedding = np.array(embedding).reshape(1, -1)

        self.documents.append(doc)
        if self.embeddings.size == 0:
            self.embeddings = embedding
        else:
            self.embeddings = np.vstack([self.embeddings, embedding])

        return f"Document added. Total documents: {len(self.documents)}."

    def search_knowledge_base(self, query: str, k: int = 3) -> List[Document]:
        """
        Searches the knowledge base for the most relevant documents.

        Args:
            query (str): The search query.
            k (int): Number of top documents to return.

        Returns:
            List[Document]: List of top matching documents.
        """
        if not self.documents:
            return []

        query_embedding = self.embedding_model.embed_query(query)
        query_embedding = np.array(query_embedding).reshape(1, -1)

        similarities = np.dot(self.embeddings, query_embedding.T).flatten()
        top_indices = similarities.argsort()[-k:][::-1]
        return [self.documents[i] for i in top_indices]

# Global knowledge base instance
kb = KnowledgeBase()

def add_to_knowledge_base(content: str) -> str:
    """
    Tool wrapper for adding content to the knowledge base.

    Args:
        content (str): Text to add.

    Returns:
        str: Confirmation message.
    """
    return kb.add_to_knowledge_base(content)

def search_knowledge_base(query: str) -> List[Document]:
    """
    Tool wrapper for searching the knowledge base.

    Args:
        query (str): Search query.

    Returns:
        List[Document]: Matching documents.
    """
    return kb.search_knowledge_base(query)