# -*- encoding:utf-8 -*-
import json
import re
import fitz
from PyPDF2 import PdfReader
import textwrap
from flask_app.general.llm.大模型通用函数 import read_txt_to_string
from flask_app.general.json_utils import clean_json_string
from flask_app.general.llm.model_continue_query import process_continue_answers
from flask_app.general.截取pdf通用函数 import create_get_text_function
from flask_app.general.llm.通义千问long import upload_file, qianwen_long_stream
from flask_app.general.llm.qianwen_plus import qianwen_plus
from flask_app.general.读取文件.clean_pdf import extract_common_header, clean_page_content
from flask_app.general.format_change import docx2pdf, pdf2docx
import concurrent.futures


# 正则表达式判断原文中是否有商务、服务、其他要求
def find_exists(truncate_file, required_keys):
    # if not truncate_file:
    #     return ["技术要求", "商务要求", "服务要求", "其他要求"]
    common_header = extract_common_header(truncate_file)  # 假设该函数已定义
    try:
        pdf_document = PdfReader(truncate_file)
        pdf_lib = 'pypdf2'
    except Exception as e:
        print(f"使用 PyPDF2 读取失败，切换到 fitz。错误信息: {e}")
        pdf_document = fitz.open(truncate_file)
        pdf_lib = 'fitz'
    get_text = create_get_text_function(pdf_lib, pdf_document)
    # 获取总页数
    if pdf_lib == 'pypdf2':
        total_pages = len(pdf_document.pages)
    else:  # fitz
        total_pages = pdf_document.page_count
    # 定义正则模式
    begin_pattern = re.compile(
        r'(?:^第[一二三四五六七八九十百千]+(?:章|部分)\s*'  # 匹配“第X章”或“第X部分”
        r'[\u4e00-\u9fff、()（）]*?'  # 匹配允许的字符
        r'(?:(?:服务|项目|商务|技术)[\u4e00-\u9fff、()（）]*?要求|'  # 匹配“服务”、“项目”、“商务”或“技术”后跟“要求”
        r'(?:采购|需求)[\u4e00-\u9fff、()（）]*?)'  # 匹配“采购”或“需求”
        r'\s*$|'  # 匹配行尾
        r'^第[一二三四五六七八九十百千]+(?:章|部分)(?!.*说明).*?'  # 匹配“第X章”后带“采购内容”等，排除“说明”
        r'(?:采购内容|采购要求|需求).*|'  # 匹配“采购内容”或“采购要求”关键词
        r'^[一二三四五六七八九十百千]+、\s*采购清单)'  # 匹配“一、采购清单”
        r'\s*$',  # 匹配行尾
        re.MULTILINE
    )
    end_pattern = re.compile(
        r'第[一二三四五六七八九十1-9]+(?:章|部分)\s*[\u4e00-\u9fff、()（）]+\s*$', re.MULTILINE)

    # 处理第一页和最后一页
    first_page_text = get_text(0)
    last_page_text = get_text(total_pages - 1)

    # 清理页面内容
    first_page_clean = clean_page_content(first_page_text, common_header)
    last_page_clean = clean_page_content(last_page_text, common_header)

    # 在第一页寻找起始位置
    start_match = re.search(begin_pattern, first_page_clean)
    if not start_match:
        print("未找到开始模式，返回完整第一页")
        first_content = first_page_clean
    else:
        start_index = start_match.end()
        first_content = first_page_clean[start_index:]

    # 在最后一页寻找结束位置
    end_match = re.search(end_pattern, last_page_clean)
    if not end_match:
        print("未找到结束模式，返回完整最后一页")
        last_content = last_page_clean
    else:
        last_content = last_page_clean[:end_match.start()]

    # 获取中间页面的内容
    middle_content = ""
    if total_pages > 2:
        for page_num in range(1, total_pages - 1):
            page_text = get_text(page_num)
            cleaned_text = clean_page_content(page_text, common_header)
            middle_content += cleaned_text + "\n"

    # 组合所有内容
    relevant_text = first_content + "\n" + middle_content + "\n" + last_content
    relevant_text = re.sub(r'\s+', ' ', relevant_text)

    # print(f"提取的内容范围:\n{relevant_text}")

    # 匹配所需的要求
    matched_requirements = []
    punctuation = r"[，。？！、；：,.?!]*"
    for req in required_keys:
        if re.search(req, relevant_text):
            # 替换特定的关键词
            if req in [r"总\s*体\s*要\s*求", r"建\s*设\s*要\s*求",r"整\s*体\s*要\s*求"]:
                # 匹配到“总体要求”或“建设要求”时，添加“技术要求”而非原关键词
                tech_req = r"技\s*术\s*要\s*求"
                if tech_req not in matched_requirements:
                    matched_requirements.append(tech_req)
            elif req in [r"培\s*训\s*要\s*求", r"质\s*保\s*要\s*求", r"售\s*后\s*要\s*求"]:
                # 匹配到“培训要求”、“质量保证要求”或“售后要求”时，添加“服务要求”而非原关键词
                service_req = r"服\s*务\s*要\s*求"
                if service_req not in matched_requirements:
                    matched_requirements.append(service_req)
            elif req in [r"进\s*度\s*要\s*求", r"工\s*期\s*要\s*求"]:
                # 匹配到“进度要求”或“工期要求”时，添加“商务要求”而非原关键词
                busi_req = r"商\s*务\s*要\s*求"
                if busi_req not in matched_requirements:
                    matched_requirements.append(busi_req)
            elif req == r"服\s*务\s*要\s*求":
                # 处理“服务要求”时的特殊逻辑
                lines = [line for line in relevant_text.split('\n') if re.search(req, line)]
                pattern = r"技\s*术" + punctuation + req
                if any(re.search(pattern, line) for line in lines):
                    combined_req = r"技\s*术\s*、\s*服\s*务\s*要\s*求"
                    if combined_req not in matched_requirements:
                        matched_requirements.append(combined_req)
                else:
                    if req not in matched_requirements:
                        matched_requirements.append(req)
            else:
                # 对于其他匹配的关键词，直接添加
                if req not in matched_requirements:
                    matched_requirements.append(req)

    # 以下部分为原有逻辑，用于替换和去除 \s*，保留注释以便以后切换
    # # 去除 \s*，仅返回原始关键词
    # clean_requirements = [re.sub(r'\\s\*', '', req) for req in matched_requirements]

    # # 判断互斥关系：如果有"技术、服务要求"，删除"技术要求"和"服务要求"
    # if "技术、服务要求" in clean_requirements:
    #     clean_requirements = [req for req in clean_requirements if req not in ["技术要求", "服务要求"]]
    # # 确保最终返回的列表仅包含指定的五项
    # # allowed_requirements = {"技术要求", "服务要求", "商务要求", "其他要求", "技术、服务要求"}
    # # final_requirements = [req for req in clean_requirements if req in allowed_requirements]

        # 简化后的逻辑：默认添加 '商务要求' 和 '服务要求'
        # 1. 如果存在“技术、服务要求”，则不返回“服务要求”
        # 2. 否则，确保“商务要求”和“服务要求”在列表中
    clean_requirements = [re.sub(r'\\s\*', '', req) for req in matched_requirements]
    if "技术、服务要求" in clean_requirements:
        # 如果存在“技术、服务要求”，则移除“服务要求”并确保添加“技术、服务要求”
        clean_requirements = [req for req in clean_requirements if req != "服务要求" and req != "技术、服务要求"]
        if "技术、服务要求" not in clean_requirements:
            clean_requirements.append("技术、服务要求")
    else:
        # 默认添加“商务要求”和“服务要求”如果它们不在列表中
        if "商务要求" not in clean_requirements:
            clean_requirements.append("商务要求")
        if "服务要求" not in clean_requirements:
            clean_requirements.append("服务要求")
    # 去除重复项
    clean_requirements = list(set(clean_requirements))

    # 最终返回清理后的要求列表
    return clean_requirements

