from typing import Any, Dict, List, Optional, Sequence, Union

from langchain_core.documents import Document

from alayaflow.component.model import ModelManager
from alayaflow.utils.logger import AlayaFlowLogger

logger = AlayaFlowLogger()


class JinaRerankerComponent:
    def __init__(self, *, model_id: str, top_n: Optional[int] = None,):
        self.model_id = model_id
        self.top_n = top_n
        self._model_manager = ModelManager()
    
    def _to_documents(
        self,
        docs: Sequence[Union[Document, Dict[str, Any], str]]
    ) -> List[Document]:
        out: List[Document] = []
        for d in docs:
            if isinstance(d, Document):
                out.append(d)
            elif isinstance(d, dict):
                out.append(
                    Document(
                        page_content=str(d.get("page_content") or d.get("content") or ""),
                        metadata=dict(d.get("metadata") or {}),
                    )
                )
            else:
                out.append(Document(page_content=str(d), metadata={}))
        return out
    
    def _get_reranker(self) -> Any:
        bind_kwargs: Dict[str, Any] = {}
        if self.top_n is not None:
            bind_kwargs["top_n"] = int(self.top_n)
        return self._model_manager.get_model(self.model_id, runtime_config=bind_kwargs)
    
    def __call__(
        self,
        *,
        query: str,
        docs: Sequence[Union[Document, Dict[str, Any], str]],
    ) -> List[Document]:
        reranker = self._get_reranker()

        documents = self._to_documents(docs)

        if not query or not documents:
            logger.debug(
                "JinaReranker short-circuit.\n"
                f"query_empty={not bool(query)} docs_empty={not bool(documents)}"
            )
            return documents

        reranked_docs: List[Document] = reranker.compress_documents(
            documents=documents,
            query=query,
        )

        logger.debug(
            "Invoke JinaReranker.\n"
            f"reranker: {self.model_id}\n"
            f"top_n: {self.top_n}\n"
            f"query: {query}\n"
            f"in_docs: {len(documents)}\n"
            f"out_docs: {len(reranked_docs)}"
        )

        return reranked_docs
    

if __name__ == "__main__":
    from alayaflow.api import Flow
    import os
    from dotenv import load_dotenv
    load_dotenv()
    flow = Flow()
    flow.init({
        "alayahub_url": "http://your-alayahub-url",
    })
    jina_api_key = os.getenv("JINA_API_KEY")
    if not jina_api_key:
        raise ValueError("请设置环境变量 JINA_API_KEY")
    
    flow.register_models([
        {
            "name": "Jina Rerank",
            "model_id": "jina-reranker",
            "provider_name": "Jina",
            "model_type": "JinaRerank",
            "model_name": "jina-reranker-v3",
            "base_url": "null",
            "api_key": jina_api_key,
        }
    ])

    jina_reranker = JinaRerankerComponent(model_id="jina-reranker", top_n=3)
    query = "如何冲泡一杯好喝的绿茶？"
    docs = [
        "冲泡美味绿茶的关键是控制水温。最好使用80-85摄氏度的热水，避免使用沸腾的100度开水，以免烫伤茶叶导致苦涩。先将茶叶放入杯中，注入热水，浸泡约2-3分钟后即可饮用。",
        "绿茶的口感与水质、茶具都有关系。建议选用软水（如矿泉水或过滤水），并使用陶瓷或玻璃杯具冲泡，这样能更好地激发绿茶的清香。",
        "绿茶的种类繁多，如龙井、碧螺春、毛峰等。不同产地的绿茶风味各异，建议多尝试找到自己喜欢的品种。",
        "饮茶是中国传统文化的重要组成部分，已有数千年的历史。茶道不仅关乎饮品，更是一种修身养性的艺术。",
        "咖啡因对人体有提神作用，但过量摄入可能导致心悸。相比之下，绿茶中的茶氨酸能带来更温和的清醒感。",
    ]
    result = jina_reranker(query=query, docs=docs)
    print(result)

