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

"""
RAG管道测试服务器
提供HTTP API供前端调用RAG链路
"""

import sys
import os
import json
import time
import asyncio
from http.server import HTTPServer, SimpleHTTPRequestHandler
from urllib.parse import parse_qs, urlparse
import threading

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
sys.path.insert(0, project_root)

from core.construction_review.component.infrastructure.milvus import MilvusConfig, MilvusManager
from core.construction_review.component.infrastructure.parent_tool import (
    enhance_with_parent_docs_grouped,
    extract_query_pairs_results
)
from foundation.ai.rag.retrieval.entities_enhance import entity_enhance
from foundation.ai.rag.retrieval.query_rewrite import query_rewrite_manager
from foundation.observability.logger.loggering import review_logger as logger
from foundation.observability.monitoring.rag import rag_monitor
from core.construction_review.component.ai_review_engine import AIReviewEngine
from core.base.task_models import TaskFileInfo

# 全局Milvus Manager
milvus_manager = None
# 全局AI审查引擎
ai_review_engine = None


def init_milvus():
    """初始化Milvus Manager"""
    global milvus_manager
    if milvus_manager is None:
        print("📌 初始化Milvus Manager...")
        milvus_manager = MilvusManager(MilvusConfig())
        print("✅ Milvus Manager 初始化成功")
    return milvus_manager


def init_ai_review_engine():
    """初始化AI审查引擎"""
    global ai_review_engine
    if ai_review_engine is None:
        print("📌 初始化AI审查引擎...")
        # 创建测试用的TaskFileInfo
        file_info_dict = {
            'file_id': "test_file_id",
            'callback_task_id': "test_task_id",
            'user_id': "test_user",
            'file_name': "测试文件.docx",
            'file_type': 'docx',
            'file_content': b'',
            'review_config': [],
            'review_item_config': {}
        }
        task_info = TaskFileInfo(file_info_dict)
        ai_review_engine = AIReviewEngine(task_info)
        print("✅ AI审查引擎初始化成功")
    return ai_review_engine


def run_async(coro):
    """在合适的环境中运行异步函数"""
    try:
        loop = asyncio.get_running_loop()
        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor() as executor:
            future = executor.submit(asyncio.run, coro)
            return future.result()
    except RuntimeError:
        return asyncio.run(coro)


