from io import BytesIO
import logging
import os
import re
import tempfile
from fastapi import FastAPI, File, Path, UploadFile, HTTPException, Form, Request, Header
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from dotenv import load_dotenv
from pathlib import Path as PATH
from starlette.datastructures import Headers
import httpx
import pdfplumber
from markdown_process.markdown_preprocess import process_pdf_logic, convert_doc_to_pdf, convert_docx_to_pdf
from markdown_process.split_markdown import hierarchical_save, split_text, save_md, split_page_chunk
from markdown_process.chunk_postprocess import context_chunk_process
from database.milvus_api import Milvus_Database, sanitize_collection_name
from modelsAPI.model_api import OpenaiAPI, query_rewrite, prepare_rag_prompt, rerank_text, rerank_3d,resetQuery
from typing import List, Optional
from collections import defaultdict
from fuzzywuzzy import fuzz, process
from collections import defaultdict,Counter
import uuid
import hanlp
from contentprocess.contentprocess import extract_unique_nouns_from_file
from itertools import groupby
import requests
app = FastAPI(max_request_size=1024 * 1024 * 10)
load_dotenv()

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allow all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allow all HTTP methods
    allow_headers=["*"],  # Allow all headers
)
print("正在加载 HanLP 模型...")
HanLPmodel = hanlp.load(hanlp.pretrained.mtl.CLOSE_TOK_POS_NER_SRL_DEP_SDP_CON_ELECTRA_SMALL_ZH)

@app.get("/debug/routes")
async def list_routes():
    return [{"path": route.path, "methods": route.methods}
            for route in app.routes]


# class MarkdownRequest(BaseModel):
#     pdf_file: UploadFile = File(...)  # 接收Markdown格式的文本
#     database_name: Optional[str] = "JXTest"  # 接收类型列表，类型应该是 ['01', '02', '03']
#     collection_name: Optional[str] = "MarkdownTest")#

class SearchRequest(BaseModel):
    query: Optional[str] = None
    top_k: Optional[int] = 5
    database_name: Optional[str] = "JXTest"
    collection_name: Optional[str] = "RAGTest"
    query_vector: Optional[List[float]] = None


@app.post("/search")
async def search_markdown(request: SearchRequest):
    logger.info(f"Received search query: {request.query}")
    # 在这里实现搜索逻辑
    milvus_db = Milvus_Database(
        user=os.getenv("Milvus_USER", "user"),
        password=os.getenv("Milvus_PASSWORD", "password"),
        uri=os.getenv("Milvus_URI")
    )
    # query 改写
    query = query_rewrite(query=request.query)
    query_vector = OpenaiAPI.get_embeddings(list(query))
    exist_database = milvus_db.list_databases()
    if request.database_name not in exist_database:
        logger.error(f"database {request.database_name} not exist, please create it first")
        raise HTTPException(400, f" {request.database_name} not exist, please create it first")
    milvus_db.using_database(database_name=request.database_name)

    search_results = milvus_db.search(collection_name=request.collection_name, query_vector=query_vector,
                                      top_k=request.top_k)
    reranker_result_list, relevance_scores = OpenaiAPI.rerank_query(query=request.query, documents=search_results,
                                                                    top_n=request.top_k)
    rag_prompt = prepare_rag_prompt(query=request.query, documents=reranker_result_list)
    return StreamingResponse(OpenaiAPI.open_api_chat_stream(query=rag_prompt), media_type="text/plain")


class CreateCollectionRequest(BaseModel):
    collection_name: str
    database_name: str
    chunk_type: str = "01"


@app.post("/create_collection")
async def create_collection(request: CreateCollectionRequest):
    """
    创建数据库和集合
    chunk_type: 01 表示智能存储方式，02表示按照章节存储，03表示按字数存储
    """
    collection_name = request.collection_name
    collection_name = sanitize_collection_name(collection_name)
    database_name = request.database_name
    chunk_type = request.chunk_type

    logger.info(f"Received create collection request: {database_name},{collection_name}")
    milvus_db = Milvus_Database(user=os.getenv("Milvus_USER", "user"),
                                password=os.getenv("Milvus_PASSWORD", "password"), uri=os.getenv("Milvus_URI"))
    try:
        # 创建数据库与collections
        exist_database = milvus_db.list_databases()
        if database_name not in exist_database:
            res_milvus_create = milvus_db.create_database(database_name=database_name)
            logger.info(f"create database result: {res_milvus_create}")
        else:
            logger.info(f"database {database_name} already exist")

        milvus_db.using_database(database_name=database_name)
        exist_collection = milvus_db.has_collection(collection_name=collection_name, database_name=database_name)
        if not exist_collection:
            if chunk_type == "03":
                save_title = False
            else:
                save_title = True
            res_milvus_create = milvus_db.create_hybrid_collection(collection_name=collection_name,
                                                                   database_name=database_name, save_title=save_title)
            logger.info(f"create collection result: {res_milvus_create}")

            return JSONResponse(content={"message": "success", "code": "200"}, status_code=200)
        else:
            logger.info(f"collection {collection_name} already exist")
            raise HTTPException(400, f"collection {collection_name} already exist")
    except Exception as e:
        logger.error(f"Error creating collection: {str(e)}")
        raise HTTPException(500, "Error creating collection")
    finally:

        milvus_db.close()


@app.get("/has_collection")
async def collection(database_name, collection_name):
    milvus_db = Milvus_Database(user=os.getenv("Milvus_USER", "user"),
                                password=os.getenv("Milvus_PASSWORD", "password"), uri=os.getenv("Milvus_URI"))
    try:
        collection = milvus_db.has_collection(collection_name=collection_name, database_name=database_name)
        return {"message": "success", "code": "200", "data": {"has_collection": collection}}
    except Exception as e:
        logger.error(f"Error getting collections: {str(e)}")
        raise HTTPException(500, "Error getting collections")
    finally:
        milvus_db.close()


