import threading
from typing import Any

from flask import Flask, current_app
from pydantic import BaseModel, Field
from sqlalchemy import select

from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
from core.model_manager import ModelManager
from core.rag.datasource.retrieval_service import RetrievalService
from core.rag.entities.citation_metadata import RetrievalSourceMetadata
from core.rag.models.document import Document as RagDocument
from core.rag.rerank.rerank_model import RerankModelRunner
from core.rag.retrieval.retrieval_methods import RetrievalMethod
from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool
from dify_graph.model_runtime.entities.model_entities import ModelType
from extensions.ext_database import db
from models.dataset import Dataset, Document, DocumentSegment

default_retrieval_model: dict[str, Any] = {
    "search_method": RetrievalMethod.SEMANTIC_SEARCH,
    "reranking_enable": False,
    "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""},
    "top_k": 2,
    "score_threshold_enabled": False,
}


class DatasetMultiRetrieverToolInput(BaseModel):
    query: str = Field(..., description="dataset multi retriever and rerank")


class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
    """Tool for querying multi dataset."""

    name: str = "dataset_"
    args_schema: type[BaseModel] = DatasetMultiRetrieverToolInput
    description: str = "dataset multi retriever and rerank. "
    dataset_ids: list[str]
    reranking_provider_name: str
    reranking_model_name: str

    @classmethod
    def from_dataset(cls, dataset_ids: list[str], tenant_id: str, **kwargs):
        return cls(
            name=f"dataset_{tenant_id.replace('-', '_')}", tenant_id=tenant_id, dataset_ids=dataset_ids, **kwargs
        )

    def _run(self, query: str) -> str:
        threads = []
        all_documents: list[RagDocument] = []
        for dataset_id in self.dataset_ids:
            retrieval_thread = threading.Thread(
                target=self._retriever,
                kwargs={
                    "flask_app": current_app._get_current_object(),  # type: ignore
                    "dataset_id": dataset_id,
                    "query": query,
                    "all_documents": all_documents,
                    "hit_callbacks": self.hit_callbacks,
                },
            )
            threads.append(retrieval_thread)
            retrieval_thread.start()
        for thread in threads:
            thread.join()
        # do rerank for searched documents
        model_manager = ModelManager()
        rerank_model_instance = model_manager.get_model_instance(
            tenant_id=self.tenant_id,
            provider=self.reranking_provider_name,
            model_type=ModelType.RERANK,
            model=self.reranking_model_name,
        )

        rerank_runner = RerankModelRunner(rerank_model_instance)
        all_documents = rerank_runner.run(query, all_documents, self.score_threshold, self.top_k)

        for hit_callback in self.hit_callbacks:
            hit_callback.on_tool_end(all_documents)

        document_score_list = {}
        for item in all_documents:
            if item.metadata and item.metadata.get("score"):
                document_score_list[item.metadata["doc_id"]] = item.metadata["score"]

        document_context_list = []
        index_node_ids = [document.metadata["doc_id"] for document in all_documents if document.metadata]
        document_segment_stmt = select(DocumentSegment).where(
            DocumentSegment.dataset_id.in_(self.dataset_ids),
            DocumentSegment.completed_at.isnot(None),
            DocumentSegment.status == "completed",
            DocumentSegment.enabled == True,
            DocumentSegment.index_node_id.in_(index_node_ids),
        )
        segments = db.session.scalars(document_segment_stmt).all()

        if segments:
            index_node_id_to_position = {id: position for position, id in enumerate(index_node_ids)}
            sorted_segments = sorted(
                segments, key=lambda segment: index_node_id_to_position.get(segment.index_node_id, float("inf"))
            )
            for segment in sorted_segments:
                if segment.answer:
                    document_context_list.append(f"question:{segment.get_sign_content()} answer:{segment.answer}")
                else:
                    document_context_list.append(segment.get_sign_content())
            if self.return_resource:
                context_list: list[RetrievalSourceMetadata] = []
                resource_number = 1
                for segment in sorted_segments:
                    dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first()
                    document_stmt = select(Document).where(
                        Document.id == segment.document_id,
                        Document.enabled == True,
                        Document.archived == False,
                    )
                    document = db.session.scalar(document_stmt)
                    if dataset and document:
                        source = RetrievalSourceMetadata(
                            position=resource_number,
                            dataset_id=dataset.id,
                            dataset_name=dataset.name,
                            document_id=document.id,
                            document_name=document.name,
                            data_source_type=document.data_source_type,
                            segment_id=segment.id,
                            retriever_from=self.retriever_from,
                            score=document_score_list.get(segment.index_node_id),
                            doc_metadata=document.doc_metadata,
                        )

                        if self.retriever_from == "dev":
                            source.hit_count = segment.hit_count
                            source.word_count = segment.word_count
                            source.segment_position = segment.position
                            source.index_node_hash = segment.index_node_hash
                        if segment.answer:
                            source.content = f"question:{segment.content} \nanswer:{segment.answer}"
                        else:
                            source.content = segment.content
                        context_list.append(source)
                    resource_number += 1

                for hit_callback in self.hit_callbacks:
                    hit_callback.return_retriever_resource_info(context_list)

            return str("\n".join(document_context_list))
        return ""

    def _retriever(
        self,
        flask_app: Flask,
        dataset_id: str,
        query: str,
        all_documents: list,
        hit_callbacks: list[DatasetIndexToolCallbackHandler],
    ):
        with flask_app.app_context():
            stmt = select(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == dataset_id)
            dataset = db.session.scalar(stmt)

            if not dataset:
                return []

            for hit_callback in hit_callbacks:
                hit_callback.on_query(query, dataset.id)

            # get retrieval model , if the model is not setting , using default
            retrieval_model = dataset.retrieval_model or default_retrieval_model

            if dataset.indexing_technique == "economy":
                # use keyword table query
                documents = RetrievalService.retrieve(
                    retrieval_method=RetrievalMethod.KEYWORD_SEARCH,
                    dataset_id=dataset.id,
                    query=query,
                    top_k=retrieval_model.get("top_k") or 4,
                )
                if documents:
                    all_documents.extend(documents)
            else:
                if self.top_k > 0:
                    # retrieval source
                    documents = RetrievalService.retrieve(
                        retrieval_method=retrieval_model["search_method"],
                        dataset_id=dataset.id,
                        query=query,
                        top_k=retrieval_model.get("top_k") or 4,
                        score_threshold=retrieval_model.get("score_threshold", 0.0)
                        if retrieval_model["score_threshold_enabled"]
                        else 0.0,
                        reranking_model=retrieval_model.get("reranking_model", None)
                        if retrieval_model["reranking_enable"]
                        else None,
                        reranking_mode=retrieval_model.get("reranking_mode") or "reranking_model",
                        weights=retrieval_model.get("weights", None),
                    )

                    all_documents.extend(documents)