async def professional_review_test(review_content: str, check_type: str = "both") -> dict:
    """
    专业性审查完整测试 - RAG召回 + AI审查

    Args:
        review_content: 待审查内容
        check_type: 审查类型 ("non_parameter", "parameter", "both")

    Returns:
        完整的审查结果，包含RAG召回和AI审查结果
    """
    global milvus_manager, ai_review_engine

    # 初始化
    if milvus_manager is None:
        init_milvus()
    if ai_review_engine is None:
        init_ai_review_engine()

    trace_id = f"professional_review_{int(time.time() * 1000)}"
    rag_monitor.start_trace(trace_id, metadata={
        "content_length": len(review_content),
        "check_type": check_type,
        "stage": "professional_review_test"
    })

    logger.info(f"[专业性审查测试] 开始处理, trace_id: {trace_id}, 类型: {check_type}")

    try:
        # ==================== Step 1: RAG召回 ====================
        logger.info("[专业性审查测试] Step 1: 开始RAG召回")

        # 1.1 查询提取
        query_pairs = query_rewrite_manager.query_extract(review_content)

        # 检查query_pairs是否为None或空
        if not query_pairs:
            logger.error("[专业性审查测试] 查询提取失败或返回空结果")
            return {
                "status": "error",
                "trace_id": trace_id,
                "error": "查询提取失败,请检查模型服务是否正常",
                "rag_summary": {
                    "query_pairs_count": 0,
                    "entity_results_count": 0
                },
                "review_results": []
            }

        logger.info(f"[专业性审查测试] 提取到 {len(query_pairs)} 个查询对")

        # 1.2 实体增强检索
        bfp_result_lists = entity_enhance.entities_enhance_retrieval(query_pairs)

        if not bfp_result_lists:
            logger.warning("[专业性审查测试] 实体检索未返回结果")
            return _build_professional_empty_result(trace_id, "no_rag_results")

        # 1.3 父文档增强
        try:
            enhancement_result = enhance_with_parent_docs_grouped(
                milvus_manager,
                bfp_result_lists,
                score_threshold=0.3,
                max_parents_per_pair=3
            )
            enhanced_results = enhancement_result['enhanced_results']
            logger.info(f"[专业性审查测试] 父文档增强完成: {enhancement_result.get('enhanced_pairs', 0)}/{enhancement_result.get('total_pairs', 0)} 个查询对")
        except Exception as e:
            logger.error(f"[专业性审查测试] 父文档增强失败: {e}", exc_info=True)
            enhanced_results = bfp_result_lists

        # 1.4 提取结果
        entity_results = extract_query_pairs_results(enhanced_results, query_pairs, score_threshold=0.5) if enhanced_results else []

        if not entity_results:
            logger.warning("[专业性审查测试] 没有结果通过阈值过滤（得分>0.5）")
            return _build_professional_empty_result(trace_id, "no_high_score_results")

        logger.info(f"[专业性审查测试] RAG召回完成，获得 {len(entity_results)} 个高分结果")

        # ==================== Step 2: 专业性审查 ====================
        logger.info("[专业性审查测试] Step 2: 开始专业性审查")

        review_results = []

        for idx, entity_result in enumerate(entity_results):
            entity = entity_result.get('entity', '')
            combined_query = entity_result.get('combined_query', '')
            text_content = entity_result.get('text_content', '')
            file_name = entity_result.get('file_name', '')
            final_score = entity_result.get('final_score', 0)

            logger.info(f"[专业性审查测试] 处理实体 [{idx+1}/{len(entity_results)}]: {entity}, 得分: {final_score:.4f}")

            # 准备审查参数
            trace_id_idx = f"{trace_id}_entity_{idx}"
            review_references = text_content  # 召回的参考文档
            reference_source = file_name
            state = {
                "callback_task_id": "test_task_id",
                "progress_manager": None
            }
            stage_name = "专业性审查测试"

            # 执行审查
            entity_review_result = {
                "entity": entity,
                "combined_query": combined_query,
                "reference_source": reference_source,
                "rag_score": final_score,
                "non_parameter_result": None,
                "parameter_result": None
            }

            # 2.1 非参数合规性检查
            if check_type in ["non_parameter", "both"]:
                try:
                    logger.info(f"[专业性审查测试] 执行非参数审查: {entity}")
                    non_param_result = await ai_review_engine.check_non_parameter_compliance(
                        trace_id_idx=trace_id_idx,
                        review_content=review_content,
                        review_references=review_references,
                        reference_source=reference_source,
                        state=state,
                        stage_name=stage_name,
                        entity_query=combined_query  # 传入实体查询条文
                    )
                    # 将ReviewResult对象转换为字典
                    if hasattr(non_param_result, '__dict__'):
                        entity_review_result["non_parameter_result"] = {
                            'success': non_param_result.success,
                            'details': non_param_result.details,
                            'error_message': non_param_result.error_message,
                            'execution_time': non_param_result.execution_time
                        }
                    else:
                        entity_review_result["non_parameter_result"] = non_param_result
                    logger.info(f"[专业性审查测试] 非参数审查完成: {entity}")
                except Exception as e:
                    logger.error(f"[专业性审查测试] 非参数审查失败: {e}", exc_info=True)
                    entity_review_result["non_parameter_result"] = {"error": str(e)}

            # 2.2 参数合规性检查
            if check_type in ["parameter", "both"]:
                try:
                    logger.info(f"[专业性审查测试] 执行参数审查: {entity}")
                    param_result = await ai_review_engine.check_parameter_compliance(
                        trace_id_idx=trace_id_idx,
                        review_content=review_content,
                        review_references=review_references,
                        reference_source=reference_source,
                        state=state,
                        stage_name=stage_name,
                        entity_query=combined_query  # 传入实体查询条文
                    )
                    # 将ReviewResult对象转换为字典
                    if hasattr(param_result, '__dict__'):
                        entity_review_result["parameter_result"] = {
                            'success': param_result.success,
                            'details': param_result.details,
                            'error_message': param_result.error_message,
                            'execution_time': param_result.execution_time
                        }
                    else:
                        entity_review_result["parameter_result"] = param_result
                    logger.info(f"[专业性审查测试] 参数审查完成: {entity}")
                except Exception as e:
                    logger.error(f"[专业性审查测试] 参数审查失败: {e}", exc_info=True)
                    entity_review_result["parameter_result"] = {"error": str(e)}

            review_results.append(entity_review_result)

        # ==================== Step 3: 构建结果 ====================
        final_result = {
            "status": "success",
            "trace_id": trace_id,
            "check_type": check_type,
            "rag_summary": {
                "query_pairs_count": len(query_pairs),
                "entity_results_count": len(entity_results),
                "query_pairs": query_pairs,
                "entity_results": entity_results
            },
            "review_results": review_results,
            "total_entities_reviewed": len(review_results)
        }

        # 专业性审查直接返回结果,不需要包装
        logger.info(f"✅ 专业性审查测试完成, trace_id: {trace_id}")
        return final_result

    except Exception as e:
        logger.error(f"[专业性审查测试] 处理失败: {e}", exc_info=True)
        return {
            "status": "error",
            "trace_id": trace_id,
            "error": str(e)
        }
    finally:
        rag_monitor.end_trace(trace_id)