@app.delete("/delete_collection/{collection_name}")
async def delete_collection(collection_name: str, database_name: Optional[str]):
    milvus_db = Milvus_Database(user=os.getenv("Milvus_USER", "user"),
                                password=os.getenv("Milvus_PASSWORD", "password"), uri=os.getenv("Milvus_URI"))
    try:
        collection_name = sanitize_collection_name(collection_name)
        milvus_db.drop_collection(collection_name=collection_name, database_name=database_name)
        return JSONResponse(content={"message": "success", "code": "200"}, status_code=200)
    except Exception as e:
        logger.error(f"Error deleting collection: {str(e)}")
        raise HTTPException(500, "Error deleting collection")
    finally:
        milvus_db.close()


@app.get("/get_collections")
async def get_collections(database_name):
    milvus_db = Milvus_Database(user=os.getenv("Milvus_USER", "user"),
                                password=os.getenv("Milvus_PASSWORD", "password"), uri=os.getenv("Milvus_URI"))
    try:
        collections = milvus_db.list_collections(database_name=database_name)
        if collections:  # 确保列表非空
            collections.pop(0)
        return {"message": "success", "code": "200", "data": {"collections": collections}}
    except Exception as e:
        logger.error(f"Error getting collections: {str(e)}")
        raise HTTPException(500, "Error getting collections")
    finally:
        milvus_db.close()


@app.post("/upload_pdf")
async def process_pdf(pdf_file: UploadFile = File(...), database_name: Optional[str] = Form("JXTest"),
                      collection_name: Optional[str] = Form("MarkdownTest"), chunk_type: str = "01",
                      chunk_size: int = 256, chunk_overlap: int = 50):
    """
    接收 PDF 文件并转发到解析服务
    chunk_type: list[str] 01: 智能分割 ； 02: 按照章节分割； 03: 按照字数分割
    chunk_size: int 按照字数分割
    chunk_overlap: int 分块重叠
    """
    logger.info(
        f"Received PDF file: {pdf_file.filename},database_name : {database_name},collection_name : {collection_name}")
    # 验证文件类型
    if not pdf_file.content_type.startswith("application/pdf"):
        logger.error(f"Invalid file type: {pdf_file.content_type}")
        raise HTTPException(400, "仅支持 PDF 文件")
    filename = pdf_file.filename
    save_title = True
    milvus_db = Milvus_Database(user=os.getenv("Milvus_USER", "user"),
                                password=os.getenv("Milvus_PASSWORD", "password"), uri=os.getenv("Milvus_URI"))
    try:
        # 创建数据库与collections

        exist_database = milvus_db.list_databases()
        if database_name not in exist_database:
            res_milvus_create = milvus_db.create_database(database_name=database_name)
            logger.info(f"create database result: {res_milvus_create}")
        else:
            logger.info(f"database {database_name} already exist")

        milvus_db.using_database(database_name=database_name)
        exist_collection = milvus_db.has_collection(collection_name=collection_name)
        if not exist_collection:
            # 创建向量检索集合
            # res_milvus_create = milvus_db.create_collections(collection_name,drop_=False)
            # 创建混检集合
            if chunk_type == "03":
                save_title = False
            res_milvus_create = milvus_db.create_hybrid_collection(collection_name=collection_name,
                                                                   database_name=database_name, save_title=save_title)
            logger.info(f"create collection result: {res_milvus_create}")
        else:
            logger.info(f"collection {collection_name} already exist")

        md_content, content_chunk_list = await process_pdf_logic(pdf_file)
        logger.info("PDF file processed successfully")
        # 以下部分需调整
        md_header_splits = split_text(resource=filename, mdDocs=md_content, logger=logger, database_name=database_name,
                                      collection_name=collection_name, chunk_type=chunk_type, chunk_length=chunk_size,
                                      chunk_overlap=chunk_overlap)
        # md_header_splits = split_by_md_header(mdDocs=md_content,logger=logger,database_name=database_name,collection_name=collection_name)
        logger.info("Markdown content split successfully")

        return {"message": "PDF save to Milvus successfully", "code": 200,
                "data": {"md_head": md_header_splits, "md_content": md_content}}

    except httpx.HTTPStatusError as e:
        logger.error(f"HTTP status error from parsing service: {e.response.text}")
        raise HTTPException(e.response.status_code, f"解析服务返回错误: {e.response.text}")
    except httpx.RequestError as e:
        logger.error(f"Failed to connect to parsing service: {str(e)}")
        raise HTTPException(503, f"无法连接到解析服务: {str(e)}")
    except Exception as e:
        logger.error(f"Internal server error: {str(e)}")
        raise HTTPException(500, f"服务器内部错误: {str(e)}")
    finally:
        milvus_db.close()


@app.delete("/delete_resource/{resource}")
async def delete_resource(resource: str, collection_name, database_name: str, ):
    """
    删除指定数据库和集合
    """
    logger.info(
        f"delete resource : {resource}, delete database_name : {database_name},collection_name : {collection_name}")
    milvus_db = Milvus_Database(
        user=os.getenv("Milvus_USER", "user"),
        password=os.getenv("Milvus_PASSWORD", "password"),
        uri=os.getenv("Milvus_URI")
    )
    try:
        milvus_db.using_database(database_name=database_name)
        collection_name = sanitize_collection_name(collection_name)
        res = milvus_db.delete_resource(collection_name=collection_name, resource=resource, database_name=database_name)

        return {"code": 200, "message": "success", "data": {"result": res}}
    except Exception as e:
        logger.error(f"Delete failed: {str(e)}", exc_info=True)
        raise HTTPException(
            status_code=500,
            detail=f"删除失败: {str(e)}"
        )
    finally:
        milvus_db.close()


