cas / fw-pms-ai

"""
补货建议 Agent

使用 part_ratio + SQL Agent + LangGraph
"""

import logging
import time
import uuid
from typing import Optional, List
from datetime import date, datetime
from decimal import Decimal

from langgraph.graph import StateGraph, END

from .state import AgentState
from .nodes import (
    fetch_part_ratio_node,
    sql_agent_node,
    allocate_budget_node,
    should_retry_sql,
)
from ..models import ReplenishmentTask, TaskStatus, ReplenishmentPartSummary
from ..services import ResultWriter

logger = logging.getLogger(__name__)


class ReplenishmentAgent:
    """补货建议 Agent"""

    def __init__(self):
        self._graph = None
        self._result_writer = ResultWriter()

    @property
    def graph(self) -> StateGraph:
        """获取工作流图"""
        if self._graph is None:
            self._graph = self._build_graph()
        return self._graph

    def _build_graph(self) -> StateGraph:
        """
        构建 LangGraph 工作流

        工作流结构：
        fetch_part_ratio → sql_agent → allocate_budget → END
        """
        workflow = StateGraph(AgentState)

        workflow.add_node("fetch_part_ratio", fetch_part_ratio_node)
        workflow.add_node("sql_agent", sql_agent_node)
        workflow.add_node("allocate_budget", allocate_budget_node)

        workflow.set_entry_point("fetch_part_ratio")
        workflow.add_edge("fetch_part_ratio", "sql_agent")

        workflow.add_conditional_edges(
            "sql_agent",
            should_retry_sql,
            {
                "retry": "sql_agent",
                "continue": "allocate_budget",
            }
        )

        workflow.add_edge("allocate_budget", END)

        return workflow.compile()

    def run(
        self,
        group_id: int,
        dealer_grouping_id: int,
        dealer_grouping_name: str,
        brand_grouping_id: Optional[int] = None,
        brand_grouping_name: str = "",
        statistics_date: Optional[str] = None,
    ) -> AgentState:
        """
        执行补货建议生成
        
        Args:
            group_id: 集团ID
            dealer_grouping_id: 商家组合ID
            dealer_grouping_name: 商家组合名称
            brand_grouping_id: 品牌组合ID
            brand_grouping_name: 品牌组合名称
            statistics_date: 统计日期
        """
        task_no = f"AI-{uuid.uuid4().hex[:12].upper()}"
        if statistics_date is None:
            statistics_date = date.today().strftime("%Y-%m-%d")

        logger.info(
            f"开始执行补货建议: task_no={task_no}, "
            f"dealer_grouping={dealer_grouping_name}"
        )

        # 初始化状态
        initial_state: AgentState = {
            "task_no": task_no,
            "group_id": group_id,
            "brand_grouping_id": brand_grouping_id,
            "brand_grouping_name": brand_grouping_name,
            "dealer_grouping_id": dealer_grouping_id,
            "dealer_grouping_name": dealer_grouping_name,
            "statistics_date": statistics_date,
            "part_ratios": [],
            "sql_queries": [],
            "sql_results": [],
            "sql_retry_count": 0,
            "sql_execution_logs": [],
            "base_ratio": Decimal("1.1"),
            "allocated_details": [],
            "details": [],
            "llm_suggestions": [],
            "part_results": [],
            "llm_provider": "",
            "llm_model": "",
            "llm_prompt_tokens": 0,
            "llm_completion_tokens": 0,

            "status": "running",
            "error_message": "",
            "start_time": time.time(),
            "end_time": 0,
            "current_node": "",
            "next_node": "fetch_part_ratio",
        }

        # 创建任务记录
        task = ReplenishmentTask(
            task_no=task_no,
            group_id=group_id,
            dealer_grouping_id=dealer_grouping_id,
            dealer_grouping_name=dealer_grouping_name,
            brand_grouping_id=brand_grouping_id,
            statistics_date=statistics_date,
            status=TaskStatus.RUNNING,
        )
        self._result_writer.save_task(task)

        try:
            # 执行工作流
            final_state = self.graph.invoke(initial_state)

            # 更新任务状态
            execution_time = int((final_state.get("end_time", time.time()) - final_state["start_time"]) * 1000)
            actual_amount = sum(d.suggest_amount for d in final_state.get("details", []))

            task.status = TaskStatus.SUCCESS
            task.actual_amount = actual_amount
            task.part_count = len(final_state.get("details", []))
            task.shop_count = len(set(d.shop_id for d in final_state.get("details", [])))
            task.base_ratio = final_state.get("base_ratio", Decimal("0"))
            task.llm_provider = final_state.get("llm_provider", "")
            task.llm_model = final_state.get("llm_model", "")
            task.llm_prompt_tokens = final_state.get("llm_prompt_tokens", 0)
            task.llm_completion_tokens = final_state.get("llm_completion_tokens", 0)
            task.llm_total_tokens = task.llm_prompt_tokens + task.llm_completion_tokens
            task.llm_analysis_summary = final_state.get("llm_analysis_summary", "")
            task.execution_time_ms = execution_time

            self._result_writer.update_task(task)

            logger.info(
                f"补货建议执行完成: task_no={task_no}, "
                f"parts={task.part_count}, amount={actual_amount}, "
                f"time={execution_time}ms"
            )

            return final_state

        except Exception as e:
            logger.error(f"补货建议执行失败: task_no={task_no}, error={e}")

            task.status = TaskStatus.FAILED
            task.error_message = str(e)
            task.execution_time_ms = int((time.time() - initial_state["start_time"]) * 1000)
            self._result_writer.update_task(task)

            raise

        finally:
            self._result_writer.close()


    def _save_part_summaries(
        self,
        task_no: str,
        group_id: int,
        dealer_grouping_id: int,
        statistics_date: str,
        part_results: list,
    ):
        """保存配件汇总"""
        from .sql_agent import PartAnalysisResult
        
        summaries = []
        for pr in part_results:
            if not isinstance(pr, PartAnalysisResult):
                continue
            summary = ReplenishmentPartSummary(
                task_no=task_no,
                group_id=group_id,
                dealer_grouping_id=dealer_grouping_id,
                part_code=pr.part_code,
                part_name=pr.part_name,
                unit=pr.unit,
                cost_price=pr.cost_price,
                total_storage_cnt=pr.total_storage_cnt,
                total_avg_sales_cnt=pr.total_avg_sales_cnt,
                group_current_ratio=pr.group_current_ratio,
                total_suggest_cnt=pr.total_suggest_cnt,
                total_suggest_amount=pr.total_suggest_amount,
                shop_count=pr.shop_count,
                need_replenishment_shop_count=pr.need_replenishment_shop_count,
                part_decision_reason=pr.part_decision_reason,
                priority=pr.priority,
                llm_confidence=pr.confidence,
                statistics_date=statistics_date,
            )
            summaries.append(summary)
        
        if summaries:
            self._result_writer.save_part_summaries(summaries)
            logger.info(f"保存配件汇总: count={len(summaries)}")

    def run_single_part(
        self,
        group_id: int,
        dealer_grouping_id: int,
        part_code: str,
        statistics_date: Optional[str] = None,
    ) -> AgentState:
        """
        执行单种配件补货建议生成

        Args:
            group_id: 集团ID
            dealer_grouping_id: 商家组合ID
            part_code: 配件编码
            statistics_date: 统计日期
        """
        task_no = f"AIS-{uuid.uuid4().hex[:12].upper()}"
        if statistics_date is None:
            statistics_date = date.today().strftime("%Y-%m-%d")

        # 查询商家组合名称
        from ..services import DataService
        data_service = DataService()
        try:
            groupings = data_service.get_dealer_groupings(group_id)
            grouping = next(
                (g for g in groupings if g["id"] == dealer_grouping_id), None
            )
            dealer_grouping_name = grouping["name"] if grouping else f"商家组合{dealer_grouping_id}"
        finally:
            data_service.close()

        logger.info(
            f"开始执行单配件补货建议: task_no={task_no}, "
            f"part_code={part_code}, dealer_grouping={dealer_grouping_name}"
        )

        # 覆盖已有的该配件补货建议
        try:
            self._result_writer.delete_details_by_part(dealer_grouping_id, part_code)
            self._result_writer.delete_part_summaries_by_part(dealer_grouping_id, part_code)
            logger.info(f"已覆盖旧补货建议: dealer_grouping_id={dealer_grouping_id}, part_code={part_code}")
        except Exception as e:
            logger.warning(f"覆盖旧数据时出错（继续执行）: {e}")

        # 获取单配件 part_ratio 数据
        from .sql_agent import SQLAgent
        sql_agent = SQLAgent()
        try:
            part_ratios = sql_agent.fetch_part_ratios_by_part_code(
                group_id=group_id,
                dealer_grouping_id=dealer_grouping_id,
                statistics_date=statistics_date,
                part_code=part_code,
            )
        finally:
            sql_agent.close()

        if not part_ratios:
            logger.warning(f"未找到配件数据: part_code={part_code}, dealer_grouping_id={dealer_grouping_id}")
            raise ValueError(f"未找到配件数据: part_code={part_code}")

        # 初始化状态（预填充 part_ratios，跳过 fetch_part_ratio 节点的重新查询）
        initial_state: AgentState = {
            "task_no": task_no,
            "group_id": group_id,
            "brand_grouping_id": None,
            "brand_grouping_name": "",
            "dealer_grouping_id": dealer_grouping_id,
            "dealer_grouping_name": dealer_grouping_name,
            "statistics_date": statistics_date,
            "part_ratios": part_ratios,
            "sql_queries": [],
            "sql_results": [],
            "sql_retry_count": 0,
            "sql_execution_logs": [],
            "base_ratio": Decimal("1.1"),
            "allocated_details": [],
            "details": [],
            "llm_suggestions": [],
            "part_results": [],
            "llm_provider": "",
            "llm_model": "",
            "llm_prompt_tokens": 0,
            "llm_completion_tokens": 0,
            "status": "running",
            "error_message": "",
            "start_time": time.time(),
            "end_time": 0,
            "current_node": "",
            "next_node": "sql_agent",
        }

        # 创建任务记录
        task = ReplenishmentTask(
            task_no=task_no,
            group_id=group_id,
            dealer_grouping_id=dealer_grouping_id,
            dealer_grouping_name=dealer_grouping_name,
            statistics_date=statistics_date,
            status=TaskStatus.RUNNING,
        )
        self._result_writer.save_task(task)

        try:
            # 构建简化工作流（跳过 fetch_part_ratio，直接从 sql_agent 开始）
            workflow = StateGraph(AgentState)
            workflow.add_node("sql_agent", sql_agent_node)
            workflow.add_node("allocate_budget", allocate_budget_node)
            workflow.set_entry_point("sql_agent")
            workflow.add_conditional_edges(
                "sql_agent",
                should_retry_sql,
                {"retry": "sql_agent", "continue": "allocate_budget"},
            )
            workflow.add_edge("allocate_budget", END)
            graph = workflow.compile()

            final_state = graph.invoke(initial_state)

            # 更新任务状态
            execution_time = int(
                (final_state.get("end_time", time.time()) - final_state["start_time"]) * 1000
            )
            actual_amount = sum(d.suggest_amount for d in final_state.get("details", []))

            task.status = TaskStatus.SUCCESS
            task.actual_amount = actual_amount
            task.part_count = len(final_state.get("details", []))
            task.shop_count = len(set(d.shop_id for d in final_state.get("details", [])))
            task.base_ratio = final_state.get("base_ratio", Decimal("0"))
            task.llm_provider = final_state.get("llm_provider", "")
            task.llm_model = final_state.get("llm_model", "")
            task.llm_prompt_tokens = final_state.get("llm_prompt_tokens", 0)
            task.llm_completion_tokens = final_state.get("llm_completion_tokens", 0)
            task.llm_total_tokens = task.llm_prompt_tokens + task.llm_completion_tokens
            task.execution_time_ms = execution_time

            self._result_writer.update_task(task)

            logger.info(
                f"单配件补货建议完成: task_no={task_no}, part_code={part_code}, "
                f"details={task.part_count}, amount={actual_amount}, time={execution_time}ms"
            )

            return final_state

        except Exception as e:
            logger.error(f"单配件补货建议失败: task_no={task_no}, error={e}")

            task.status = TaskStatus.FAILED
            task.error_message = str(e)
            task.execution_time_ms = int(
                (time.time() - initial_state["start_time"]) * 1000
            )
            self._result_writer.update_task(task)

            raise

    def run_for_all_groupings(self, group_id: int):
        """
        为所有商家组合执行补货建议
        """
        from ..services import DataService

        data_service = DataService()
        try:
            groupings = data_service.get_dealer_groupings(group_id)
            logger.info(f"获取商家组合: group_id={group_id}, count={len(groupings)}")

            for idx, grouping in enumerate(groupings):
                logger.info(f"[{idx+1}/{len(groupings)}] 开始处理商家组合: {grouping['name']} (id={grouping['id']})")
                try:
                    self.run(
                        group_id=group_id,
                        dealer_grouping_id=grouping["id"],
                        dealer_grouping_name=grouping["name"],
                    )
                    logger.info(f"[{grouping['name']}] 执行完成")

                except Exception as e:
                    logger.error(f"商家组合执行失败: {grouping['name']}, error={e}", exc_info=True)
                    continue

        finally:
            data_service.close()