from typing import TypedDict

from sqlalchemy.orm import Session

from core.model_manager import ModelInstance, ModelManager
from core.rag.data_post_processor.reorder import ReorderRunner
from core.rag.index_processor.constant.query_type import QueryType
from core.rag.models.document import Document
from core.rag.rerank.entity.weight import KeywordSetting, VectorSetting, Weights
from core.rag.rerank.rerank_base import BaseRerankRunner
from core.rag.rerank.rerank_factory import RerankRunnerFactory
from core.rag.rerank.rerank_type import RerankMode
from extensions.otel import trace_span
from graphon.model_runtime.entities.model_entities import ModelType
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError


class RerankingModelDict(TypedDict):
    reranking_provider_name: str
    reranking_model_name: str


class VectorSettingDict(TypedDict):
    vector_weight: float
    embedding_provider_name: str
    embedding_model_name: str


class KeywordSettingDict(TypedDict):
    keyword_weight: float


class WeightsDict(TypedDict):
    vector_setting: VectorSettingDict
    keyword_setting: KeywordSettingDict


class DataPostProcessor:
    """Interface for data post-processing document."""

    def __init__(
        self,
        tenant_id: str,
        reranking_mode: str,
        reranking_model: RerankingModelDict | None = None,
        weights: WeightsDict | None = None,
        reorder_enabled: bool = False,
        *,
        session: Session,
    ):
        self.rerank_runner = self._get_rerank_runner(
            reranking_mode, tenant_id, reranking_model, weights, session=session
        )
        self.reorder_runner = self._get_reorder_runner(reorder_enabled)

    @trace_span()
    def invoke(
        self,
        query: str,
        documents: list[Document],
        score_threshold: float | None = None,
        top_n: int | None = None,
        query_type: QueryType = QueryType.TEXT_QUERY,
    ) -> list[Document]:
        if self.rerank_runner:
            documents = self.rerank_runner.run(query, documents, score_threshold, top_n, query_type)

        if self.reorder_runner:
            documents = self.reorder_runner.run(documents)

        return documents

    def _get_rerank_runner(
        self,
        reranking_mode: str,
        tenant_id: str,
        reranking_model: RerankingModelDict | None = None,
        weights: WeightsDict | None = None,
        *,
        session: Session,
    ) -> BaseRerankRunner | None:
        if reranking_mode == RerankMode.WEIGHTED_SCORE and weights:
            runner = RerankRunnerFactory.create_rerank_runner(
                runner_type=reranking_mode,
                tenant_id=tenant_id,
                weights=Weights(
                    vector_setting=VectorSetting(
                        vector_weight=weights["vector_setting"]["vector_weight"],
                        embedding_provider_name=weights["vector_setting"]["embedding_provider_name"],
                        embedding_model_name=weights["vector_setting"]["embedding_model_name"],
                    ),
                    keyword_setting=KeywordSetting(
                        keyword_weight=weights["keyword_setting"]["keyword_weight"],
                    ),
                ),
            )
            return runner
        elif reranking_mode == RerankMode.RERANKING_MODEL:
            rerank_model_instance = self._get_rerank_model_instance(tenant_id, reranking_model)
            if rerank_model_instance is None:
                return None
            runner = RerankRunnerFactory.create_rerank_runner(
                runner_type=reranking_mode,
                rerank_model_instance=rerank_model_instance,
                session=session,
            )
            return runner
        return None

    def _get_reorder_runner(self, reorder_enabled) -> ReorderRunner | None:
        if reorder_enabled:
            return ReorderRunner()
        return None

    def _get_rerank_model_instance(
        self, tenant_id: str, reranking_model: RerankingModelDict | None
    ) -> ModelInstance | None:
        if reranking_model:
            try:
                model_manager = ModelManager.for_tenant(tenant_id=tenant_id)
                reranking_provider_name = reranking_model.get("reranking_provider_name")
                reranking_model_name = reranking_model.get("reranking_model_name")
                if not reranking_provider_name or not reranking_model_name:
                    return None
                rerank_model_instance = model_manager.get_model_instance(
                    tenant_id=tenant_id,
                    provider=reranking_provider_name,
                    model_type=ModelType.RERANK,
                    model=reranking_model_name,
                )
                return rerank_model_instance
            except InvokeAuthorizationError:
                return None
        return None
