from langgraph.runtime import Runtime

from alayaflow.component.rerank_node import JinaRerankerComponent
from alayaflow.utils.logger import AlayaFlowLogger

from ..schemas import TaskState, WorkflowContext

logger = AlayaFlowLogger()


def create_rerank_node():
    RERANK_MODEL_ID= "jina-reranker"
    TOP_N = 5
    
    def rerank_node(state: TaskState, runtime: Runtime[WorkflowContext]):
        query = state.get("query")
        docs = state.get("docs", [])
        
        if not query or not docs:
            logger.debug(f"Rerank节点跳过: query为空={not bool(query)}, docs为空={not bool(docs)}")
            return {"docs": docs}
        
        model_id = RERANK_MODEL_ID
        if model_id is None:
            model_id = getattr(runtime.context, "rerank_model_id", None)
            if model_id is None:
                logger.warning("Rerank模型ID未配置，跳过rerank步骤")
                return {"docs": docs}
        
        try:
            reranker = JinaRerankerComponent(model_id=model_id, top_n=TOP_N)
            reranked_docs = reranker(query=query, docs=docs)
            
            if reranked_docs:
                if hasattr(reranked_docs[0], 'page_content'):
                    docs_list = [doc.page_content for doc in reranked_docs]
                else:
                    docs_list = [str(doc) for doc in reranked_docs]
            else:
                docs_list = []
            
            logger.debug(
                f"Rerank完成: 输入文档数={len(docs)}, "
                f"输出文档数={len(docs_list)}, model_id={model_id}"
            )
            
            return {"docs": docs_list}
        except Exception as e:
            logger.error(f"Rerank错误 {type(e).__name__}: {e}")
            return {"docs": docs}
    
    return rerank_node



