#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
集成测试QueryRewriteManager.query_extract方法
"""

import sys
import os
import time
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from foundation.ai.rag.retrieval.query_rewrite import QueryRewriteManager
from foundation.observability.logger.loggering import server_logger as logger


def test_query_extract_with_construction_content():
    """
    测试使用施工方案内容的query_extract
    """
    print("="*60)
    print("测试QueryRewriteManager.query_extract方法")
    print("="*60)

    query_rewrite_manager = QueryRewriteManager()

    # 测试数据: 基坑工程
    review_content = """（1）应将架桥机随机文件和有关技术资料准备齐全，认真组织安装人员学习阅
读本方案，并以此为依据结合实际工况及地貌拟定有关安装施工方案。
（2）安装前，应对设备散件进行全面检查、清理，如发现有损伤、腐蚀或其它
缺陷，应在安装前予以处理，合格后方可安装。
（3）应对架桥机运梁轨道进行如下检查：
1）架桥机运梁轨道基础应有足够的承压能力，应能满足架桥机运梁平车载重安
全运行。轨道下部如有支垫是否具有强大的刚性和一定的密度，特殊端部的支垫必
须采用刚性支垫，严禁使用腐蚀的枕木。
2）轨道钢轨正面、侧面的不平度不应大于1/1500，全长范围内不应大于10毫
米。
3）轨道安装的允许偏差：
①轨道实际中心线对轨道设计中心线的位置偏移允许偏差为3毫米。
②轨距允许偏差为±5毫米。
③轨道纵向不平度应小于 1/1500，且全行程不超过10毫米。
④同一断面上两轨道的标高相对偏差不超过5毫米。
4）轨道接头应符合下列要求：
①接头左、右、上三面错位不应大于2毫米。
②两平行轨道接口的位置应错开，其错开距离不应等于架桥机运梁平车前后车
轮的轮距。
第19页
四川公路桥梁建设集团有限公司 镇广C4项目经理部JQ220t-40m架桥机安装及拆除专项施工方案
③接头间隙应为1~2毫米，伸缩缝接头间隙应符合设计要求，其偏差不应大于
±1毫米。
④轨道的悬空部位是否得到了加强。
5）施工期间的基本要求及安全规定
①架桥机安装属于高空作业，施工前技术员及现场指挥员向参加施工的所有人
员详细介绍安装工序、技术要求和指挥信号。
②严禁在施工的架桥机下面逗留通过，与安装施工无关人员不准擅自进入施工
现场。
③现场施工人员必须按有关安全规定佩戴好安全用品（安全带、安全帽、绝缘
鞋等）。
④施工现场使用的氧气瓶、乙炔瓶必须保证立放并固定好，两瓶之间距离不得
小于5m。
6）架桥机组装
①安装原则：按从下至上的顺序进行安装，且必须在下部结构安装固定牢固以
后，才能进行上部结构的安装。同时，在安装过程中，在运输及吊装能力允许的情
况下，尽量保持设备各部件的整体性，一方面能有效减少危险作业点，另一方面有
利于提高设备现场安装工效。
②安装场地：架桥拼装在桥台台背后路基上，拼装场地选址位于路基上，拼装
场地面积宽15m，长 100m，架桥机主体最大宽度为9m（移动轨道 18m），路基宽度
满足安装场地所需，路基及便道顶面承载力满足接卸拼装所需的承载要求。
③辅助机械设备：根据施工现场实际情况，选择1台25T汽车吊作为主要安装
设备，架桥机主梁拼装采用分节进行吊装，整机重量最大的组件为三节主梁同时起
吊，总重18t，采用2 台50t汽车吊满足吊装要求。"""

    print(f"输入内容: {review_content}")
    print(f"内容长度: {len(review_content)}")

    try:
        start_time = time.time()
        result = query_rewrite_manager.query_extract(review_content)
        end_time = time.time()
        elapsed_time = end_time - start_time

        print(f"\n[OK] 提取完成，耗时: {elapsed_time:.2f}秒")
        print(f"返回结果类型: {type(result)}")
        print(f"返回结果长度: {len(str(result))}")
        print(f"返回结果: {result}")

        if result:
            print("[OK] 成功提取到查询信息")

            # 尝试解析JSON响应
            try:
                import json
                if isinstance(result, str):
                    # 提取JSON部分
                    if '```json' in result:
                        start = result.find('```json') + 7
                        end = result.find('```', start)
                        if end != -1:
                            json_str = result[start:end].strip()
                        else:
                            json_str = result[start:].strip()
                    else:
                        json_str = result.strip()

                    parsed_result = json.loads(json_str)
                    print(f"[OK] JSON解析成功，提取到 {len(parsed_result)} 个实体")

                    for i, entity in enumerate(parsed_result):
                        entity_name = entity.get('entity', 'N/A')
                        keywords = entity.get('search_keywords', [])
                        background = entity.get('background', 'N/A')
                        parameter = entity.get('parameter', 'N/A')

                        print(f"  实体{i+1}: {entity_name}")
                        print(f"    关键词: {keywords}")
                        print(f"    背景: {background}")
                        print(f"    参数: {parameter}")
                        print()
            except Exception as e:
                print(f"[WARN] JSON解析失败: {e}")
        else:
            print("[ERROR] 未返回有效结果")

    except Exception as e:
        print(f"[ERROR] 测试失败: {str(e)}")
        logger.error(f"测试失败: {str(e)}")
        import traceback
        traceback.print_exc()


def main():
    """
    主测试函数
    """
    print("开始集成测试 QueryRewriteManager.query_extract 方法")

    # 执行核心测试
    test_query_extract_with_construction_content()

    print("\n集成测试完成")


if __name__ == "__main__":
    main()