def generate_template(required_keys,full_text, type=1):
    # 定义每个键对应的示例内容
    example_content1 = {
        "技术要求": ["相关技术要求1", "相关技术要求2"],
        "服务要求": ["服务要求1", "服务要求2", "服务要求3"],
        "商务要求": ["商务要求1", "商务要求2"],
        "其他要求": ["关于项目采购的其他要求1...", "关于项目采购的其他要求2...", "关于项目采购的其他要求3..."],
        "技术、服务要求": ["相关技术、服务要求内容1", "相关技术、服务要求内容2", "相关技术、服务要求内容3"]
    }
    example_content2 = {
        "技术要求": {
            "子因素名1": ["相关技术要求1", "相关技术要求2"]
        },
        "服务要求": {
            "子因素名1": ["服务要求1"],
            "子因素名2": ["服务要求2", "服务要求3"]
        },
        "商务要求": {
            "子因素名1": ["商务要求1"],
            "子因素名2": ["商务要求2"]
        },
        "其他要求": {
            "子因素名1": ["关于项目采购的其他要求1...", "关于项目采购的其他要求2..."],
            "子因素名2": ["关于项目采购的其他要求3..."]
        },
        "技术、服务要求": {
            "子因素名1": ["相关技术、服务要求内容1"],
            "子因素名2": ["相关技术、服务要求内容2", "相关技术、服务要求内容3"]
        }
    }
    # 将 required_keys 转换为集合以便于操作
    keys = set(required_keys)
    type_to_keys_map = {
        1: ["服务要求", "商务要求", "其他要求"],
        2: ["技术要求", "技术、服务要求"]
    }
    # 根据 type 获取对应的 all_possible_keys
    chosen_keys = type_to_keys_map.get(type, [])
    another_keys_list = type_to_keys_map.get(3 - type, [])  # 3 - type 将 type 1 映射到 2，反之亦然
    another_keys_str = ', '.join([f"'{key}'" for key in another_keys_list])
    # 处理互斥关系：如果 "技术要求" 和 "服务要求" 同时存在，则移除 "技术、服务要求"
    if "技术要求" in keys and "服务要求" in keys:
        keys.discard("技术、服务要求")
    # 如果 "技术、服务要求" 存在，则移除 "技术要求" 和 "服务要求"
    elif "技术、服务要求" in keys:
        keys.discard("技术要求")
        keys.discard("服务要求")

    # 确保 keys 中只包含允许的键
    keys = keys.intersection(chosen_keys)
    # 按照预定义的顺序排序键，以保持一致性
    sorted_keys = [key for key in chosen_keys if key in keys]

    # 如果没有任何键被选中，返回""
    if not sorted_keys:
        return ""

    # 生成模板的通用部分
    def generate_prompt_instruction(keys_str, outer_keys_str, another_keys_str, type):
        if type == 1:
            specific_instructions = textwrap.dedent(
                """6. 补充要求（商务要求提取）:
        -商务要求提取范围：
            如果章节开头位置或采购清单中，除了列出货物名称，还描述了如工期要求、进度要求、品牌要求等商务要求，需提取这些内容
            若文档标题包含“工期要求”、“进度要求”等商务要求相关的关键字，应提取对应内容。
        -商务要求的组织形式：
            嵌套键值对形式：添加至 '商务要求' 的键值部分，嵌套键名为对应的子标题，保留 三角▲、五角星★ 等特殊符号（若有）。若不存在这些内容，无需额外添加，避免返回空的嵌套键值对'工期要求':[]
            直接添加具体内容：不采用嵌套键值对形式，将具体内容直接作为字符串列表的一部分添加到 '商务要求' 的键值部分。
    7. 补充要求（服务要求提取）：
        -在提取'服务要求'的时候，若原文（包含正文和表格）中存在'安装要求'、'售后要求'、'维护要求'、'培训要求、质量要求'等服务相关的标题及内容，不要遗漏这部分的'服务要求'：
            嵌套键值对形式：添加至 '服务要求' 的键值部分，嵌套键名为对应的子标题，保留 三角▲、五角星★ 等特殊符号（若有）。若不存在这些内容，无需额外添加，避免返回空的嵌套键值对'安装要求':[]
            直接添加具体内容：不采用嵌套键值对形式，将具体内容直接作为字符串列表的一部分添加到 '服务要求' 的键值部分。
    8. 避免重复提取：
        若正文某部分已被提取，则无需再次重复提取。例如：提取'服务要求'时得到了'售后要求'相关内容，那么在提取'商务要求'时无需再提取该内容。如果文档中未明确列出某类要求，直接返回空列表 []。
    **限制内容**：
        - **避免提取技术要求**：在提取这些要求时，确保不包含任何与技术规格、功能或性能相关的内容。
        - **避免提取资格审查相关内容**：在提取这些要求时，确保不包含任何与资格审查、符合性审查、形式评审有关的内容。
                """
            )
        else:
            specific_instructions = textwrap.dedent(
                f"""6. 技术要求提取规则：
        -在提取技术要求或技术、服务要求时，你无需从采购清单或表格中提取具体设备、采购标的的技术要求以及参数要求，你仅需定位到原文中包含'技术要求'或'技术、服务要求'关键字的标题，并提取该标题下的整体技术要求内容；
    7. 补充要求（技术要求提取）：
        -在提取{outer_keys_str}的时候，若原文中存在如“总体要求”、“建设要求”等子标题，不要遗漏这部分的'技术要求':
            嵌套键值对形式：添加至 {outer_keys_str}的键值部分，嵌套键名为对应的子标题，保留 三角▲、五角星★ 等特殊符号（若有）。若不存在这些内容，无需额外添加，避免返回空的嵌套键值对'安装要求':[]
            直接添加具体内容：不采用嵌套键值对形式，将具体内容直接作为字符串列表的一部分添加到 {outer_keys_str}的键值部分。
    8. 在提取'技术要求'时，注意不要提取有关'安装、售后、维护、运维、培训、质保、工期、进度'等要求，它们不属于'技术要求'。
    **限制内容**：
        - **避免提取商务要求**：在提取技术要求时，确保不包含任何与商务要求相关的内容，例如'安装、售后、维护、运维、培训、质保、工期、进度'等要求，它们不属于'技术要求'。
        - **避免提取资格审查相关内容**：在提取技术要求时，确保不包含任何与资格审查、符合性审查、形式评审有关的内容。
                """
            )
        return textwrap.dedent(
            f"""请你根据该货物类招标文件中的采购要求部分内容，请告诉我该项目采购的{keys_str}分别是什么，请以json格式返回结果，外层键名是{outer_keys_str}，默认情况下键值为字符串列表，每个字符串表示具体的一条要求，可以按原文中的序号作划分（若有序号的话），请按原文内容回答，不要擅自增删内容。
    
    要求与指南：
    1. **提取范围**：仅提取文件中与采购要求直接相关的内容，重点关注针对投标人、中标人、供应商等投标相关主体的具体要求，是整体要求，而非针对具体采购物品的技术参数或功能要求。
        -避免提取如行政性要求（投标文件的提交方式、截止时间、地点等）、招标活动流程、答疑相关内容等。
        -避免提取 {another_keys_str} 中的内容。
    2. **如果在相应要求下发现只有标题性质的子标题,却没有实际的具体要求，那么可以**：
        -忽略掉这些标题,不要将它们与下面的具体要求合并在一起。
        -也可以将它们作为该要求下的嵌套键名,但字符串列表中只提取实际的具体要求。
    3. **内容提取规则**：
        -**保留原始格式和符号**：字符串列表中的每个字符串内容需与原文内容保持一致，保留前面的三角▲、五角星★或其他特殊符号和序号（如果有）。不得擅自添加、删减这些符号。
        -如果文档中有明确的标题或子标题，其前面带有三角▲、五角星★等特殊符号，则键名中应完整保留这些符号，与原文保持一致。
        -表格形式处理：
            -注意请不要返回Markdown表格语法，可以使用冒号':'将相关信息拼接在一起，如"交付期:合同签订之日起30天内。"；或将其组织为嵌套键值对形式，最多允许一层嵌套。
            -表格中出现的特殊符号如▲★需要添加至相应键值中。
    4. **JSON 的结构要求**：
        - 默认情况无需嵌套键值对，键值为字符串列表；
        - 如果文档中有明确的子标题或者表格形式的各行表示子要求，则采用嵌套键值对形式，嵌套键名是各子要求，内层键值为字符串列表，表示该子要求下的具体要求。最多允许一层嵌套。
            - 每个外层键（大要求）对应的值可以是：
                a. 一个字符串列表，表示具体的一条条要求。若只有一条要求，也用字符串列表表示。
                b. 一个对象（字典），其键为子要求，值为字符串列表。
                c. 如果文档中没有找到相关的要求，键值为空列表[]，无需进行推测或强行生成答案。
            - 最多只允许一层嵌套。
    5. **优先定位规则**：
        -请优先且准确定位正文部分包含以下关键字的标题或表格：{outer_keys_str}，在其之后提取'XX要求'相关内容，尽量避免在无关地方提取内容。
    {specific_instructions}""")

    # 过滤示例内容
    def filter_example_content(example_content, keys):
        return {k: v for k, v in example_content.items() if k in keys}

    def format_example(example_content):
        return json.dumps(example_content, indent=4, ensure_ascii=False)

    filtered_example_content1 = filter_example_content(example_content1, sorted_keys)
    filtered_example_content2 = filter_example_content(example_content2, sorted_keys)
    tech_json_example1_str = format_example(filtered_example_content1)
    tech_json_example2_str = format_example(filtered_example_content2)
    keys_str = '、'.join(sorted_keys)
    outer_keys_str = ', '.join([f"'{key}'" for key in sorted_keys])
    prompt_instruction = generate_prompt_instruction(keys_str, outer_keys_str, another_keys_str, type)
    # 完整的用户查询模板，包含两份示例输出
    user_query_template = f"""
    {prompt_instruction}
    以下为示例输出，仅供格式参考：
    示例 1,无嵌套键值对：
    {tech_json_example1_str}
    示例 2，嵌套键值对形式：
    {tech_json_example2_str}
    """
    if full_text:
        user_query_template = f"文件内容：{full_text}\n" + user_query_template
    return user_query_template