@app.post("/get_markdown")
async def get_markdown(file: UploadFile = File(...)):
    filename = file.filename
    file_ext = os.path.splitext(filename)[1].lower()
    headers = Headers({"content-type": "application/pdf"})
    if file.content_type.startswith("application/pdf"):
        logger.info(f"Received file: {filename}")
        md_content, content_chunk_list = await process_pdf_logic(file)
        logger.info("PDF file processed successfully")
        if not md_content:
            content = await file.read()
            await file.seek(0)
            md_content = ""
            try:
                with pdfplumber.open(BytesIO(content)) as pdf:
                    for page in pdf.pages:
                        md_content += page.extract_text() or ""
            except Exception as e:
                logger.error(f"Failed to parse PDF: {e}")
                raise HTTPException(
                    status_code=400,
                    detail="PDF file read failed"
                )
    # elif file.content_type == EXT_TO_CONTENT_TYPE[".docx"]:
    elif file_ext == ".docx":
        # 处理DOCX文件
        converted_path = await convert_docx_to_pdf(os.getenv("CONVERT_DOC_URL", "http://192.168.0.46:59070"), file)

        # 读取转换后的DOCX文件
        with open(converted_path, 'rb') as f:
            converted_file = UploadFile(filename=f"{PATH(filename).stem}.pdf", file=f, headers=headers)
            md_content, content_chunk_list = await process_pdf_logic(converted_file)
        os.unlink(converted_path)
        if not md_content:
            content = await file.read()
            await file.seek(0)
            md_content = ""
            try:
                with pdfplumber.open(BytesIO(content)) as pdf:
                    for page in pdf.pages:
                        md_content += page.extract_text() or ""
            except Exception as e:
                logger.error(f"Failed to parse PDF: {e}")
                raise HTTPException(
                    status_code=400,
                    detail="PDF file read failed")
    # elif file.content_type == EXT_TO_CONTENT_TYPE[".doc"]:
    elif file_ext == ".doc":
        # 处理DOC文件，先转换为DOCX
        try:
            logger.info(f"Converting DOC file: {filename}")
            converted_path = await convert_doc_to_pdf(os.getenv("CONVERT_DOC_URL", "http://192.168.0.46:59070"), file)

            # 读取转换后的DOCX文件
            with open(converted_path, 'rb') as f:
                converted_file = UploadFile(filename=f"{PATH(filename).stem}.pdf", file=f, headers=headers)
                md_content, content_chunk_list = await process_pdf_logic(converted_file)

            # 获取图片信息
            img_information = [chunk for chunk in content_chunk_list if chunk.get("type") == "image"]

            # 清理临时文件
            os.unlink(converted_path)
            if not md_content:
                content = await file.read()
                await file.seek(0)
                md_content = ""
                try:
                    with pdfplumber.open(BytesIO(content)) as pdf:
                        for page in pdf.pages:
                            md_content += page.extract_text() or ""
                except Exception as e:
                    logger.error(f"Failed to parse PDF: {e}")
                    raise HTTPException(
                        status_code=400,
                        detail="PDF file read failed")

        except Exception as e:
            logger.error(f"Failed to convert DOC: {e}")
            raise HTTPException(
                status_code=400,
                detail=f"DOC conversion failed: {str(e)}"
            )
    else:
        raise HTTPException(
            status_code=400,
            detail="Unsupported file type. Only PDF, DOC and DOCX are supported."
        )
    img_information = [chunk for chunk in content_chunk_list if chunk.get("type") == "image"]

    return {
        "message": "success",
        "code": 200,
        "data": {
            "markdown_content": md_content,
            "content_chunk_list": content_chunk_list,
            "img_information": img_information
        }
    }


class SplitContextRequest(BaseModel):
    file_name: str = "测试文件"
    context: Optional[str] = "测试文本"
    chunk_type: str = "01"
    chunk_size: int = 1024
    chunk_overlap: int = 100
    max_level: Optional[int] = 3
    chunk_list: List[dict] = []


@app.post("/split_result")
async def split_result(file: UploadFile = File(...),
                       chunk_type: str = Form("01"),
                       chunk_size: int = Form(1024),
                       chunk_overlap: int = Form(100),
                       max_level: Optional[int] = Form(3),
                       chunk_list: Optional[List[dict]] = Form([])
                       ):
    """
    拆分结果
    """
    filename = file.filename
    file_ext = os.path.splitext(filename)[1]
    headers = Headers({"content-type": "application/pdf"})
    if file.content_type.startswith("application/pdf"):
        logger.info(f"Received file: {filename}")
        md_content, content_chunk_list = await process_pdf_logic(file)
        logger.info("PDF file processed successfully")
    elif file_ext == ".docx":
        # 处理DOCX文件
        converted_path = await convert_docx_to_pdf(os.getenv("CONVERT_DOC_URL", "http://192.168.0.46:59070"), file)

        # 读取转换后的DOCX文件
        with open(converted_path, 'rb') as f:
            converted_file = UploadFile(filename=f"{PATH(filename).stem}.pdf", file=f, headers=headers)
            md_content, content_chunk_list = await process_pdf_logic(converted_file)
        os.unlink(converted_path)
    elif file_ext == ".docx":
        # 处理DOC文件，先转换为DOCX
        try:
            logger.info(f"Converting DOC file: {filename}")
            converted_path = await convert_doc_to_pdf(os.getenv("CONVERT_DOC_URL", "http://192.168.0.46:59070"), file)

            # 读取转换后的DOCX文件
            with open(converted_path, 'rb') as f:
                converted_file = UploadFile(filename=f"{PATH(filename).stem}.pdf", file=f, headers=headers)
                md_content, content_chunk_list = await process_pdf_logic(converted_file)

            # 清理临时文件
            os.unlink(converted_path)

        except Exception as e:
            logger.error(f"Failed to convert DOC: {e}")
            raise HTTPException(
                status_code=400,
                detail=f"DOC conversion failed: {str(e)}"
            )
    else:
        raise HTTPException(
            status_code=400,
            detail="Unsupported file type. Only PDF, DOC and DOCX are supported."
        )
    if not md_content:
        content = await file.read()
        await file.seek(0)
        md_content = ""
        try:
            with pdfplumber.open(BytesIO(content)) as pdf:
                for page in pdf.pages:
                    md_content += page.extract_text() or ""
        except Exception as e:
            logger.error(f"Failed to parse PDF: {e}")
            raise HTTPException(
                status_code=400,
                detail="PDF file read failed")
    if chunk_list:
        md_header_splits = split_page_chunk(
            resource=filename,
            content_chunk_list=content_chunk_list,
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            logger=logger,
            max_level=max_level,
            chunk_type=chunk_type
        )
    else:
        md_header_splits = split_text(
            resource=filename,
            mdDocs=md_content,
            logger=logger,
            chunk_type=chunk_type,
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            max_level=max_level
        )
    return {
        "message": "success",
        "code": 200,
        "data": {
            "split_result": md_header_splits,
        }
    }


