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

"""
@Project   : lq-agent-api
@File      : parent_tool.py
@IDE       : VsCode
@Author    : 王旭明
@Date      : 2025-12-30
@Description: 父文档增强工具 - 提供父文档查询和拼接功能
"""

from typing import Any, Dict, List, Optional

from foundation.observability.logger.loggering import review_logger as logger


# =============================== 配置 ===============================

PARENT_COLLECTION = "rag_parent_hybrid"  # 父文档集合名称
PARENT_OUTPUT_FIELDS = ["parent_id", "text"]  # 父文档查询字段


# =============================== 核心工具函数 ===============================

def fetch_parent_document(
    milvus_manager,
    parent_id: str,
    collection_name: str = PARENT_COLLECTION,
    output_fields: List[str] = PARENT_OUTPUT_FIELDS
) -> Optional[Dict[str, Any]]:
    """
    查询单个父文档

    Args:
        milvus_manager: MilvusManager 实例
        parent_id: 父文档ID
        collection_name: 集合名称
        output_fields: 需要返回的字段列表

    Returns:
        父文档字典,如果不存在返回None
    """
    try:
        rows = milvus_manager.condition_query(
            collection_name=collection_name,
            filter=f"parent_id == '{parent_id}'",
            output_fields=output_fields,
            limit=1,
        )

        if not rows or not rows[0]:
            logger.debug(f"[父文档工具] 父文档 {parent_id} 不存在")
            return None

        # 只返回需要的字段
        result = {k: rows[0].get(k) for k in output_fields if k in rows[0]}
        logger.debug(f"[父文档工具] 父文档 {parent_id} 查询成功, 内容长度: {len(result.get('text', ''))}")
        return result

    except Exception as e:
        logger.error(f"[父文档工具] 查询父文档 {parent_id} 失败: {e}")
        return None


def fetch_parent_chunks_by_parent_id(
    milvus_manager,
    parent_id: str,
    collection_name: str = "rag_parent_hybrid",  # 子块集合
    output_fields: List[str] = None
) -> Optional[List[Dict[str, Any]]]:
    """
    通过 parent_id 召回所有子块片段

    Args:
        milvus_manager: MilvusManager 实例
        parent_id: 父文档ID
        collection_name: 集合名称 (默认为子块集合 rag_children_hybrid)
        output_fields: 需要返回的字段列表

    Returns:
        子块片段列表，按 pk 排序，如果不存在返回 None
    """
    if output_fields is None:
        output_fields = ["pk", "text", "parent_id", "file_name", "title"]

    try:
        rows = milvus_manager.condition_query(
            collection_name=collection_name,
            filter=f"parent_id == '{parent_id}'",
            output_fields=output_fields,
            limit=1000,  # 足够大的数字，获取所有片段
        )
        
        if not rows:
            logger.warning(f"[父文档工具] parent_id {parent_id} 没有召回任何片段")
            return None

        # 转换为字典列表
        chunks = []
        for row in rows:
            chunk = {}
            for field in output_fields:
                if field in row:
                    value = row[field]
                    # 处理字节类型
                    if isinstance(value, bytes):
                        value = value.decode('utf-8')
                    chunk[field] = value
            chunks.append(chunk)

        logger.info(f"[父文档工具] parent_id {parent_id} 召回了 {len(chunks)} 个片段")
        return chunks

    except Exception as e:
        logger.error(f"[父文档工具] 召回 parent_id {parent_id} 的片段失败: {e}")
        return None



