agents/brain.py

from agents.planner import detect_intent
from agents.reasoner import (
    run_slotting_analysis,
    run_picking_optimization
)
import pandas as pd


class AutoWarehouseAgent:

    def run(self, message, slotting_df, picking_df):
        """
        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 prints
        print(f"🧠 DETECTED TASK: {task}")
        print("📦 Slotting DF:", slotting_df)
        print("🚚 Picking DF:", picking_df)

        # ------------------------------------------------------
        # 1️⃣ SLOTTING
        # ------------------------------------------------------
        if task == "slotting":
            explanation, slot_plan = run_slotting_analysis(message, slotting_df)

            return {
                "report": f"### 📦 Slotting Optimization Report\n\n{explanation}",
                "route_image": None,
                "slotting_table": slot_plan
            }

        # ------------------------------------------------------
        # 2️⃣ PICKING
        # ------------------------------------------------------
        if task == "picking":
            explanation, route_img = run_picking_optimization(message, picking_df)

            return {
                "report": f"### 🚚 Picking Route Optimization\n\n{explanation}",
                "route_image": route_img,
                "slotting_table": pd.DataFrame()
            }

        # ------------------------------------------------------
        # 3️⃣ FULL REPORT (both 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
            }

        # ------------------------------------------------------
        # ❌ FALLBACK
        # ------------------------------------------------------
        return {
            "report": (
                "### ❓ Unable to Understand Request\n"
                "I couldn't determine whether your request is about:\n"
                "- Slotting optimization\n"
                "- Picking route planning\n"
                "- Full warehouse report\n\n"
                "Try rewriting your query."
            ),
            "route_image": None,
            "slotting_table": pd.DataFrame()
        }