@app.post("/split_context")
async def split_context(request: SplitContextRequest):
    """
    拆分结果
    """
    filename = request.file_name
    context = request.context
    chunk_type = request.chunk_type
    chunk_size = request.chunk_size
    chunk_overlap = request.chunk_overlap
    max_level = request.max_level
    chunk_list = request.chunk_list
    if chunk_list:
        md_header_splits = split_page_chunk(
            resource=filename,
            content_chunk_list=chunk_list,
            chunk_type=chunk_type,
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            logger=logger,
            max_level=max_level
        )
    else:
        if context:
            md_header_splits = split_text(resource=filename, mdDocs=context,
                                          logger=logger, chunk_type=chunk_type,
                                          chunk_size=chunk_size, chunk_overlap=chunk_overlap,
                                          max_level=max_level)
    return {"message": "sucess", "code": 200, "data": {"split_result": md_header_splits}}


@app.post("/split_context_chunk")
async def split_context_chunk(request: SplitContextRequest):
    """
    拆分结果
    """
    filename = request.file_name
    context = request.context
    chunk_type = request.chunk_type
    chunk_size = request.chunk_size
    chunk_overlap = request.chunk_overlap
    max_level = request.max_level
    chunk_list = request.chunk_list
    md_header_splits = split_page_chunk(
        resource=filename,
        content_chunk_list=chunk_list,
        chunk_type=chunk_type,
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        logger=logger,
        max_level=max_level
    )
    context_chunk_splits = context_chunk_process(markdown_context=context, chunks=md_header_splits)
    # context_chunk_splits = OpenaiAPI().context_chunk_process(markdown_context=context,chunks=md_header_splits)
    return {"message": "sucess", "code": 200, "data": {"split_result": context_chunk_splits}}


class InsertRequest(BaseModel):
    split_list: list
    collection_name: str
    database_name: str  # 默认值


@app.post("/insert_md")
async def insert_md(request: InsertRequest):
    """
    插入md
    """
    collection_name = request.collection_name
    logger.info(f"get {collection_name}")
    database_name = request.database_name
    split_list = request.split_list
    milvus_db = Milvus_Database(
        user=os.getenv("Milvus_USER", "user"),
        password=os.getenv("Milvus_PASSWORD", "password"),
        uri=os.getenv("Milvus_URI")
    )
    for split_result in split_list:
        reource = split_result.get("resource")
        text = milvus_db.get_text_by_resource(collection_name=collection_name, database_name=database_name,
                                              resource=reource, logger=logger)
        if text:
            logger.info("已经存在相同的内容")
            return {"code": 200, "message": "success", "data": {"result": split_list}}
        else:
            logger.info("内容不在库内，跳过后续查询")
            break
    try:
        # if SPLIT_LIST:
        #     logger.info(f"Inserting {SPLIT_LIST}")
        res = save_md(results=request.split_list, collection_name=collection_name, database_name=request.database_name)
        # res = save_md(results=SPLIT_LIST,collection_name=collection_name,database_name=request.database_name)
        return {"code": 200, "message": "success", "data": {"result": res}}
        # else:
        #     return {"code":200,"message":"success","data":{"result":split_list}}
    except Exception as e:
        logger.error(f"Insert failed: {str(e)}", exc_info=True)
        raise HTTPException(400, f"插入失败: {str(e)}")


@app.post("/hierarchical_insert")
async def hierarchical_insert(request: InsertRequest):
    collection_name = request.collection_name
    logger.info(f"get {collection_name}")
    database_name = request.database_name
    split_list = request.split_list
    milvus_db = Milvus_Database(
        user=os.getenv("Milvus_USER", "user"),
        password=os.getenv("Milvus_PASSWORD", "password"),
        uri=os.getenv("Milvus_URI")
    )
    for split_result in split_list:
        reource = split_result.get("resource")
        text = milvus_db.get_text_by_resource(collection_name=collection_name, database_name=database_name,
                                              resource=reource, logger=logger)
        if text:
            logger.info("已经存在相同的内容")
            return {"code": 200, "message": "success", "data": {"result": split_list}}
        else:
            logger.info("内容不在库内，跳过后续查询")
            break
    try:
        res = hierarchical_save(results=request.split_list, collection_name=collection_name,
                                database_name=request.database_name)
        return {"code": 200, "message": "success", "data": {"result": res}}
        # else:
        #     return {"code":200,"message":"success","data":{"result":split_list}}
    except Exception as e:
        logger.error(f"Insert failed: {str(e)}", exc_info=True)
        raise HTTPException(400, f"插入失败: {str(e)}")


@app.post("/search_source")
async def search_source(request: SearchRequest):
    """
    搜索与给定查询最相关的 Markdown 文档
    """
    logger.info(f"Received search query: {request.query},collection_name: {request.collection_name}")

    # 初始化 Milvus 数据库
    milvus_db = Milvus_Database(
        user=os.getenv("Milvus_USER", "user"),
        password=os.getenv("Milvus_PASSWORD", "password"),
        uri=os.getenv("Milvus_URI")
    )

    # 获取查询的向量表示
    # query = query_rewrite(request.query)
    query = request.query
    query = resetQuery(query)
    logger.info(f"改写query:{query}")
    query_vector = OpenaiAPI.get_embeddings(list(query))

    collection_name = request.collection_name
    final_results = []
    # 搜索相关文档
    search_results, search_resources, recall_results = milvus_db.hybrid_search(
        collection_name=collection_name,
        database_name=request.database_name,
        query_vector=query_vector,
        top_k=request.top_k,
        query=query,
        embedding_rate=0.4
    )
    logger.info(f"文本混合检索结果: {recall_results}")
    if not recall_results:
        search_results, search_resources, recall_results = milvus_db.embedding_search(
            collection_name=collection_name,
            query_vector=query_vector,
            database_name=request.database_name,
            top_k=request.top_k
        )
        logger.info(f"文本纯向量检索结果: {search_results}")

    for recall_result in recall_results:
        text_1 = recall_result["text"]
        text_2 = query
        score = OpenaiAPI.embedding_cosine_similarity(text_1, text_2)
        recall_result["score"] = score
    # recall_results = sorted(recall_results, key=lambda x: x["score"], reverse=True)
    # for idx, item in enumerate(recall_results, start=1):
    #   item["index"] = idx
    logger.info("------------------------------------")
    logger.info(f"第一轮精排结果: {recall_results}")
    recall_results = rerank_text(query, recall_results)
    recall_results = sorted(recall_results, key=lambda x: x["score"], reverse=True)
    for idx, item in enumerate(recall_results, start=1):
        item["index"] = idx
    threshold = 0.2
    recall_results = [item for item in recall_results if item.get('score', 0) > threshold]
    logger.info("------------------------------------")
    logger.info(f"第二轮精排结果: {recall_results}")

    return {"code": 200, "message": "success", "data": recall_results, "query_vector": query_vector}