def get_business_requirements(procurement_path, processed_filepath, model_type):
    required_keys = ["技\s*术\s*要\s*求", "商\s*务\s*要\s*求", "服\s*务\s*要\s*求", "其\s*他\s*要\s*求",
                     "总\s*体\s*要\s*求", "整\s*体\s*要\s*求","建\s*设\s*要\s*求", "进\s*度\s*要\s*求", "工\s*期\s*要\s*求",
                     "质\s*保\s*要\s*求", "培\s*训\s*要\s*求", "售\s*后\s*要\s*求"]

    # 将 doc/docx 转换为 pdf
    procurement_pdf_path = procurement_path
    if procurement_path.lower().endswith(('.doc', '.docx')):
        procurement_pdf_path = docx2pdf(procurement_path)

    # 查找包含的关键词
    contained_keys = find_exists(procurement_pdf_path, required_keys)
    print(contained_keys)
    if not contained_keys:
        return {}

    # 读取文件全文
    full_text = read_txt_to_string(processed_filepath)

    # 生成业务查询和技术查询
    busi_user_query = generate_template(contained_keys, full_text, 1)
    tech_user_query = generate_template(contained_keys, full_text, 2)
    # 初始化结果存储
    final_res = {}

    # 如果是非模型调用，需要提前上传文件并获取 file_id
    file_id = None
    if model_type!=4:
        procurement_docx_path=procurement_path
        if procurement_path.lower().endswith('.pdf'):
            procurement_docx_path = pdf2docx(procurement_path)
        file_id = upload_file(procurement_docx_path)  # 只上传一次文件，避免冗余调用

    # 并行处理业务和技术查询
    questions_to_continue = []  # 存储需要调用 continue_answer 的 (original_query, parsed)
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
        futures = []
        future_to_query = {}  # 创建一个字典来映射 future 到 original_query

        # 将所有需要处理的查询放入一个列表中
        queries = [
            ('busi_user_query', busi_user_query),
            ('tech_user_query', tech_user_query)
        ]
        for query_name, user_query in queries:
            if user_query:
                if model_type == 4:
                    # 如果是模型调用，直接使用 qianwen_plus
                    future = executor.submit(qianwen_plus, user_query, True)
                else:
                    # 使用 qianwen_long_stream 并传入 file_id
                    future = executor.submit(qianwen_long_stream, file_id, user_query, 2, 1, True)
                futures.append(future)
                future_to_query[future] = user_query  # 映射 future 到 user_query
            # 收集需要继续回答的问题
        initial_results = {}
        max_tokens = 7900 if model_type == 4 else 5900
        # 获取结果
        for future in concurrent.futures.as_completed(futures):
            original_query = future_to_query[future]  # 获取对应的 original_query
            try:
                result = future.result()
                if result:  # 确保结果不为空
                    message = result[0]
                    parsed = clean_json_string(message)
                    total_tokens = result[1]
                    if not parsed and total_tokens > max_tokens:
                        questions_to_continue.append((original_query, message))
                    else:
                        initial_results.update(parsed)
            except Exception as e:
                print(f"An error occurred: {e}")
    # 处理需要继续回答的问题
    if questions_to_continue:
        continued_results = process_continue_answers(questions_to_continue, model_type, file_id)
        final_res.update(continued_results)

    # 合并初步结果
    final_res.update(initial_results)

    return final_res

