agents/brain.py

from agents.planner import detect_intent
from agents.reasoner import (
    run_slotting_analysis,
    run_picking_optimization,
    run_demand_forecast,
    run_replenishment_analysis,
    run_rebalancing_analysis,
    run_workforce_optimization,
    run_dock_scheduling
)
import pandas as pd


class AutoWarehouseAgent:

    def run(self, message, slotting_df, picking_df):
        """
        Central warehouse AI engine.
        Returns a structured dictionary:
        {
            "report": markdown text,
            "route_image": image or None,
            "slotting_table": dataframe or empty df
        }
        """

        task = detect_intent(message.lower().strip())

        # Debug logs
        print(f"🧠 DETECTED TASK: {task}")
        print("📦 Slotting DF:", slotting_df)
        print("🚚 Picking DF:", picking_df)

        # ------------------------------------------------------
        # 1️⃣ SLOTING OPTIMIZATION
        # ------------------------------------------------------
        if task == "slotting":
            explanation, slot_plan = run_slotting_analysis(message, slotting_df)
            return {
                "report": explanation,
                "route_image": None,
                "slotting_table": slot_plan
            }

        # ------------------------------------------------------
        # 2️⃣ PICKING ROUTE OPTIMIZATION
        # ------------------------------------------------------
        if task == "picking":
            explanation, route_img = run_picking_optimization(message, picking_df)
            return {
                "report": explanation,
                "route_image": route_img,
                "slotting_table": pd.DataFrame()
            }

        # ------------------------------------------------------
        # 3️⃣ FULL REPORT — Slotting + Picking
        # ------------------------------------------------------
        if task == "report":
            exp1, slot_plan = run_slotting_analysis(message, slotting_df)
            exp2, route_img = run_picking_optimization(message, picking_df)

            combined_report = (
                "### 🧾 Full Warehouse Intelligence Report\n\n"
                "#### 📦 Slotting Optimization\n" + exp1 +
                "\n\n---\n\n"
                "#### 🚚 Picking Optimization\n" + exp2
            )

            return {
                "report": combined_report,
                "route_image": route_img,
                "slotting_table": slot_plan
            }

        # ------------------------------------------------------
        # 4️⃣ DEMAND FORECASTING
        # ------------------------------------------------------
        if task == "forecast":
            explanation, forecast_plot, forecast_table = run_demand_forecast(message, slotting_df)
            return {
                "report": explanation,
                "route_image": forecast_plot,
                "slotting_table": forecast_table
            }

        # ------------------------------------------------------
        # 5️⃣ REPLENISHMENT ANALYSIS
        # ------------------------------------------------------
        if task == "replenishment":
            explanation, repl_table = run_replenishment_analysis(message, slotting_df)
            return {
                "report": explanation,
                "route_image": None,
                "slotting_table": repl_table
            }

        # ------------------------------------------------------
        # 6️⃣ INVENTORY REBALANCING
        # ------------------------------------------------------
        if task == "rebalancing":
            explanation, move_table = run_rebalancing_analysis(message, slotting_df)
            return {
                "report": explanation,
                "route_image": None,
                "slotting_table": move_table
            }

        # ------------------------------------------------------
        # 7️⃣ WORKFORCE OPTIMIZATION
        # ------------------------------------------------------
        if task == "workforce":
            explanation, workforce_table = run_workforce_optimization(message, slotting_df)
            return {
                "report": explanation,
                "route_image": None,
                "slotting_table": workforce_table
            }

        # ------------------------------------------------------
        # 8️⃣ DOCK SCHEDULING OPTIMIZATION
        # ------------------------------------------------------
        if task == "dock":
            explanation, dock_table = run_dock_scheduling(message, slotting_df)
            return {
                "report": explanation,
                "route_image": None,
                "slotting_table": dock_table
            }

        # ------------------------------------------------------
        # ❌ FALLBACK HANDLER
        # ------------------------------------------------------
        return {
            "report": (
                "### ❓ Unable to Understand Your Request\n"
                "Please try queries related to:\n"
                "- Slotting optimization\n"
                "- Picking optimization\n"
                "- Forecasting\n"
                "- Replenishment\n"
                "- Inventory rebalancing\n"
                "- Workforce planning\n"
                "- Dock scheduling\n"
                "- Full warehouse report\n"
            ),
            "route_image": None,
            "slotting_table": pd.DataFrame()
        }