# 图谱检索服务地址（可通过环境变量覆盖）
GRAPH_SEARCH_URL = os.getenv("GRAPH_SEARCH_URL", "http://192.168.0.46:59085/search_graph")


async def call_graph_search(query: str, top_k: int = 10, graph_timeout: float = 30.0):
    """
    异步调用图谱检索接口，并解析返回的 data.results
    """
    if not query:
        return []

    try:
        # 使用较长超时时间，避免复杂图谱查询被过早中断
        async with httpx.AsyncClient(timeout=graph_timeout) as client:
            resp = await client.post(
                GRAPH_SEARCH_URL,
                json={
                    "query": query,
                    "top_k": top_k,
                    "graph_timeout": graph_timeout,
                    # 如需入口节点可在这里扩展： "entry_nodes": ...
                },
            )
            resp.raise_for_status()
            result = resp.json()
    except httpx.RequestError as e:
        # 打印更详细的错误信息，便于排查（连接失败 / 读超时等）
        logger.error(
            f"[call_graph_search] 请求图谱检索服务失败: {e.__class__.__name__} - {e}",
            exc_info=True
        )
        return []
    except httpx.HTTPStatusError as e:
        logger.error(f"[call_graph_search] 图谱检索服务返回异常状态码: {e.response.status_code}, body={e.response.text}")
        return []
    except Exception as e:
        logger.error(f"[call_graph_search] 解析图谱检索结果失败: {e}", exc_info=True)
        return []

    # 按照约定结构提取 data.results
    try:
        data = result.get("data", {})
        if isinstance(data, dict):
            graph_results = data.get("results", [])
        else:
            graph_results = []
        if not isinstance(graph_results, list):
            graph_results = []
        return graph_results
    except Exception as e:
        logger.error(f"[call_graph_search] 提取 data.results 失败: {e}", exc_info=True)
        return []



"""
    reranker_result_list =  OpenaiAPI.rerank_query(
        query=query,
        documents=search_results,
        top_n=request.top_k
    )
    # print("###########")
    # print("reranker_result_list:",reranker_result_list)


    for res in reranker_result_list:
        index = res["index"]
        resource,img_path = search_resources[index]["resource"],search_resources[index]["img_path"]
        res["resource"] = resource
        res["img_path"] = img_path

    #内容去重
    # 基于text字段去重
    seen_texts = set()
    unique_results = []
    for res in reranker_result_list:
        if res["text"] not in seen_texts:
            seen_texts.add(res["text"])
            unique_results.append(res)
    milvus_db.close()
    return {"code": 200, "message": "success", "data": unique_results}
"""
class SearchBatchRequest(BaseModel):
    queries: Optional[List[str]] = None  # 新增：支持批量文本
    query_vectors: Optional[List[List[float]]] = None  # 新增：支持批量向量
    top_k: Optional[int] = 5
    database_name: Optional[str] = "XIAN"
    collection_name: Optional[str] = "RAGTest"

@app.post("/search_batch")
async def search_source(request: SearchBatchRequest):
    """
    搜索与给定查询最相关的 Markdown 文档
    """
    # ✅ 修复：使用 request.queries 而不是 request.query
    logger.info(f"Received search queries: {request.queries}, collection_name: {request.collection_name}")

    milvus_db = Milvus_Database(
        user=os.getenv("Milvus_USER", "user"),
        password=os.getenv("Milvus_PASSWORD", "password"),
        uri=os.getenv("Milvus_URI")
    )

    # 处理文本或向量输入
    if request.queries:
        query_texts = request.queries
        try:
            result = OpenaiAPI.batch_embeddings(query_texts)
            query_vector = [item['embedding'] for item in result]
        except Exception as e:
            logger.error(f"Embedding failed: {e}")
            raise HTTPException(status_code=500, detail="Failed to generate embeddings")
    elif request.query_vectors:
        query_texts = request.queries  # 可能为空
        query_vector = request.query_vectors
    else:
        raise HTTPException(status_code=400, detail="You must provide either 'queries' or 'query_vectors'")

    # 执行混合搜索
    search_results, search_resources, recall_results = milvus_db.hybrid_search_batch(
        collection_name=request.collection_name,
        database_name=request.database_name,
        query_vectors=query_vector,
        top_k=request.top_k,
        queries=query_texts,  # 可为 None
        embedding_rate=0.1
    )
    logger.info(f"3D 模型混合检索结果: {recall_results}")

    return {"code": 200, "message": "success", "data": recall_results}

class Search3DtextRequest(BaseModel):
    content : str
    search3d_result : list
    top_k: Optional[int] = 5
    database_name: Optional[str] = "XIAN"
    collection_name: Optional[str] = "b9d47d5a6d66cbe2cd44a8ea4ad021fc"


