"""
Result Aggregation Function（解码工作流版本）

结果汇总函数：将所有解构结果汇总为最终的JSON结构，包括三点解构（灵感点、目的点、关键点）、选题理解、脚本理解（实质、形式）等所有数据。
"""

from typing import Dict, Any, List, Optional

from src.components.functions.base import BaseFunction
from src.states.what_deconstruction_state import WhatDeconstructionState, ElementNode, FinalOutput
from src.utils.logger import get_logger

logger = get_logger(__name__)


class ResultAggregationFunction(BaseFunction[WhatDeconstructionState, Dict[str, Any]]):
    """结果汇总函数（解码工作流版本）

    功能：
    - 将各Agent的解构结果组装为JSON结构
    - 汇总三点解构结果（灵感点、目的点、关键点）
    - 汇总选题理解结果
    - 汇总脚本理解结果（实质、形式、段落划分等）
    - 验证JSON格式
    """

    def __init__(
        self,
        name: str = "result_aggregation_function",
        description: str = "汇总所有解构结果为最终的JSON结构（包括三点解构、选题理解、脚本理解等所有数据）"
    ):
        super().__init__(name, description)

    def execute(
        self,
        input_data: WhatDeconstructionState,
        context: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """执行结果汇总

        Args:
            input_data: 工作流状态，包含所有解构结果
            context: 上下文信息

        Returns:
            最终的树状JSON结构
        """
        try:
            # 提取视频相关结果
            video_url = input_data.get("video", "")
            
            # 从 state 中获取标题和正文（优先从直接字段获取，兼容text字典）
            title = input_data.get("title", "")
            body_text = input_data.get("body_text", "")
            if not title or not body_text:
                text_data = input_data.get("text", {})
                title = title or text_data.get("title", "")
                body_text = body_text or text_data.get("body", "")

            # 提取三点解构结果
            # 灵感点：可能是数组或包含数组的对象
            inspiration_points_raw = input_data.get("inspiration_points", {})
            if isinstance(inspiration_points_raw, list):
                inspiration_points = inspiration_points_raw
            elif isinstance(inspiration_points_raw, dict):
                # 可能是包含 points 或直接是数组的字典
                inspiration_points = inspiration_points_raw.get("points", inspiration_points_raw.get("inspiration_points", []))
            else:
                inspiration_points = []
            
            # 目的点：保持原始结构（包含 perspective, purposes, total_count）
            purpose_point = input_data.get("purpose_point", {})
            
            # 关键点：保持原始结构（包含 key_points, total_count, root_count 等）
            key_points = input_data.get("key_points", {})

            # 提取选题理解结果 - 保持完整结构（包含主题、描述、覆盖情况等所有字段）
            topic_selection_understanding = input_data.get("topic_selection_understanding", {})

            # 提取脚本理解结果
            # 段落划分相关
            content_category = input_data.get("内容品类", "未知")
            section_list = input_data.get("段落列表", [])
            
            # 实质和形式列表
            substance_list = input_data.get("实质列表", [])
            form_list = input_data.get("形式列表", [])
            
            # 如果没有实质列表，尝试从 substance_final_elements 获取
            if not substance_list:
                substance_list = input_data.get("substance_final_elements", [])
            
            # 如果没有形式列表，尝试从 form_final_elements 获取
            if not form_list:
                form_list = input_data.get("form_final_elements", [])

            # 构建视频基本信息
            video_info = {
                "视频URL": video_url,
                "标题": title,
                "正文": body_text
            }

            # 组装最终JSON（包含所有数据，参考 output_demo_script.json 格式）
            final_result = {
                "视频信息": video_info,
                "三点解构": {
                    "灵感点": inspiration_points,  # 数组格式
                    "目的点": purpose_point,  # 对象，包含 perspective, purposes, total_count
                    "关键点": key_points  # 对象，包含 key_points, total_count, root_count
                },
                "选题理解": topic_selection_understanding,  # 保持完整结构，包含 主题、描述、覆盖情况等
                "脚本理解": {
                    "内容品类": content_category,
                    "段落列表": section_list,
                    "实质列表": substance_list,  # 数组格式
                    "形式列表": form_list  # 数组格式
                }
            }

            # 验证JSON格式
            self._validate_result(final_result)

            substance_count = len(substance_list) if isinstance(substance_list, list) else 0
            form_count = len(form_list) if isinstance(form_list, list) else 0
            logger.info(
                f"结果汇总完成：选题 + 三点解构 + 脚本理解（实质:{substance_count}个, 形式:{form_count}个）"
            )

            return final_result

        except Exception as e:
            logger.error(f"结果汇总失败: {e}", exc_info=True)
            # 返回错误结果（保持格式一致性）
            return {
                "视频信息": {},
                "三点解构": {
                    "灵感点": [],
                    "目的点": {},
                    "关键点": {}
                },
                "选题理解": {},
                "脚本理解": {
                    "内容品类": "未知",
                    "段落列表": [],
                    "实质列表": [],
                    "形式列表": []
                },
                "错误": f"汇总失败: {str(e)}"
            }

    def _build_element_tree(
        self,
        id: str,
        result: Dict[str, Any],
        element_type: str
    ) -> ElementNode:
        """递归构建元素树（视频分析版本中不再使用，保留以保持代码完整性）

        Args:
            id: 节点ID
            result: 解构结果
            element_type: 元素类型（image/text）

        Returns:
            元素节点
        """
        element: ElementNode = {
            "id": id,
            "what": result.get("what", ""),
            "position": result.get("position", ""),  # 添加position字段
            "split_dimension": result.get("split_dimension", ""),  # 添加split_dimension字段
            "描述": result.get("description", result.get("描述", {})),
            "子节点元素关系": result.get("relationships", result.get("子节点元素关系", [])),
            "元素重要性权重": result.get("importance_weight", result.get("元素重要性权重", 0.5)),
            "权重理由": result.get("weight_reason", result.get("权重理由", "")),
            "子节点元素": []
        }

        # 添加图片链接（仅图片元素）
        if element_type == "image" and "image_path" in result:
            element["图片链接"] = result["image_path"]
        elif "图片链接" in result:
            element["图片链接"] = result["图片链接"]

        # 递归处理子节点
        children = result.get("children", result.get("子节点元素", []))
        if children:
            for idx, child in enumerate(children):
                child_element = self._build_element_tree(
                    id=f"{id}_{idx + 1}",
                    result=child,
                    element_type=element_type
                )
                element["子节点元素"].append(child_element)

        return element

    def _validate_result(self, result: Dict[str, Any]) -> None:
        """验证JSON格式（解码工作流版本）

        Args:
            result: 最终结果

        Raises:
            ValueError: 格式不正确时抛出异常
        """
        # 检查必需字段
        if "视频信息" not in result:
            logger.warning("缺少'视频信息'字段")
            result["视频信息"] = {}

        if "三点解构" not in result:
            logger.warning("缺少'三点解构'字段")
            result["三点解构"] = {
                "灵感点": [],
                "目的点": {},
                "关键点": {}
            }

        if "选题理解" not in result:
            logger.warning("缺少'选题理解'字段")
            result["选题理解"] = {}

        if "脚本理解" not in result:
            logger.warning("缺少'脚本理解'字段")
            result["脚本理解"] = {
                "内容品类": "未知",
                "段落列表": [],
                "实质列表": [],
                "形式列表": []
            }

        # 验证三点解构结构
        three_points = result.get("三点解构", {})
        if not isinstance(three_points, dict):
            logger.warning("'三点解构'必须是字典")
            result["三点解构"] = {
                "灵感点": [],
                "目的点": {},
                "关键点": {}
            }
        else:
            # 确保三点字段存在，灵感点是数组
            if "灵感点" not in three_points:
                logger.warning("三点解构缺少'灵感点'字段")
                three_points["灵感点"] = []
            elif not isinstance(three_points["灵感点"], list):
                logger.warning("'灵感点'必须是列表")
                three_points["灵感点"] = []
            
            if "目的点" not in three_points:
                logger.warning("三点解构缺少'目的点'字段")
                three_points["目的点"] = {}
            elif not isinstance(three_points["目的点"], dict):
                logger.warning("'目的点'必须是字典")
                three_points["目的点"] = {}
            
            if "关键点" not in three_points:
                logger.warning("三点解构缺少'关键点'字段")
                three_points["关键点"] = {}
            elif not isinstance(three_points["关键点"], dict):
                logger.warning("'关键点'必须是字典")
                three_points["关键点"] = {}

        # 验证选题理解（保持完整结构，不做拆解）
        if "选题理解" not in result:
            result["选题理解"] = {}
        elif not isinstance(result["选题理解"], dict):
            logger.warning("'选题理解'必须是字典")
            result["选题理解"] = {}

        # 验证脚本理解结构
        script_understanding = result.get("脚本理解", {})
        if not isinstance(script_understanding, dict):
            logger.warning("'脚本理解'必须是字典")
            result["脚本理解"] = {
                "内容品类": "未知",
                "段落列表": [],
                "实质列表": [],
                "形式列表": []
            }
        else:
            # 确保脚本理解字段存在
            if "内容品类" not in script_understanding:
                script_understanding["内容品类"] = "未知"
            if "段落列表" not in script_understanding:
                script_understanding["段落列表"] = []
            if "实质列表" not in script_understanding:
                script_understanding["实质列表"] = []
            if "形式列表" not in script_understanding:
                script_understanding["形式列表"] = []
            
            # 确保列表类型正确
            for field in ["段落列表", "实质列表", "形式列表"]:
                if not isinstance(script_understanding.get(field), list):
                    logger.warning(f"脚本理解中的'{field}'必须是列表")
                    script_understanding[field] = []

    def _validate_element(self, element: Dict[str, Any]) -> None:
        """验证元素节点格式（视频分析版本中不再使用，保留以保持代码完整性）

        Args:
            element: 元素节点

        Raises:
            ValueError: 格式不正确时抛出异常
        """
        required_fields = ["id", "what", "描述", "子节点元素关系", "元素重要性权重", "权重理由", "子节点元素"]

        for field in required_fields:
            if field not in element:
                raise ValueError(f"元素节点缺少'{field}'字段")

        # 验证权重范围
        weight = element["元素重要性权重"]
        if not isinstance(weight, (int, float)) or weight < 0 or weight > 1:
            logger.warning(f"元素权重不在[0, 1]范围内: {weight}，已修正")
            element["元素重要性权重"] = max(0, min(1, weight))

        # 递归验证子节点
        for child in element["子节点元素"]:
            self._validate_element(child)