def _build_professional_empty_result(trace_id: str, reason: str) -> dict:
    """构建专业性审查的空结果"""
    return {
        "status": "no_results",
        "trace_id": trace_id,
        "reason": reason,
        "message": "未获取到有效的RAG召回结果",
        "rag_summary": {
            "query_pairs_count": 0,
            "entity_results_count": 0
        },
        "review_results": []
    }


def rag_enhanced_check(query_content: str) -> dict:
    """
    RAG增强检查 - 完整链路（使用装饰器监控版本）
    复用 ai_review_engine.py rag_enhanced_check 方法的逻辑
    """
    global milvus_manager
    if milvus_manager is None:
        init_milvus()

    # 开始追踪会话
    trace_id = f"rag_{int(time.time() * 1000)}"
    rag_monitor.start_trace(trace_id, metadata={
        "content_length": len(query_content),
        "content_preview": query_content[:100] if len(query_content) > 100 else query_content,
        "stage": "rag_enhanced_check"
    })

    logger.info(f"[RAG增强] 开始处理, trace_id: {trace_id}, 内容长度: {len(query_content)}")

    try:
        # Step 1: 查询提取（使用装饰器监控）
        @rag_monitor.monitor_step(
            step_name="1_query_extract",
            capture_input=True,
            capture_output=True,
            input_transform=lambda x: {
                "content_length": len(query_content),
                "content_preview": query_content[:200] if len(query_content) > 200 else query_content
            }
        )
        def step1_query_extract(content):
            return query_rewrite_manager.query_extract(content)

        query_pairs = step1_query_extract(query_content)
        logger.info(f"[RAG增强] 提取到 {len(query_pairs)} 个查询对")

        # Step 2: 实体增强检索
        @rag_monitor.monitor_step(
            step_name="2_entity_enhance_retrieval",
            capture_input=True,
            capture_output=True,
            output_transform=lambda x: {
                "results_count": len(x) if x else 0,
                "has_results": bool(x)
            } if isinstance(x, list) else x
        )
        def step2_entity_retrieval(pairs):
            return entity_enhance.entities_enhance_retrieval(pairs)

        bfp_result_lists = step2_entity_retrieval(query_pairs)

        # 检查检索结果
        if not bfp_result_lists:
            logger.warning("[RAG增强] 实体检索未返回结果")
            final_result = _build_empty_result()
            return _finalize_pipeline_result(trace_id, final_result)

        # Step 3: 父文档增强（使用分组增强策略）
        @rag_monitor.monitor_step(
            step_name="3_parent_doc_enhancement",
            capture_input=True,
            capture_output=True,
            input_transform=lambda x: {
                "bfp_results_count": len(bfp_result_lists),
                "score_threshold": 0.3,
                "max_parents_per_pair": 3
            },
            output_transform=lambda x: {
                "enhanced_count": x.get("enhanced_count", 0),
                "parent_docs_count": len(x.get("parent_docs", [])),
                "enhanced_pairs": x.get("enhanced_pairs", 0),
                "total_pairs": x.get("total_pairs", 0)
            } if isinstance(x, dict) else {"error": str(x)}
        )
        def step3_parent_enhancement(results):
            return enhance_with_parent_docs_grouped(
                milvus_manager,
                results,
                score_threshold=0.3,  # bfp_rerank_score 阈值
                max_parents_per_pair=3  # 每个查询对最多3个父文档
            )

        try:
            enhancement_result = step3_parent_enhancement(bfp_result_lists)
            enhanced_results = enhancement_result['enhanced_results']
            enhanced_count = enhancement_result['enhanced_count']
            parent_docs = enhancement_result['parent_docs']
            enhanced_pairs = enhancement_result.get('enhanced_pairs', 0)
            total_pairs = enhancement_result.get('total_pairs', 0)

            logger.info(f"[RAG增强] 分组增强完成: {enhanced_pairs}/{total_pairs} 个查询对进行了增强")
            logger.info(f"[RAG增强] 成功增强 {enhanced_count} 个结果，使用了 {len(parent_docs)} 个父文档")

        except Exception as e:
            logger.error(f"[RAG增强] 父文档增强失败: {e}", exc_info=True)
            enhanced_results = bfp_result_lists

        # Step 4: 提取结果（按查询对区分，只保留得分>0.5的结果）
        @rag_monitor.monitor_step(
            step_name="4_extract_query_pairs_results",
            capture_input=True,
            capture_output=True,
            input_transform=lambda x: {
                "enhanced_results_count": len(enhanced_results) if enhanced_results else 0,
                "query_pairs_count": len(query_pairs),
                "score_threshold": 0.5
            },
            output_transform=lambda x: {
                "entities_count": len(x) if x else 0,
                "has_results": bool(x)
            } if isinstance(x, list) else x
        )
        def step4_extract_results(results, pairs):
            return extract_query_pairs_results(results, pairs, score_threshold=0.5) if results else []

        entity_results = step4_extract_results(enhanced_results, query_pairs)

        # 构建最终结果
        final_result = _build_final_result(entity_results)

        return _finalize_pipeline_result(trace_id, final_result)

    except Exception as e:
        logger.error(f"[RAG增强] 处理失败: {e}", exc_info=True)
        raise
    finally:
        # 结束追踪并保存（即使出错也要保存）
        rag_monitor.end_trace(trace_id)


