import gradio as gr
import pandas as pd
from agents.brain import AutoWarehouseAgent

agent = AutoWarehouseAgent()

# -------------------------------------------------------------
# Convert UI DataFrame input to Pandas
# -------------------------------------------------------------
def df_from_input(df):
    return pd.DataFrame(df) if df else pd.DataFrame()


# -------------------------------------------------------------
# MAIN AGENT RUNNER
# -------------------------------------------------------------
def run_agent(message, slotting_df, picking_df):
    print("📥 INPUT MESSAGE:", message)
    print("📦 SLOT DATA RECEIVED:", slotting_df)
    print("🚚 PICKING DATA RECEIVED:", picking_df)

    try:
        raw = agent.run(message, slotting_df, picking_df)
        print("🔍 RAW RESULT FROM AGENT:", raw)

        # -----------------------------------------------------
        # NORMALIZE OUTPUT DICTIONARY
        # -----------------------------------------------------
        if isinstance(raw, str):
            # If only text returned → wrap it
            result = {
                "report": raw,
                "route_image": None,
                "slotting_table": pd.DataFrame(),
            }

        elif isinstance(raw, dict):
            # Normalize missing fields
            result = {
                "report": raw.get("report", "No report generated."),
                "route_image": raw.get("route_image", None),
                "slotting_table": raw.get("slotting_table", pd.DataFrame()),
            }
        else:
            # Unexpected type
            result = {
                "report": "⚠️ Invalid agent output format.",
                "route_image": None,
                "slotting_table": pd.DataFrame(),
            }

        # UI expects EXACTLY 3 outputs
        return (
            result["report"],
            result["route_image"],
            result["slotting_table"],
        )

    except Exception as e:
        import traceback
        print("❌ ERROR OCCURRED:")
        traceback.print_exc()

        return (
            f"### ❌ Error\n{str(e)}",
            None,
            pd.DataFrame(),
        )


# -------------------------------------------------------------
# BUILD UI
# -------------------------------------------------------------
def build_ui():
    with gr.Blocks() as demo:

        # Title
        gr.Markdown(
            """
            <h1 style='color:#FF6A00;text-align:center;'>
                Procelevate Autonomous Warehouse Operator (v2.0)
            </h1>
            <p style='text-align:center;font-size:16px;'>
                Slotting • Picking • Forecasting • Replenishment • Rebalancing • Workforce • Dock Scheduling
            </p>
            """
        )

        # User input
        query = gr.Textbox(
            label="Enter your warehouse request",
            placeholder="Examples: 'Rebalance inventory', 'Forecast next 7 days', 'Optimize slotting', 'How many workers needed?'"
        )

        # Slotting input table
        gr.Markdown("### 📦 Slotting Data (SKU Master)")
        slot_df = gr.Dataframe(
            headers=["SKU", "Velocity", "Frequency"],
            value=[
                ["A123", "Fast", 120],
                ["B555", "Medium", 60],
                ["C888", "Slow", 8],
            ],
        )

        # Picking input table
        gr.Markdown("### 🚚 Picking Data")
        pick_df = gr.Dataframe(
            headers=["Aisle", "Rack"],
            value=[[5, 14], [3, 10], [12, 7]],
        )

        # Run button
        btn = gr.Button("Run Warehouse Agent", variant="primary")

        # Outputs
        out_report = gr.Markdown(label="📄 Report")
        out_route = gr.Image(label="📊 Charts / Diagrams",type="pil")
        out_slot = gr.Dataframe(label="📘 Output Table")

        # Execution wiring
        btn.click(
            run_agent,
            inputs=[query, slot_df, pick_df],
            outputs=[out_report, out_route, out_slot],
        )

    return demo


# -------------------------------------------------------------
# LAUNCH APP
# -------------------------------------------------------------
demo = build_ui()
demo.launch()