@app.post("/search_3Dtext")
async def search_source(request: Search3DtextRequest):
    content = request.content
    search3d_result = request.search3d_result
    model = HanLPmodel
    allowed_resources = list(set([item['resource'] for item in search3d_result]))
    SCORE_THRESHOLD = 0.3
    logger.info(f"Received search content: {content}, collection_name: {request.collection_name}")

    milvus_db = Milvus_Database(
        user=os.getenv("Milvus_USER", "user"),
        password=os.getenv("Milvus_PASSWORD", "password"),
        uri=os.getenv("Milvus_URI")
    )
    # 处理文本或向量输入
    if content:
        query_texts = extract_unique_nouns_from_file(content,model=model)
        try:
            result = OpenaiAPI.batch_embeddings(query_texts)
            query_vector = [item['embedding'] for item in result]
        except Exception as e:
            logger.error(f"Embedding failed: {e}")
            raise HTTPException(status_code=500, detail="Failed to generate embeddings")
    else:
        raise HTTPException(status_code=400, detail="You must provide either 'content' ")

    # 执行混合搜索
    search_results, search_resources, recall_results = milvus_db.hybrid_search_batch(
        collection_name=request.collection_name,
        database_name=request.database_name,
        query_vectors=query_vector,
        top_k=request.top_k,
        queries=query_texts,  # 可为 None
        embedding_rate=0.1
    )
    search_result = []
    for query_results in recall_results:
        filtered_query_results = [
            item for item in query_results
            if item["score"] >= SCORE_THRESHOLD and item["resource"] in allowed_resources
        ]
        search_result.append(filtered_query_results)
    search_result = [sublist for sublist in search_result if sublist]
   # for item in search_result[0]:
    #    if item["score"] >= 0.9:
     #       continue
      #  else:
       #     item["score"] += 0.1
    if search_result:
        for item in search_result[0]:
            if item["score"] < 0.9:
                item["score"] += 0.1

    logger.info(f"大模型生成内容零部件3D检索召回:{search_result}")
    if search_result:
        return {"code": 200, "message": "success", "data": search_result}
    else:
        return {"code": 200, "message": "success", "data": []}

@app.post("/search_3D")
async def search_3D(request: SearchRequest):
    """
    搜索与给定查询最相关的 Markdown 文档
    """
    logger.info(f"Received search query: {request.query},collection_name: {request.collection_name}")
    # 初始化 Milvus 数据库
    milvus_db = Milvus_Database(
        user=os.getenv("Milvus_USER", "user"),
        password=os.getenv("Milvus_PASSWORD", "password"),
        uri=os.getenv("Milvus_URI")
    )

    # 获取查询的向量表示
    # query = query_rewrite(request.query)
    if request.query:
        query = request.query
        # query_vector =  OpenaiAPI.get_embeddings(list(query))
    # else:
    #     query_vector = request.query_vector
    if request.query_vector:
        query_vector = request.query_vector
    else:
        query_vector = OpenaiAPI.get_embeddings(list(query))

    # 检查数据库是否存在
    # exist_database = milvus_db.list_databases()
    # if request.database_name not in exist_database:
    #     logger.error(f"database {request.database_name} not exist, please create it first")
    #     raise HTTPException(400, f" {request.database_name} not exist, please create it first")

    # 使用指定数据库
    # milvus_db.using_database(database_name=request.database_name)
    # collection_name = sanitize_collection_name(request.collection_name)
    collection_name = request.collection_name
    search_results, search_resources, recall_results = milvus_db.hybrid_search(
        collection_name=collection_name,
        database_name=request.database_name,
        query_vector=query_vector,
        top_k=request.top_k,
        query=request.query,
        embedding_rate=0.1
    )
    logger.info(f"3D 模型混合检索结果: {recall_results}")

    threshold = 0.2
    logger.info(f"阈值:{threshold}")

    search_result = [item for item in recall_results if item.get('score', 0) > threshold]
    logger.info(f"精排召回第一轮:{search_result}")
    search_result = rerank_3d(query, search_result)
    search_result = [item for item in search_result if item.get('score', 0) > threshold]
    logger.info(f"精排召回第二轮:{search_result}")
    if search_result:
        return {"code": 200, "message": "success", "data": search_result}
    else:
        return {"code": 200, "message": "success", "data": []}


class SearchChunkRequest(BaseModel):
    resource: str
    database_name: str
    collection_name: str


@app.post("/search_chunk")
async def search_chunk(request: SearchChunkRequest):
    milvus_db = Milvus_Database(user=os.getenv("Milvus_USER", "user"),
                                password=os.getenv("Milvus_PASSWORD", "password"), uri=os.getenv("Milvus_URI"))
    results = milvus_db.get_chunk_by_resource(collection_name=request.collection_name,
                                              database_name=request.database_name, resource=request.resource)
    return {"code": 200, "message": "success", "data": results}


@app.post("/search")
async def search_markdown(request: Request, query: str, top_k: Optional[int] = 5,
                          database_name: Optional[str] = "JXTest", collection_name: Optional[str] = "RAGTest"):
    """
    搜索与给定查询最相关的 Markdown 文档
    """
    logger.info(f"Received search query: {query}")
    # 在这里实现搜索逻辑，例如使用 Milvus 进行向量搜索
    milvus_db = Milvus_Database(user=os.getenv("Milvus_USER", "user"),
                                password=os.getenv("Milvus_PASSWORD", "password"), uri=os.getenv("Milvus_URI"))
    query_vector = OpenaiAPI.get_embeddings(list(query))
    exist_database = milvus_db.list_databases()
    if database_name not in exist_database:
        logger.error(f"database {database_name} not exist, please create it first")
        raise HTTPException(400, f" {database_name} not exist, please create it first")
    # 混搜部分
    search_results = milvus_db.hybrid_search(query_vector=query_vector, query=query, top_k=top_k,
                                             collection_name=collection_name, database_name=database_name)
    reranker_result_list, relevance_scores = OpenaiAPI.rerank_query(query=query, documents=search_results, top_n=top_k)
    rag_prompt = prepare_rag_prompt(query=query, documents=reranker_result_list)
    return StreamingResponse(OpenaiAPI.open_api_chat_stream(query=rag_prompt), media_type="text/plain")


@app.post("/search_context_resource")
def search_context_resource(request: SearchRequest):
    logger.info(f"Received search query: {request.query},collection_name: {request.collection_name}")

    # 初始化 Milvus 数据库
    milvus_db = Milvus_Database(
        user=os.getenv("Milvus_USER", "user"),
        password=os.getenv("Milvus_PASSWORD", "password"),
        uri=os.getenv("Milvus_URI")
    )

    # 获取查询的向量表示
    # query = query_rewrite(request.query)
    query = request.query
    query_vector = OpenaiAPI.get_embeddings(list(query))

    # 检查数据库是否存在
    # exist_database = milvus_db.list_databases()
    # if request.database_name not in exist_database:
    #     logger.error(f"database {request.database_name} not exist, please create it first")
    #     raise HTTPException(400, f" {request.database_name} not exist, please create it first")

    # 使用指定数据库
    # milvus_db.using_database(database_name=request.database_name)
    # collection_name = sanitize_collection_name(request.collection_name)
    collection_name = request.collection_name
    final_results = []
    # 搜索相关文档
    search_results, search_resources, recall_results = milvus_db.hierarchical_search(
        collection_name=collection_name,
        database_name=request.database_name,
        query_vector=query_vector,
        top_k=request.top_k,
        query=query
    )
    logger.info(f"文本混合检索结果: {recall_results}")

    if not recall_results:
        search_results, search_resources, recall_results = milvus_db.embedding_search(
            collection_name=collection_name,
            query_vector=query_vector,
            database_name=request.database_name,
            top_k=request.top_k
        )
        logger.info(f"文本纯向量检索结果: {search_results}")

    logger.info(f"recall_results: {recall_results}")
    return {"code": 200, "message": "success", "data": recall_results, "query_vector": query_vector}