def enhance_with_parent_docs_grouped(
    milvus_manager,
    bfp_result_lists: List,
    score_threshold: float = 0.5,
    max_parents_per_pair: int = 3,
    max_parent_text_length: int = 8000
) -> Dict[str, Any]:
    """
    分组增强 + 按分数筛选 + 智能扩展（每个查询对独立处理）

    核心逻辑:
    1. 每个查询对独立处理，按 bfp_rerank_score 筛选
    2. 只保留并增强高分结果，低分查询对直接跳过
    3. 用父ID从Milvus召回父文档内容
    4. 根据子块在父文档中的位置进行智能扩展：
       - Top1: 前后总共扩展 3000 字
       - Top2: 前后总共扩展 1500 字
       - Top3: 前后总共扩展 1500 字
    5. 如果总长度超过 12000 字，按比例压缩所有父块

    Args:
        milvus_manager: MilvusManager 实例
        bfp_result_lists: 检索结果列表 (二维，每个子列表对应一个查询对)
        score_threshold: bfp_rerank_score 最低阈值，低于此分数直接跳过 (默认0.5)
        max_parents_per_pair: 每个查询对最多选取的父文档数量 (默认3个，已硬编码为Top3)
        max_parent_text_length: 单个父文档最大长度限制（已弃用，改为智能扩展）

    Returns:
        Dict: 增强结果，包含:
            - enhanced_results: 增强后的结果列表 (二维，只包含高分查询对)
            - enhanced_count: 成功增强的结果数量
            - parent_docs: 使用的所有父文档列表
            - enhanced_pairs: 进行了增强的查询对数量
            - total_pairs: 原始查询对总数
    """
    enhanced_results = []
    total_enhanced_count = 0
    all_parent_docs = []
    enhanced_pairs_count = 0

    logger.info(f"[分组增强] 开始处理 {len(bfp_result_lists)} 个查询对，阈值={score_threshold}")

    for pair_idx, result_list in enumerate(bfp_result_lists):
        if not result_list:
            continue

        # 1. 按分数排序并筛选高分结果片段 
        sorted_results = sorted(result_list, key=lambda x: x.get('bfp_rerank_score', 0), reverse=True)
        high_score_results = [r for r in sorted_results if r.get('bfp_rerank_score', 0) >= score_threshold]

        if not high_score_results:
            logger.info(f"[分组增强] 查询对 {pair_idx}: 所有结果分数均低于 {score_threshold}，跳过")
            continue

        # 2. 对 parent_id 去重
        parent_ids = list(set([
            r.get('metadata', {}).get('parent_id')
            for r in high_score_results
            if r.get('metadata', {}).get('parent_id')
        ]))

        # 从第一个结果中提取 source_entity（用于日志）
        source_entity = high_score_results[0].get('source_entity', '') if high_score_results else ''

        # 限制父文档数量为 top3
        parent_ids = parent_ids[:3]
        logger.debug(f"实体：{source_entity},[分组增强] 父文档数量限制为 3，实际数量为 {len(parent_ids)},{parent_ids}")

        # 3. 用 parent_id 召回片段并按 pk 排序拼接（优化：根据子块位置扩展）
        parent_id_to_doc = {}
        parent_id_to_metadata = {}  # 存储元数据

        # 定义扩展策略：top1=3000字，top2=1500字，top3=1500字
        expansion_limits = [3000, 1500, 1500]

        for idx, pid in enumerate(parent_ids):
            # 用 parent_id 查询 Milvus，召回该父文档的所有片段
            chunks = fetch_parent_chunks_by_parent_id(milvus_manager, pid)

            if chunks:
                # 按 pk 排序
                sorted_chunks = sorted(chunks, key=lambda x: x.get('pk', 0))
                logger.debug(f"排序结果：{pid}:{[c.get('pk') for c in sorted_chunks]}")

                # 提取元数据（从第一个片段获取）
                first_chunk = sorted_chunks[0]
                file_name = first_chunk.get('file_name', '')
                title = first_chunk.get('title', '')

                # 找到高分子块在父文档中的位置
                # 从 high_score_results 中找到对应 parent_id 的子块
                matching_result = next(
                    (r for r in high_score_results if str(r.get('metadata', {}).get('parent_id')) == str(pid)),
                    None
                )

                if matching_result:
                    # 获取子块的 text_content（用于定位）
                    child_text = matching_result.get('text_content', '')

                    # 拼接所有片段为完整父文档
                    full_text = "\n".join([c.get('text', '') for c in sorted_chunks])

                    # 在父文档中找到子块的位置
                    child_position = full_text.find(child_text)

                    if child_position >= 0:
                        # 计算子块中心位置
                        child_center = child_position + len(child_text) // 2

                        # 根据 top 级别确定扩展长度
                        expansion_limit = expansion_limits[idx] if idx < len(expansion_limits) else 1500
                        half_expansion = expansion_limit // 2

                        # 计算扩展范围
                        start_pos = max(0, child_center - half_expansion)
                        end_pos = min(len(full_text), child_center + half_expansion)

                        # 提取扩展后的文本
                        expanded_text = full_text[start_pos:end_pos]

                        # 添加截断提示
                        if start_pos > 0:
                            expanded_text = f"...[前文已省略]...\n{expanded_text}"
                        if end_pos < len(full_text):
                            expanded_text = f"{expanded_text}\n...[后文已省略]..."

                        logger.info(f"[分组增强] Top{idx+1} parent_id={str(pid)[:8]}... 子块中心位置={child_center}, 扩展范围=[{start_pos}, {end_pos}], 扩展长度={len(expanded_text)}字符")
                    else:
                        # 子块未找到，回退到完整父文档
                        expanded_text = full_text
                        logger.warning(f"[分组增强] parent_id={str(pid)[:8]}... 子块未在父文档中找到，使用完整父文档")
                else:
                    # 没有匹配的高分结果，回退到完整父文档
                    expanded_text = "\n".join([c.get('text', '') for c in sorted_chunks])
                    logger.warning(f"[分组增强] parent_id={str(pid)[:8]}... 未找到匹配的高分结果，使用完整父文档")

                # 存储扩展后的文本和元数据
                parent_id_to_doc[str(pid)] = expanded_text
                parent_id_to_metadata[str(pid)] = {
                    'file_name': file_name,
                    'title': title,
                    'top_rank': idx + 1,  # Top1/Top2/Top3
                    'original_length': len("\n".join([c.get('text', '') for c in sorted_chunks])),
                    'expanded_length': len(expanded_text)
                }
                logger.debug(f"[分组增强] Top{idx+1} parent_id={str(pid)[:8]}... 召回了 {len(sorted_chunks)} 个片段，文件={file_name}，扩展后长度={len(expanded_text)}字符")

        if not parent_id_to_doc:
            logger.warning(f"[分组增强] 查询对 {pair_idx}: 所有父文档召回失败，跳过")
            continue

        # 4. 检查总长度并按比例压缩（如果超过 12000 字）
        total_length = sum(len(text) for text in parent_id_to_doc.values())
        max_total_length = 12000

        if total_length > max_total_length:
            logger.warning(f"[分组增强] 查询对 {pair_idx}: 总长度 {total_length} 字符超过限制 {max_total_length}，开始按比例压缩")

            # 计算压缩比例
            compression_ratio = max_total_length / total_length

            # 按比例压缩每个父文档
            for pid, text in parent_id_to_doc.items():
                original_length = len(text)
                target_length = int(original_length * compression_ratio)

                if target_length < original_length:
                    # 从中间截取，保留核心内容
                    start_cut = (original_length - target_length) // 2
                    compressed_text = text[start_cut:start_cut + target_length]

                    # 添加压缩提示
                    compressed_text = f"...[已按比例压缩{compression_ratio:.2%}]...\n{compressed_text}\n...[已压缩]..."
                    parent_id_to_doc[pid] = compressed_text

                    logger.info(f"[分组增强] parent_id={str(pid)[:8]}... 压缩前={original_length}字符，压缩后={len(compressed_text)}字符")

            # 重新计算总长度
            total_length_after = sum(len(text) for text in parent_id_to_doc.values())
            logger.info(f"[分组增强] 查询对 {pair_idx}: 压缩完成，总长度从 {total_length} 压缩至 {total_length_after} 字符")

        # 5. 只保留并增强高分结果（为每个结果添加头部信息和父文档）
        enhanced_list = []

        for result in high_score_results:
            parent_id = str(result.get('metadata', {}).get('parent_id'))
            if parent_id in parent_id_to_doc:
                # 获取父文档和元数据
                parent_text = parent_id_to_doc[parent_id]
                metadata = parent_id_to_metadata.get(parent_id, {})
                file_name = metadata.get('file_name', '')
                title = metadata.get('title', '')
                top_rank = metadata.get('top_rank', 0)

                # 构建头部信息
                header = f"【Top{top_rank} 参考文档】\n【文件】{file_name}\n【标题】{title}\n" if file_name or title else f"【Top{top_rank} 参考文档】\n"

                # 拼接增强内容
                original_text = result.get('text_content', '')
                enhanced_content = original_text + f"\n{header}{parent_text}\n"

                enhanced_list.append({
                    'text_content': enhanced_content,
                    'metadata': result.get('metadata', {}),
                    'hybrid_similarity': result.get('hybrid_similarity'),
                    'rerank_score': result.get('rerank_score'),
                    'bfp_rerank_score': result.get('bfp_rerank_score'),
                    'bfp_rerank_parent_id': result.get('bfp_rerank_parent_id', ''),
                    'source_entity': result.get('source_entity', ''),
                    'enhanced': True,
                    'parent_top_rank': top_rank  # 新增：记录父文档的 Top 排名
                })

        if enhanced_list:
            enhanced_results.append(enhanced_list)
            # 构建父文档列表用于统计
            for pid, text in parent_id_to_doc.items():
                all_parent_docs.append({'parent_id': pid, 'text': text})
            enhanced_pairs_count += 1
            total_enhanced_count += len(enhanced_list)
            logger.info(f"[分组增强] 查询对 {pair_idx}: 保留 {len(enhanced_list)} 个高分结果，使用 {len(parent_id_to_doc)} 个父文档")

    logger.info(f"[分组增强] 完成: {enhanced_pairs_count}/{len(bfp_result_lists)} 个查询对，{total_enhanced_count} 个结果")

    return {
        'enhanced_results': enhanced_results,
        'enhanced_count': total_enhanced_count,
        'parent_docs': all_parent_docs,
        'enhanced_pairs': enhanced_pairs_count,
        'total_pairs': len(bfp_result_lists)
    }


