# -*- coding: utf-8 -*-
"""
大纲生成工作流
基于 LangGraph 的大纲生成工作流编排

该模块位于 workflows 层，负责：
1. 构建 LangGraph 状态图
2. 定义节点间的转换关系和条件边
3. 编译和管理工作流图实例

依赖 component 层的 OutlineGenerator 提供节点实现
"""

from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from foundation.observability.logger.loggering import write_logger as logger
from ..component.state_models import OutlineGenerationState
from ..component.outline_generator import OutlineGenerator
import os
import json


class OutlineWorkflow:
    """
    大纲生成工作流

    使用 LangGraph 构建完整的大纲生成工作流，包含：
    - 模板加载
    - 大纲生成
    - 相似案例推荐（可选）
    - 相似片段推荐（可选）
    - 知识点关联（可选）

    示例：
        workflow = OutlineWorkflow()
        graph = workflow.build_graph()
        result = await graph.ainvoke(initial_state)
    """

    def __init__(self):
        """初始化大纲生成工作流"""
        self.outline_generator = OutlineGenerator()
        self.checkpoint_saver = MemorySaver()
        self.graph = None  # 编译后的图实例

    def build_graph(self) -> StateGraph:
        """
        构建大纲生成 LangGraph 工作流图

        工作流路径：
            START → start → load_template → generate_outline
                                                    ↓
                              ┌────────────────────┴────────────────────┐
                              │ 条件: should_recommend_similar          │
                              ↓                                         ↓
                    recommend_similar_cases (yes)              skip_similar (no)
                              ↓                                         ↓
                    recommend_similar_fragments                    check_knowledge
                              ↓                                         ↓
                              └────────────────────┬────────────────────┘
                                                   ↓
                              ┌────────────────────┴────────────────────┐
                              │ 条件: should_link_knowledge             │
                              ↓                                         ↓
                    link_knowledge_bases (yes)                    complete (no)
                              ↓
                          complete
                              ↓
                             END

        Returns:
            StateGraph: 编译完成的大纲生成工作流图
        """
        logger.info("开始构建大纲生成 LangGraph 工作流图")

        # 创建状态图
        workflow = StateGraph(OutlineGenerationState)

        # ==================== 添加节点 ====================

        # 开始节点
        workflow.add_node("start", self.outline_generator.start_node)

        # 模板加载节点
        workflow.add_node("load_template", self.outline_generator.load_template_node)

        # 大纲生成节点
        workflow.add_node("generate_outline", self.outline_generator.generate_outline_node)

        # 相似案例推荐节点
        workflow.add_node("recommend_similar_cases", self.outline_generator.recommend_similar_cases_node)

        # 相似片段推荐节点
        workflow.add_node("recommend_similar_fragments", self.outline_generator.recommend_similar_fragments_node)

        # 知识点关联节点
        workflow.add_node("link_knowledge_bases", self.outline_generator.link_knowledge_bases_node)

        # 完成节点
        workflow.add_node("complete", self.outline_generator.complete_node)

        # 错误处理节点
        workflow.add_node("error_handler", self._error_handler_node)

        # 终止节点
        workflow.add_node("terminate", self._terminate_node)

        # ==================== 添加边 ====================

        # 设置入口点
        workflow.set_entry_point("start")

        # start -> load_template
        workflow.add_edge("start", "load_template")

        # load_template -> generate_outline（带条件检查）
        workflow.add_conditional_edges(
            "load_template",
            self._check_terminate_or_error,
            {
                "terminate": "terminate",
                "error": "error_handler",
                "continue": "generate_outline"
            }
        )

        # generate_outline -> 条件分支（是否启用相似推荐）
        workflow.add_conditional_edges(
            "generate_outline",
            self._check_recommend_similar_or_error,
            {
                "terminate": "terminate",
                "error": "error_handler",
                "yes": "recommend_similar_cases",
                "no": "check_knowledge"
            }
        )

        # recommend_similar_cases -> recommend_similar_fragments（带条件检查）
        workflow.add_conditional_edges(
            "recommend_similar_cases",
            self._check_terminate_or_error,
            {
                "terminate": "terminate",
                "error": "error_handler",
                "continue": "recommend_similar_fragments"
            }
        )

        # recommend_similar_fragments -> 条件分支（是否启用知识点关联）
        workflow.add_conditional_edges(
            "recommend_similar_fragments",
            self._check_link_knowledge_or_error,
            {
                "terminate": "terminate",
                "error": "error_handler",
                "yes": "link_knowledge_bases",
                "no": "complete"
            }
        )

        # link_knowledge_bases -> complete（带条件检查）
        workflow.add_conditional_edges(
            "link_knowledge_bases",
            self._check_terminate_or_error,
            {
                "terminate": "terminate",
                "error": "error_handler",
                "continue": "complete"
            }
        )

        # check_knowledge 是虚拟节点，使用条件边直接处理
        # 但 LangGraph 要求条件边必须从实际节点出发
        # 所以我们需要在 generate_outline 的 "no" 分支直接到 "check_knowledge" 逻辑
        # 这里我们添加一个直通边，让流程继续到知识点关联判断

        # 为了简化，我们添加一个条件判断辅助节点
        workflow.add_node("check_knowledge", self._check_knowledge_node)

        # generate_outline 的 "no" 分支 -> check_knowledge
        # 这需要在 add_conditional_edges 中直接指定
        # 修改：我们在条件函数中处理

        # 修正：从 recommend_similar_fragments 和 check_knowledge 都连接到 complete 或 link_knowledge_bases
        # 但我们需要让 generate_outline 的 "no" 分支也能到达 check_knowledge
        workflow.add_edge("check_knowledge", "complete")

        # complete -> END
        workflow.add_edge("complete", END)

        # error_handler -> END
        workflow.add_edge("error_handler", END)

        # terminate -> END
        workflow.add_edge("terminate", END)

        # ==================== 编译工作流 ====================

        # 编译图，使用内存检查点保存器
        compiled_graph = workflow.compile(checkpointer=self.checkpoint_saver)

        # 保存工作流图到 temp/construction_write 目录
        self._save_workflow_graph(compiled_graph, "temp/construction_write/outline_workflow.png")

        logger.info("大纲生成 LangGraph 工作流图构建完成")
        return compiled_graph

    def _save_workflow_graph(self, compiled_graph: StateGraph, output_path: str):
        """
        保存 LangGraph 工作流图为 PNG 图片

        Args:
            compiled_graph: 编译后的 LangGraph 工作流图
            output_path: 输出文件路径
        """
        try:
            # 确保输出目录存在
            output_dir = os.path.dirname(output_path)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir, exist_ok=True)
                logger.info(f"创建输出目录：{output_dir}")

            # 使用 graphviz 保存图片
            # 需要安装 graphviz 和 graphviz Python 包
            try:
                graph = compiled_graph.get_graph()
                graph_image = graph.draw_mermaid_png()
                with open(output_path, "wb") as f:
                    f.write(graph_image)
                logger.info(f"工作流图已保存到：{output_path}")
            except Exception as graphviz_error:
                logger.warning(f"Graphviz 保存失败：{str(graphviz_error)}, 尝试使用 JSON 格式保存")
                # 备用方案：保存为 JSON 格式
                json_path = output_path.replace(".png", ".json")
                with open(json_path, "w", encoding="utf-8") as f:
                    json.dump(compiled_graph.get_graph().to_json(), f, indent=2, ensure_ascii=False)
                logger.info(f"工作流图已保存到 (JSON 格式): {json_path}")

        except Exception as e:
            logger.warning(f"保存工作流图失败：{str(e)}")

    # ==================== 条件判断函数 ====================

    def _check_terminate_or_error(self, state: OutlineGenerationState) -> str:
        """
        检查是否终止或出错

        Returns:
            str: "terminate", "error", 或 "continue"
        """
        if state.get("overall_task_status") == "terminated":
            return "terminate"
        if state.get("overall_task_status") == "failed" or state.get("error_message"):
            return "error"
        return "continue"

    def _check_recommend_similar_or_error(self, state: OutlineGenerationState) -> str:
        """
        检查是否启用相似推荐，或终止/出错

        Returns:
            str: "terminate", "error", "yes", 或 "no"
        """
        # 先检查终止/错误状态
        if state.get("overall_task_status") == "terminated":
            return "terminate"
        if state.get("overall_task_status") == "failed" or state.get("error_message"):
            return "error"

        # 检查是否启用相似推荐
        similarity_config = state.get("similarity_config", {})
        if similarity_config and similarity_config.get("max_cases", 0) > 0:
            return "yes"
        return "no"

    def _check_link_knowledge_or_error(self, state: OutlineGenerationState) -> str:
        """
        检查是否启用知识点关联，或终止/出错

        Returns:
            str: "terminate", "error", "yes", 或 "no"
        """
        # 先检查终止/错误状态
        if state.get("overall_task_status") == "terminated":
            return "terminate"
        if state.get("overall_task_status") == "failed" or state.get("error_message"):
            return "error"

        # 检查是否启用知识点关联
        knowledge_config = state.get("knowledge_config", {})
        if knowledge_config.get("enable_basis", True):
            return "yes"
        return "no"

    # ==================== 节点实现 ====================

    async def _check_knowledge_node(self, state: OutlineGenerationState) -> dict:
        """
        检查知识点关联条件的辅助节点

        这个节点只用于流程控制，不产生实际业务逻辑
        """
        logger.debug(f"检查知识点关联条件: {state['callback_task_id']}")

        # 根据配置决定是否跳转到 link_knowledge_bases
        knowledge_config = state.get("knowledge_config", {})
        if knowledge_config.get("enable_basis", True):
            # 启用知识点关联，需要跳转到 link_knowledge_bases
            # 但 LangGraph 中不能直接跳转，需要通过条件边
            # 这里我们只记录状态，让后续条件边处理
            return {
                "current_stage": "check_knowledge",
                "messages": state["messages"] + [{"role": "system", "content": "检查知识点关联条件"}]
            }
        else:
            # 禁用知识点关联，直接到 complete
            return {
                "current_stage": "check_knowledge",
                "messages": state["messages"] + [{"role": "system", "content": "跳过知识点关联"}]
            }

    async def _error_handler_node(self, state: OutlineGenerationState) -> dict:
        """
        错误处理节点

        Args:
            state: 当前状态

        Returns:
            dict: 错误状态
        """
        error_message = state.get("error_message", "未知错误")
        logger.error(f"大纲生成工作流错误: {state['callback_task_id']}, 错误: {error_message}")

        # 标记任务失败
        state["task_info"].fail_processing(error_message)

        # 更新进度管理器
        if state.get("progress_manager"):
            await state["progress_manager"].complete_task(
                state["callback_task_id"],
                state["user_id"],
                {
                    "overall_task_status": "failed",
                    "error": error_message,
                    "status": "failed"
                }
            )

        return {
            "current_stage": "error_handler",
            "overall_task_status": "failed",
            "messages": state["messages"] + [{
                "role": "system",
                "content": f"工作流错误: {error_message}"
            }]
        }

    async def _terminate_node(self, state: OutlineGenerationState) -> dict:
        """
        终止节点

        Args:
            state: 当前状态

        Returns:
            dict: 终止状态
        """
        logger.warning(f"大纲生成工作流终止: {state['callback_task_id']}")

        # 标记任务取消
        state["task_info"].cancel_processing()

        # 更新进度管理器
        if state.get("progress_manager"):
            await state["progress_manager"].complete_task(
                state["callback_task_id"],
                state["user_id"],
                {
                    "overall_task_status": "terminated",
                    "message": "任务已被用户终止"
                }
            )

        return {
            "current_stage": "terminated",
            "overall_task_status": "terminated",
            "messages": state["messages"] + [{
                "role": "system",
                "content": "工作流已终止"
            }]
        }

    def get_graph(self) -> StateGraph:
        """
        获取编译后的工作流图

        如果图尚未编译，则自动编译

        Returns:
            StateGraph: 编译后的工作流图
        """
        if self.graph is None:
            self.graph = self.build_graph()
        return self.graph



def should_recommend_similar(state: OutlineGenerationState) -> str:
    """
    是否应该推荐相似案例

    用于条件边判断
    """
    similarity_config = state.get("similarity_config")
    if similarity_config and similarity_config.get("max_cases", 0) > 0:
        return "yes"
    return "no"


def should_link_knowledge(state: OutlineGenerationState) -> str:
    """
    是否应该关联知识点

    用于条件边判断
    """
    knowledge_config = state.get("knowledge_config", {})
    if knowledge_config.get("enable_basis", True):
        return "yes"
    return "no"


def should_terminate_or_error(state: OutlineGenerationState) -> str:
    """
    是否终止或出错

    用于条件边判断
    """
    if state.get("overall_task_status") == "terminated":
        return "terminate"
    if state.get("overall_task_status") == "failed" or state.get("error_message"):
        return "error"
    return "continue"