def _build_empty_result() -> dict:
    """构建空结果"""
    return {
        'retrieval_status': 'no_results',
        'entity_results': [],
        'total_entities': 0,
        'vector_search': [],
        'file_name': '',
        'text_content': '',
        'metadata': {}
    }


def _build_final_result(entity_results: list) -> dict:
    """构建最终结果"""
    if not entity_results:
        return {
            'retrieval_status': 'no_results',
            'entity_results': [],
            'total_entities': 0,
            'message': '没有结果通过阈值过滤（得分>0.5）'
        }
    return {
        'retrieval_status': 'success',
        'entity_results': entity_results,
        'total_entities': len(entity_results)
    }


def _finalize_pipeline_result(trace_id: str, final_result: dict) -> dict:
    """
    完成管道结果处理，合并监控数据和业务结果

    Args:
        trace_id: 追踪会话ID
        final_result: 业务最终结果

    Returns:
        合并后的完整结果
    """
    # 获取监控数据
    trace_data = rag_monitor.get_trace_data(trace_id)

    if trace_data:
        # 合并监控数据和业务结果
        pipeline_data = {
            "trace_id": trace_id,
            "stage": "rag_enhanced_check",
            "timestamp": trace_data.get("start_time"),
            "total_execution_time": trace_data.get("total_duration", 0),
            "metadata": trace_data.get("metadata", {}),
            "steps": trace_data.get("steps", {}),
            "final_result": final_result
        }

        # 保存到旧的位置（兼容前端）
        os.makedirs(os.path.join(project_root, "temp", "rag_pipeline_server"), exist_ok=True)
        with open(os.path.join(project_root, "temp", "rag_pipeline_server", "rag_pipeline_data.json"), "w", encoding='utf-8') as f:
            json.dump(pipeline_data, f, ensure_ascii=False, indent=2, default=str)

        return pipeline_data
    else:
        # 如果没有监控数据，直接返回业务结果
        logger.warning(f"[RAG增强] 未找到追踪数据: {trace_id}")
        return final_result