def extract_query_pairs_results(bfp_result_lists: List, query_pairs: List[Dict] = None,
                                score_threshold: float = 0.5) -> List[Dict[str, Any]]:
    """
    从检索结果中提取每个查询对的最高得分结果（得分必须大于阈值），返回列表格式

    得分优先级：
    1. bfp_rerank_score (二次重排序得分，最高优先级)
    2. rerank_score (第一次重排序得分，回退选项)

    Args:
        bfp_result_lists: 检索结果列表（二维列表，每个子列表对应一个查询对）
        query_pairs: 查询对列表（用于映射实体名称、背景、参数）
        score_threshold: 得分阈值，只返回得分大于该阈值的结果（默认0.5）

    Returns:
        List[Dict]: 结果列表，每个元素包含：
        {
            'entity': str,              # 实体名称
            'background': str,          # 背景信息
            'parameter': str,           # 技术参数
            'combined_query': str,      # 组合查询条文 (entity + background + parameter)
            'file_name': str,           # 文件名
            'text_content': str,        # 检索到的文本内容
            'metadata': dict,           # 元数据
            'bfp_rerank_score': float,  # BFP重排序得分
            'rerank_score': float,      # 重排序得分
            'final_score': float,       # 最终得分
            'score_type': str,          # 得分类型
            'query_index': int          # 查询索引
        }
        - 如果所有结果得分都低于阈值，返回空列表
    """
    if not bfp_result_lists:
        logger.warning("[父文档工具] 检索结果为空")
        return []

    logger.info(f"[父文档工具] 开始提取查询对结果，得分阈值: {score_threshold}")

    entity_results = []
    total_count = 0
    filtered_count = 0

    for query_idx, results in enumerate(bfp_result_lists):
        if not results:
            logger.warning(f"[父文档工具] 查询对 {query_idx} 无检索结果")
            continue

        total_count += 1

        # 获取查询对信息
        query_pair = query_pairs[query_idx] if query_pairs and query_idx < len(query_pairs) else {}
        entity = query_pair.get('entity', '')
        background = query_pair.get('background', '')
        parameter = query_pair.get('parameter', '')

        # 组合完整查询条文: 实体 + 背景 + 参数
        combined_parts = []
        if entity:
            combined_parts.append(entity)
        if background:
            combined_parts.append(background)
        if parameter:
            combined_parts.append(parameter)
        combined_query = ' '.join(combined_parts)

        # 优先使用结果中的 source_entity，回退到 query_pairs
        first_result = results[0]
        entity_name = first_result.get('source_entity', entity or f'query_{query_idx}')

        # 在该查询对的结果中找得分最高的
        # 优先使用 bfp_rerank_score，如果不存在或为0，则使用 rerank_score
        best_result = max(results, key=lambda x: (
            x.get('bfp_rerank_score', 0.0),
            x.get('rerank_score', 0.0)
        ))

        # 确定最终使用的得分和类型
        bfp_score = best_result.get('bfp_rerank_score') or 0.0
        rerank_score = best_result.get('rerank_score') or 0.0

        # 如果 bfp_rerank_score 有效（>0），使用它；否则使用 rerank_score
        if bfp_score and bfp_score > 0:
            final_score = bfp_score
            score_type = 'bfp_rerank_score'
        else:
            final_score = rerank_score
            score_type = 'rerank_score'

        logger.info(f"[父文档工具] 实体 '{entity_name}' 最优结果: {score_type}={final_score:.6f} (bfp={bfp_score:.6f}, rerank={rerank_score:.6f})")

        # 只保留得分超过阈值的结果
        if final_score > score_threshold:
            result_item = {
                'entity': entity,
                'background': background,
                'parameter': parameter,
                'combined_query': combined_query,
                'file_name': best_result['metadata'].get('file_name', 'unknown'),
                'text_content': best_result['text_content'],
                'metadata': best_result['metadata'],
                'bfp_rerank_score': bfp_score,
                'rerank_score': rerank_score,
                'final_score': final_score,
                'score_type': score_type,
                'query_index': query_idx
            }
            entity_results.append(result_item)
            filtered_count += 1
            logger.info(f"[父文档工具] ✅ 实体 '{entity_name}' 得分 {final_score:.6f} ({score_type}) > {score_threshold}，保留结果")
            logger.info(f"[父文档工具] 组合查询条文: {combined_query}")
        else:
            logger.warning(f"[父文档工具] ❌ 实体 '{entity_name}' 得分 {final_score:.6f} ({score_type}) <= {score_threshold}，过滤结果")

    logger.info(f"[父文档工具] 提取完成: 总计 {total_count} 个查询对，{filtered_count} 个结果通过阈值过滤")

    return entity_results

