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

agent = AutoWarehouseAgent()


def df_from_input(df):
    return pd.DataFrame(df) if df else pd.DataFrame()


def run_agent(message, slotting_df, picking_df):
    print("📥 INPUT MESSAGE:", message)
    print("📄 SLOT DATA:", slotting_df)
    print("🚚 PICKING DATA:", picking_df)

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

        # -----------------------------------------------------
        # 📌 STANDARDIZE OUTPUT FORMAT
        # -----------------------------------------------------
        if isinstance(raw, str):
            # Agent returned text only → convert to dictionary
            result = {
                "report": raw,
                "route_image": None,
                "slotting_table": pd.DataFrame(),
            }

        elif isinstance(raw, dict):
            # Fill missing pieces
            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 return type
            result = {
                "report": "⚠️ Invalid agent output format.",
                "route_image": None,
                "slotting_table": pd.DataFrame(),
            }

        # Return EXACTLY 3 items to match UI 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(),
        )


def build_ui():
    with gr.Blocks() as demo:

        gr.Markdown(
            "<h1 style='color:#FF6A00;'>Procelevate Autonomous Warehouse Operator</h1>"
        )

        query = gr.Textbox(label="Enter your request")

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

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

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

        out_report = gr.Markdown()
        out_route = gr.Image()
        out_slot = gr.Dataframe()

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

    return demo


demo = build_ui()
demo.launch()