class RAGPipelineHandler(SimpleHTTPRequestHandler):
    """HTTP请求处理器"""

    def do_GET(self):
        """处理GET请求"""
        parsed = urlparse(self.path)

        if parsed.path == '/api/health':
            self.send_json_response({'status': 'ok', 'milvus_ready': milvus_manager is not None})
        elif parsed.path == '/api/data':
            # 返回最新的pipeline数据
            data_path = os.path.join(project_root, "temp", "rag_pipeline_server", "rag_pipeline_data.json")
            if os.path.exists(data_path):
                with open(data_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                self.send_json_response(data)
            else:
                self.send_json_response({'error': '数据文件不存在'}, 404)
        else:
            # 静态文件服务
            super().do_GET()

    def do_POST(self):
        """处理POST请求"""
        parsed = urlparse(self.path)

        if parsed.path == '/api/rag':
            content_length = int(self.headers['Content-Length'])
            post_data = self.rfile.read(content_length)

            try:
                body = json.loads(post_data.decode('utf-8'))
                query_content = body.get('content', '')

                if not query_content:
                    self.send_json_response({'error': '请提供content参数'}, 400)
                    return

                print(f"\n📝 收到RAG请求, 内容长度: {len(query_content)}")
                result = rag_enhanced_check(query_content)
                print(f"✅ RAG处理完成, 耗时: {result.get('total_execution_time', 0)}秒")

                self.send_json_response(result)

            except json.JSONDecodeError:
                self.send_json_response({'error': 'JSON解析失败'}, 400)
            except Exception as e:
                logger.error(f"RAG处理失败: {e}", exc_info=True)
                self.send_json_response({'error': str(e)}, 500)

        elif parsed.path == '/api/init':
            try:
                init_milvus()
                self.send_json_response({'status': 'ok', 'message': 'Milvus初始化成功'})
            except Exception as e:
                self.send_json_response({'error': str(e)}, 500)

        elif parsed.path == '/api/professional_review':
            # 新增：专业性审查测试API
            content_length = int(self.headers['Content-Length'])
            post_data = self.rfile.read(content_length)

            try:
                body = json.loads(post_data.decode('utf-8'))
                review_content = body.get('content', '')
                check_type = body.get('check_type', 'both')  # "non_parameter", "parameter", "both"

                if not review_content:
                    self.send_json_response({'error': '请提供content参数'}, 400)
                    return

                if check_type not in ['non_parameter', 'parameter', 'both']:
                    self.send_json_response({'error': 'check_type必须是 non_parameter, parameter 或 both'}, 400)
                    return

                print(f"\n📝 收到专业性审查测试请求, 类型: {check_type}, 内容长度: {len(review_content)}")

                # 运行异步函数
                result = run_async(professional_review_test(review_content, check_type))

                print(f"✅ 专业性审查测试完成")
                self.send_json_response(result)

            except json.JSONDecodeError:
                self.send_json_response({'error': 'JSON解析失败'}, 400)
            except Exception as e:
                logger.error(f"专业性审查测试失败: {e}", exc_info=True)
                self.send_json_response({'error': str(e)}, 500)

        else:
            self.send_json_response({'error': 'Not Found'}, 404)

    def do_OPTIONS(self):
        """处理CORS预检请求"""
        self.send_response(200)
        self.send_header('Access-Control-Allow-Origin', '*')
        self.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS')
        self.send_header('Access-Control-Allow-Headers', 'Content-Type')
        self.end_headers()

    def send_json_response(self, data, status=200):
        """发送JSON响应"""
        self.send_response(status)
        self.send_header('Content-Type', 'application/json; charset=utf-8')
        self.send_header('Access-Control-Allow-Origin', '*')
        self.end_headers()
        self.wfile.write(json.dumps(data, ensure_ascii=False, default=str).encode('utf-8'))

    def end_headers(self):
        self.send_header('Access-Control-Allow-Origin', '*')
        super().end_headers()


def run_server(port=8765):
    """启动服务器"""
    web_dir = os.path.dirname(os.path.abspath(__file__))
    
    # 切换到项目根目录，确保内部模块的相对路径正确
    os.chdir(project_root)
    # 确保temp目录存在
    os.makedirs(os.path.join(project_root, "temp", "rag_pipeline_server"), exist_ok=True)

    # 自定义Handler，指定静态文件目录
    class CustomHandler(RAGPipelineHandler):
        def translate_path(self, path):
            # 静态文件从web目录提供
            if path.startswith('/api'):
                return path
            # 去掉开头的/
            path = path.lstrip('/')
            return os.path.join(web_dir, path)

    server = HTTPServer(('0.0.0.0', port), CustomHandler)
    print(f"\n{'='*70}")
    print(f"🚀 RAG管道测试服务器已启动")
    print(f"{'='*70}")
    print(f"📍 访问地址: http://localhost:{port}")
    print(f"📍 工作目录: {project_root}")
    print(f"📍 API端点:")
    print(f"   POST /api/rag                    - 执行RAG检索")
    print(f"   POST /api/professional_review    - 专业性审查完整测试（RAG+AI审查）")
    print(f"   POST /api/init                   - 初始化Milvus")
    print(f"   GET  /api/data                   - 获取最新数据")
    print(f"   GET  /api/health                 - 健康检查")
    print(f"{'='*70}")
    print(f"\n💡 专业性审查测试示例:")
    print(f"   curl -X POST http://localhost:{port}/api/professional_review \\")
    print(f"     -H \"Content-Type: application/json\" \\")
    print(f"     -d '{{\"content\": \"待审查内容\", \"check_type\": \"both\"}}'")
    print(f"\n   check_type 可选值: non_parameter, parameter, both")
    print(f"{'='*70}\n")

    # 预初始化Milvus
    try:
        init_milvus()
    except Exception as e:
        print(f"⚠️  Milvus初始化失败: {e}")
        print("   可通过 POST /api/init 重新初始化\n")

    server.serve_forever()


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='RAG管道测试服务器')
    parser.add_argument('--port', type=int, default=8765, help='服务端口 (默认: 8765)')
    args = parser.parse_args()

    run_server(args.port)
