import asyncio
import json
import logging
from typing import AsyncIterable, List, Optional
from fastapi import Body, Request, HTTPException
from fastapi.concurrency import run_in_threadpool
from sse_starlette.sse import EventSourceResponse
from server.knowledge_base.kb_doc_api import search_docs
from configs import (
    TEMPERATURE,
    VECTOR_SEARCH_TOP_K,
    SCORE_THRESHOLD,
    LLM_MODELS,
    JOURNAL_KNOWLEDGE_BASE,
    MAX_TOKENS,
)
from server.utils import BaseResponse, get_prompt_template
from server.chat.policy_fun_iast import get_llm_model_response
from server.knowledge_base.kb_service.base import KBServiceFactory

async def abstract_search(
                        query: str = Body(..., description="用户输入", examples=["人工智能领域相关的文献有啥"]),
                        knowledge_base_name_list: List[str] = Body(..., description="知识库名称", examples=[["t_journal_article_bge_v1"]]),
                        # top_k: int = Body(VECTOR_SEARCH_TOP_K, description="匹配向量数", ge=1),
                        # fileName: List[str] = Body([], description="文件名称", examples=[["123.txt"]]),
                        # stream: bool = Body(False, description="流式输出"),
                        # model_name: str = Body(LLM_MODELS[0], description="LLM 模型名称。"),
                        # temperature: float = Body(TEMPERATURE, description="LLM 采样温度", ge=0.0, le=1.0),
                        # max_tokens: Optional[int] = Body(MAX_TOKENS,description="限制LLM生成Token数量，默认None代表模型最大值"),
                        # score_threshold: float = Body(SCORE_THRESHOLD,
                        #     description=(
                        #         "知识库匹配相关度阈值，取值范围在0-1之间，"
                        #         "SCORE越小，相关度越高，取到1相当于不筛选，建议设置在0.5左右"
                        #     ),
                        #     ge=0.0,
                        #     le=1.0
                        # ),
                        # request: Request = None,
                        # query_rewrite_model_name: str = Body(
                        #     'Qwen2-72B-Instruct',
                        #     description="Query rewrite model name."
                        # ),
) -> EventSourceResponse:

    async def search_iterator() -> AsyncIterable[str]:
        try:
            # 重写查询
            search_query_response = await run_in_threadpool(
                get_llm_model_response,
                strategy_name="query rewrite",
                llm_model_name=LLM_MODELS[0],
                template_prompt_name="query_rewrite",
                prompt_param_dict={"query": query, "history": []},
                temperature=0.01,
                max_tokens=512
            )

            # 解析重写后的查询
            try:
                data = json.loads(search_query_response.strip("```json\n").strip("```"))
                queries = data.get('query', [])
                search_query = ' '.join(queries)
            except json.JSONDecodeError:
                search_query = query

            logging.info(f'段落推荐输入问题: {search_query}')

            # 执行文档搜索（这里只搜索期刊知识库）
            journaldocs = await run_in_threadpool(
                search_docs,
                query=search_query,
                fileName=[],
                knowledge_base_name=JOURNAL_KNOWLEDGE_BASE,
                top_k=VECTOR_SEARCH_TOP_K,
                score_threshold=SCORE_THRESHOLD
            )
            logging.info(f"段落推荐召回资料: {journaldocs}")
            # 处理搜索结果，去重并整理
            seen_ids = set()
            journallist = []
            # journaldocs = []
            for doc in journaldocs:
                doc_id = doc.metadata.get('ID')
                content = doc.metadata.get('abstract', '').strip()
                #处理content不是完整段落的情况    
                if content:
                    if not content.endswith('。'):
                        last_period_index = content.rfind('。')
                        if last_period_index != -1:
                            content = content[:last_period_index + 1]
                        else:
                            content += '...'
                    
                    if doc_id and content and doc_id not in seen_ids:
                        journallist.append({
                            'id': doc_id,
                            'title': doc.metadata.get('title', ''),
                            'content': content
                        })
                        seen_ids.add(doc_id)

            search_results = {'journal_article': journallist}
            logging.info(f"段落推荐召回资料整理结果：{journallist}")

            # 如果没有找到相关文档，返回默认响应
            # if not journallist:
            #     search_results['journal_article'].append({
            #         'id': 'empty',
            #         'title': 'empty',
            #         'content': '暂未找到相关资料。'
            #     })

            yield json.dumps(search_results, ensure_ascii=False)

        except Exception as e:
            # 捕获并返回错误信息
            error_response = BaseResponse(code=500, msg=str(e))
            yield json.dumps(error_response.dict(), ensure_ascii=False)

    return EventSourceResponse(search_iterator())