if __name__ == "__main__":
    # truncate_file = "C:\\Users\\Administrator\\Desktop\\fsdownload\\e4be098d-b378-4126-9c32-a742b237b3b1\\ztbfile_procurement.docx"
    # truncate_file = r"C:\Users\Administrator\Desktop\货物标\output1\2-招标文件（广水市教育局封闭管理）_procurement.pdf"
    procurement_path=r'C:\Users\Administrator\Desktop\fsdownload\bbf7504f-3c75-45e5-b3e2-ab0a15ec9c14\tmp\ztbfile_procurement.pdf'
    docx_path=r'D:\flask_project\flask_app\static\output\output1\83ae3e35-9136-4402-a74f-01d7adfcbb73\invalid_added.docx'
    # truncate_file=r"C:\Users\Administrator\Desktop\new招标文件\output5\HBDL-2024-0519-001-招标文件_procurement.pdf"
    # file_id = upload_file(truncate_file)
    # processed_filepath = pdf2txt(procurement_path)
    processed_filepath=r'C:\Users\Administrator\Desktop\fsdownload\bbf7504f-3c75-45e5-b3e2-ab0a15ec9c14\tmp\extract1.txt'
    final_res= get_business_requirements(procurement_path,processed_filepath,4)
    print(json.dumps(final_res, ensure_ascii=False, indent=4))