class HightLightRequest(BaseModel):
    markdown_content: str
    chunk: str


@app.post("/get_highlight_markdown")
async def get_hightlight(request: HightLightRequest):
    logger.info(f"Received chunk :{request.chunk}")

    def find_best_match(text, chunk, threshold=10):
        """返回得分最高的单个匹配位置及得分"""
        normalized_text = text
        normalized_chunk = chunk

        # 使用滑动窗口确保找到全局最佳匹配
        best_match = None
        best_score = 0

        # 窗口大小设为待匹配文本长度的2倍
        window_size = min(len(normalized_chunk), len(normalized_text))
        step = len(normalized_chunk)  # 滑动步长

        for i in range(1, len(normalized_text) - len(normalized_chunk), step):
            window = normalized_text[i:i + window_size]

            # 使用多种匹配策略组合
            current_score = 0
            current_score += fuzz.ratio(normalized_chunk, window) * 0.3
            current_score += fuzz.partial_ratio(normalized_chunk, window) * 0.4
            current_score += fuzz.token_sort_ratio(normalized_chunk, window) * 0.3

            if current_score > best_score:
                best_score = current_score
                # 记录在原始文本中的位置（而非预处理后的位置）
                raw_start = text.find(window[i:i + len(normalized_chunk)])
                raw_end = raw_start + len(normalized_chunk)
                best_match = (raw_start, raw_end, best_score)

        # 仅返回超过阈值的匹配
        return best_match if best_score >= threshold else None

    def make_flexible_pattern(chunk):
        """生成允许换行符和空格数量不一致的正则模式"""
        # 转义特殊字符（保留 \n 和空格）
        escaped = re.escape(chunk)

        # 将 \n 替换为 \s*（匹配任意空白，包括 \n 和空格）
        flexible_pattern = escaped.replace(r"\n", r"\s+")  # \s+ 确保至少一个空白

        # 进一步优化：如果原文本有连续空格，允许匹配变体
        flexible_pattern = re.sub(r"\\ ", r"\\s+", flexible_pattern)

        # 完整模式：必须完整匹配整个chunk（避免部分匹配）
        return re.compile(rf"\b{flexible_pattern}\b", flags=re.DOTALL)

    def highlight_markdown(text, chunk):
        """在文本中找到最相似的部分并高亮"""
        logger.info(f"chunk length: {len(chunk)}")
        # 把换行符变成可选匹配（\s* 匹配任意空白）
        # pattern = re.compile(re.escape(clean_chunk))
        pattern = make_flexible_pattern(chunk)
        match = pattern.search(text)

        if match:
            matched_text = match.group()
            return text.replace(matched_text, f'<span style="background-color: #fff3cd;">{matched_text}</span>')
        else:
            logger.info("无法匹配，尝试调整换行符...")
            chunk_0 = chunk.split("；")[0]
            chunk_1 = chunk.split("\n")[0]
            if chunk_0 in text:
                start = text.find(chunk_0)
                return text[
                       :start] + f'<span style="background-color: #fff3cd;">{text[start:start + len(chunk)]}</span>' + text[
                                                                                                                       start + len(
                                                                                                                           chunk):]
            elif chunk_1 in text:
                start = text.find(chunk_1)
                return text[
                       :start] + f'<span style="background-color: #fff3cd;">{text[start:start + len(chunk)]}</span>' + text[
                                                                                                                       start + len(
                                                                                                                           chunk):]

            first_line = chunk.split('\n')[0].split("；")[0]  # 获取首行内容
            if first_line and first_line in text:
                # 找到首行位置
                start = text.find(first_line)
                # 找到chunk在原文中的结束位置（可能跨越多行）
                end = start
                for line in chunk.split('\n'):
                    line_end = text.find(line, end)
                    if line_end == -1:
                        break
                    end = line_end + len(line)

                # 确保至少匹配了首行
                if end > start:
                    return text[:start] + f'<span style="background-color: #fff3cd;">{text[start:end]}</span>' + text[
                                                                                                                 end:]
            else:
                logger.warning("仍然无法匹配！,采用模糊匹配")
                result = find_best_match(text, chunk)
                if not result:
                    return text

                start, end, score = result
                logger.info(f"最佳匹配得分: {score}")

                # 确保不重复高亮（处理嵌套匹配）
                if '<span style=' not in text[start:end]:
                    return text[:start] + \
                        f'<span style="background-color: #fff3cd;">{text[start:end]}</span>' + \
                        text[end:]
                return text

    highlighted = highlight_markdown(
        # markdown_content,
        request.markdown_content,
        # all_content,
        request.chunk
    )
    # logger.info(f"hightlighted text :{hightlighted_text}")
    return {"code": 200, "message": "success", "data": [{"highlighted_markdown": highlighted}]}


