import logging
import uuid
from typing import List, Optional, Any
from functools import partial
from os import cpu_count
from concurrent.futures import ThreadPoolExecutor

from pydantic import Field, PrivateAttr
from sqlalchemy import Engine
from llama_index.core.base.llms.types import ChatResponse

from autoflow.chunkers.base import Chunker
from autoflow.chunkers.helper import get_chunker_for_datatype
from autoflow.configs.knowledge_base import IndexMethod
from autoflow.data_types import DataType, guess_datatype
from autoflow.knowledge_base.prompts import QA_WITH_KNOWLEDGE_PROMPT_TEMPLATE
from autoflow.knowledge_graph.index import KnowledgeGraphIndex
from autoflow.loaders.base import Loader
from autoflow.loaders.helper import get_loader_for_datatype
from autoflow.models.llms import LLM
from autoflow.models.embedding_models import EmbeddingModel
from autoflow.models.llms.dspy import get_dspy_lm_by_llm
from autoflow.models.rerank_models import RerankModel
from autoflow.types import BaseComponent, SearchMode
from autoflow.storage.doc_store import DocumentSearchResult, Document

logger = logging.getLogger(__name__)


class KnowledgeBase(BaseComponent):
    _llm: LLM = PrivateAttr()
    namespace: Optional[str] = Field(default=None)
    name: Optional[str] = Field(default=None)
    description: Optional[str] = Field(default=None)
    index_methods: List[IndexMethod] = Field(default=[IndexMethod.VECTOR_SEARCH])

    def __init__(
        self,
        db_engine: Engine = None,
        namespace: Optional[str] = None,
        name: Optional[str] = None,
        description: Optional[str] = None,
        index_methods: Optional[List[IndexMethod]] = None,
        llm: Optional[LLM] = None,
        embedding_model: Optional[EmbeddingModel] = None,
        rerank_model: Optional[RerankModel] = None,
        max_workers: Optional[int] = None,
    ):
        super().__init__(
            namespace=namespace,
            name=name,
            description=description,
            index_methods=index_methods,
        )
        self._db_engine = db_engine
        self._llm = llm
        self._embedding_model = embedding_model
        self._reranker_model = rerank_model
        self._init_stores()
        self._init_indexes()
        self._max_workers = max_workers or cpu_count()

    def _init_stores(self):
        from autoflow.storage.doc_store.tidb_doc_store import TiDBDocumentStore
        from autoflow.storage.graph_store.tidb_graph_store import TiDBGraphStore
        from pytidb import TiDBClient

        self._tidb_client = TiDBClient(self._db_engine)
        self._doc_store = TiDBDocumentStore(
            client=self._tidb_client,
            embedding_model=self._embedding_model,
            namespace=self.namespace,
        )
        self._kg_store = TiDBGraphStore(
            client=self._tidb_client,
            embedding_model=self._embedding_model,
            namespace=self.namespace,
        )

    def _init_indexes(self):
        self._dspy_lm = get_dspy_lm_by_llm(self._llm)
        self._kg_index = KnowledgeGraphIndex(
            kg_store=self._kg_store,
            dspy_lm=self._dspy_lm,
            embedding_model=self._embedding_model,
        )

    def class_name(self):
        return "KnowledgeBase"

    def documents(self):
        return self._doc_store;

    def knowledge_graph(self):
        return self._kg_store

    def add(
        self,
        source: str | list[str],
        data_type: Optional[DataType] = None,
        loader: Optional[Loader] = None,
        chunker: Optional[Chunker] = None,
    ) -> List[Document]:
        if data_type is None:
            data_type = guess_datatype(source)
        if data_type is None:
            raise ValueError("Please provide a valid data type.")

        if loader is None:
            loader = get_loader_for_datatype(data_type)

        with ThreadPoolExecutor(max_workers=self._max_workers) as executor:
            build_index_for_document = partial(
                self.build_index_for_document, chunker=chunker
            )

            results = executor.map(build_index_for_document, loader.load(source))

        return_documents = []
        for result in results:
            return_documents.append(result)
        return return_documents

    def build_index_for_document(
        self,
        document: Document,
        chunker: Optional[Chunker] = None,
    ) -> List[Document]:
        """
        Build index for a document.

        Args:
            document: The document to build index for.
            chunker: The chunker to use to chunk the document.

        Returns:
            A list of documents that are the result of indexing the original document.
        """
        # TODO: handle duplicate documents.
        if chunker is None:
            chunker = get_chunker_for_datatype(document.data_type)

        chunked_document = chunker.chunk(document)
        self.add_document(chunked_document)

        if IndexMethod.KNOWLEDGE_GRAPH in self.index_methods:

            def add_chunk_to_kg(chunk):
                logger.info("Adding chunk <id: %s> to knowledge graph.", chunk.id)
                self._kg_index.add_chunk(chunk)

            with ThreadPoolExecutor(max_workers=self._max_workers) as executor:
                list(executor.map(add_chunk_to_kg, chunked_document.chunks))

        return chunked_document

    # Document management.

    def add_document(self, document: Document):
        self._doc_store.add([document])

    def add_documents(self, documents: List[Document]):
        return self._doc_store.add(documents)

    def list_documents(self) -> List[Document]:
        return self._doc_store.list()

    def get_document(self, doc_id: uuid.UUID) -> Document:
        return self._doc_store.get(doc_id)

    def delete_document(self, doc_id: uuid.UUID) -> None:
        return self._doc_store.delete(doc_id)

    # Search

    def search(self):
        # TODO: Support one interface search documents and knowledge graph at the same time.
        raise NotImplementedError()

    def search_documents(
        self,
        query: str,
        mode: SearchMode = "vector",
        similarity_threshold: Optional[float] = None,
        num_candidate: Optional[int] = None,
        top_k: Optional[int] = 5,
        **kwargs: Any,
    ) -> DocumentSearchResult:
        return self._doc_store.search(
            query=query,
            mode=mode,
            similarity_threshold=similarity_threshold,
            num_candidate=num_candidate,
            top_k=top_k,
            **kwargs,
        )

    def search_knowledge_graph(
        self,
        query: str,
        depth: int = 2,
        metadata_filters: Optional[dict] = None,
        **kwargs,
    ):
        return self._kg_index.retrieve(
            query=query,
            depth=depth,
            metadata_filters=metadata_filters,
            **kwargs,
        )

    # Generation.

    def ask(self, question: str) -> ChatResponse:
        result = self.search_documents(
            query=question,
            similarity_threshold=0.4,
            top_k=5,
        )
        chunks = result.chunks
        knowledge_graph = self.search_knowledge_graph(
            query=question,
        )
        messages = QA_WITH_KNOWLEDGE_PROMPT_TEMPLATE.format_messages(
            llm=self._llm,
            query_str=question,
            chunks=chunks,
            knowledge_graph=knowledge_graph,
        )
        return self._llm.chat(messages)

    # Knowledge Base Operation.

    def reset(self):
        self._doc_store.reset()
        self._kg_store.reset()