import shutil
from fastapi.staticfiles import StaticFiles
app.mount("/converted_pdfs", StaticFiles(directory="converted_pdfs"), name="converted_pdfs")
@app.post("/word_to_pdf")
async def get_markdown(file: UploadFile = File(...)):
    filename = file.filename
    file_ext = os.path.splitext(filename)[1].lower()
    file_name = os.path.splitext(filename)[0].lower()
    pdf_url = None

    if file_ext == ".docx" or file_ext == ".doc":
        converted_path = await convert_docx_to_pdf(
            os.getenv("CONVERT_DOC_URL", "http://192.168.0.46:59070"), file
        )

        try:
            output_dir = "converted_pdfs"
            os.makedirs(output_dir, exist_ok=True)

            pdf_filename = f"{file_name}_{uuid.uuid4().hex[:8]}.pdf"
            persistent_pdf_path = os.path.join(output_dir, pdf_filename)
            shutil.copy2(converted_path, persistent_pdf_path)

            # 构造可访问的 HTTP URL
            base_url = os.getenv("BASE_URL", "http://192.168.0.46:59079")
            pdf_url = f"{base_url}/converted_pdfs/{pdf_filename}"
            logger.info(f"DOCX converted to PDF, accessible at: {pdf_url}")

        except Exception as e:
            logger.error(f"Failed to save PDF file: {e}")
            raise HTTPException(status_code=500, detail="Failed to process and save PDF.")

    else:
        raise HTTPException(
            status_code=400,
            detail="Unsupported file type. Only DOC and DOCX are supported."
        )

    return {
        "message": "success",
        "code": 200,
        "data": {
            "pdf_url": pdf_url,  # 返回完整可下载链接
        }
    }


class hightlightboxRequest(BaseModel):
    split_chunk: List[dict] = []
    content_middle: List[dict] = []


@app.post("/get_highlight_bbox")
async def get_hightlight(request: hightlightboxRequest):
    split_chunk = request.split_chunk
    content_middle = request.content_middle


    def sort_by_page_idx_frequency(data):

        page_idx_counts = Counter(item['page_idx'] for item in data)

        # 获取出现频率最高的前三个 page_idx（如果有并列，取前三个）
        top_three_page_idxs = {page_idx for page_idx, count in page_idx_counts.most_common(3)}

        # 筛选出所有属于前三高频 page_idx 的原始记录
        filtered_data = [item for item in data if item['page_idx'] in top_three_page_idxs]

        # 在这些高频记录中找出最大的 page_idx
        max_page_idx = max(top_three_page_idxs)  # 因为 top_three_page_idxs 是前三高频的键

        return filtered_data, max_page_idx

    def find_and_filter_records(record_list, text):
        # 尝试找到第一个“content”作为text前缀的记录的索引
        start_index = next((i for i, record in enumerate(record_list) if
                            record.get('content', '').strip() and text.startswith(record['content'])), None)
        # 如果找到了这样的记录，返回从这条记录开始到最后的所有记录
        if start_index is not None:
            record_list = record_list[start_index:]

        # 如果没有找到以“content”开头的记录，则尝试找到第一个“content”作为text后缀的记录的索引
        end_index = next((i for i, record in enumerate(record_list) if
                          record.get('content', '').strip() and text.endswith(record['content'])), None)
        #
        # # 如果找到了这样的记录，返回从这条记录开始到最后的所有记录
        if end_index is not None:
            record_list = record_list[:end_index+1]

        # 如果既没有找到匹配的前缀也没有找到匹配的后缀，则返回原list
        return record_list



    def getbbox(result_devide, content_middle):
        for idx, item in enumerate(content_middle):
            item['id'] = idx + 1  # 从1开始编号

        # for item in content_middle:
        #     print(item)
        resource = result_devide[0]['resource']
        for re in result_devide:
            if 'chunklist' not in re:
                re['chunklist'] = []
            temp_chunk_list = []

            for record in content_middle:
                if record['type'] == 'text':
                    if record['content'] in re['text']:
                        temp_chunk_list.append(record)
                elif record['type'] == 'image':
                    if record['image_path'] in re['text']:
                        temp_chunk_list.append(record)
                elif record['type'] == 'table':
                    if record['html'] in re['text']:
                        temp_chunk_list.append(record)
            re['chunklist'] = temp_chunk_list
        result_devide = [item for item in result_devide if item['chunklist']]

        length_result_devide = len(result_devide)
        result = {}

        for i in range(length_result_devide):
            temp_uuid = str(uuid.uuid4())
            result[temp_uuid] = []
            chunk = result_devide[i]['text']
            result_devide[i]["id"] = temp_uuid

            # 假设输入数据存储在名为data的列表中
            data = result_devide[i]['chunklist']  # 您的数据列表

            sorted_data = sorted(data, key=lambda x: x['id'])

            # 步骤2：使用 groupby 找出连续 id 的组
            groups = []
            current_group = []

            for item in sorted_data:
                if not current_group:
                    current_group.append(item)
                else:
                    # 检查当前 id 是否与前一个连续
                    if item['id'] == current_group[-1]['id'] + 1:
                        current_group.append(item)
                    else:
                        # 不连续，保存当前组并开始新组
                        if len(current_group) >= 1:  # 可设最小长度过滤
                            groups.append(current_group)
                        current_group = [item]

            # 别忘了最后一个组
            if current_group:
                groups.append(current_group)

            # 按组的长度从大到小排序
            sorted_groups = sorted(groups, key=len, reverse=True)

            sorted_data, most_common_page_idx1 = sort_by_page_idx_frequency(sorted_groups[0])
            sorted_filtered_items = sorted(
                (item for sublist in sorted_groups for item in sublist if
                 most_common_page_idx1 - 1 <= item['page_idx'] <= most_common_page_idx1 + 1),
                key=lambda x: x['id']
            )
            sorted_filtered_items = find_and_filter_records(sorted_filtered_items, chunk)
            if not sorted_filtered_items :
                most_common_page_idx = 0
            else:
                most_common_page_idx = sorted_filtered_items[0]['page_idx']

            for chunk in sorted_filtered_items:
                chunk['resource'] = resource
                result[temp_uuid].append(chunk)
            result_devide[i]['page_idx'] = most_common_page_idx
            # result[temp_uuid] =sorted_data
        #
        data = []
        for ins in result_devide:
            data.append(
                {
                    'id': ins['id'],
                    'Header_1': ins['Header_1'],
                    'Header_2': ins['Header_2'],
                    'Header_3': ins['Header_3'],
                    'text': ins['text'],
                    'resource': resource,
                    'img_path': ins['img_path'],
                    'origin_text': ins['origin_text'],
                    'page': None,
                    'page_idx': ins['page_idx']
                }

            )
        return data, result

    data, result = getbbox(split_chunk, content_middle)
    if not data or not result:
        return {"code": 200, "message": "success", "data": []}
    else:
        return {"code": 200, "message": "success", "data": [{"data": data, "result": result}]}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=9079)