Blame for gradio_app.py

"""Interactive Gradio frontend for the Store OpenEnv simulation."""

from __future__ import annotations

from typing import Any

try:
    import gradio as gr
except ImportError as exc:  # pragma: no cover
    raise ImportError(
        "gradio is required for the UI. Install project dependencies with `uv sync`."
    ) from exc

try:
    from .models import OrderRequest, PriceUpdate, StoreAction, StoreObservation, dump_model
    from .server.store_environment import StoreEnvironment
except ImportError:  # pragma: no cover
    from models import OrderRequest, PriceUpdate, StoreAction, StoreObservation, dump_model
    from server.store_environment import StoreEnvironment

SELLERS = ["Seller A", "Seller B", "Seller C", "Seller D"]

CSS = """
/* ── Global ── */
.gradio-container {
    background: linear-gradient(135deg, #0f172a 0%, #1e293b 50%, #0f172a 100%) !important;
    font-family: 'Inter', system-ui, -apple-system, sans-serif !important;
}

/* ── Header ── */
.app-header {
    text-align: center;
    padding: 28px 16px 12px;
}
.app-header h1 {
    font-size: 2rem;
    font-weight: 700;
    background: linear-gradient(135deg, #60a5fa, #a78bfa, #f472b6);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin: 0;
}
.app-header p {
    color: #94a3b8;
    font-size: 0.95rem;
    margin-top: 4px;
}

/* ── Panels ── */
.panel-card {
    background: rgba(30, 41, 59, 0.85) !important;
    border: 1px solid rgba(148, 163, 184, 0.12) !important;
    border-radius: 16px !important;
    padding: 20px !important;
    backdrop-filter: blur(12px);
}

/* ── Section headings ── */
.section-title {
    color: #e2e8f0 !important;
    font-size: 1.05rem !important;
    font-weight: 600 !important;
    border-bottom: 1px solid rgba(148, 163, 184, 0.15);
    padding-bottom: 8px;
    margin-bottom: 12px !important;
}

/* ── Metric cards ── */
.metric-row {
    display: flex;
    gap: 10px;
    flex-wrap: wrap;
    margin-bottom: 8px;
}
.metric-card {
    flex: 1;
    min-width: 120px;
    background: rgba(51, 65, 85, 0.6);
    border: 1px solid rgba(148, 163, 184, 0.1);
    border-radius: 12px;
    padding: 14px 16px;
    text-align: center;
}
.metric-card .label {
    font-size: 0.72rem;
    text-transform: uppercase;
    letter-spacing: 0.6px;
    color: #94a3b8;
    margin-bottom: 4px;
}
.metric-card .value {
    font-size: 1.25rem;
    font-weight: 700;
    color: #f1f5f9;
}
.metric-card .value.positive { color: #4ade80; }
.metric-card .value.negative { color: #f87171; }

/* ── Status banner ── */
.status-banner {
    padding: 12px 16px;
    border-radius: 10px;
    font-size: 0.92rem;
    margin-bottom: 10px;
}
.status-banner.success {
    background: rgba(34, 197, 94, 0.12);
    border: 1px solid rgba(34, 197, 94, 0.25);
    color: #86efac;
}
.status-banner.error {
    background: rgba(239, 68, 68, 0.12);
    border: 1px solid rgba(239, 68, 68, 0.25);
    color: #fca5a5;
}
.status-banner.info {
    background: rgba(59, 130, 246, 0.12);
    border: 1px solid rgba(59, 130, 246, 0.25);
    color: #93c5fd;
}

/* ── Buttons ── */
.btn-primary {
    background: linear-gradient(135deg, #3b82f6, #6366f1) !important;
    border: none !important;
    color: #fff !important;
    font-weight: 600 !important;
    border-radius: 10px !important;
    padding: 10px 20px !important;
    transition: opacity 0.2s !important;
}
.btn-primary:hover { opacity: 0.88 !important; }

.btn-secondary {
    background: rgba(51, 65, 85, 0.7) !important;
    border: 1px solid rgba(148, 163, 184, 0.2) !important;
    color: #cbd5e1 !important;
    font-weight: 500 !important;
    border-radius: 10px !important;
    padding: 10px 20px !important;
}
.btn-secondary:hover { background: rgba(71, 85, 105, 0.7) !important; }

.btn-danger {
    background: rgba(239, 68, 68, 0.15) !important;
    border: 1px solid rgba(239, 68, 68, 0.3) !important;
    color: #fca5a5 !important;
    font-weight: 500 !important;
    border-radius: 10px !important;
}

/* ── Order rows ── */
.order-row {
    background: rgba(51, 65, 85, 0.4);
    border: 1px solid rgba(148, 163, 184, 0.1);
    border-radius: 10px;
    padding: 10px 14px;
    margin-bottom: 6px;
}

/* ── Fix Gradio element readability ── */
.gradio-container label { color: #cbd5e1 !important; }
.gradio-container .gr-input, .gradio-container input,
.gradio-container select, .gradio-container textarea {
    background: rgba(30, 41, 59, 0.9) !important;
    color: #f1f5f9 !important;
    border: 1px solid rgba(148, 163, 184, 0.2) !important;
    border-radius: 8px !important;
}
.gradio-container .gr-input:focus, .gradio-container input:focus {
    border-color: #3b82f6 !important;
    box-shadow: 0 0 0 2px rgba(59, 130, 246, 0.15) !important;
}
.gradio-container .gr-check-radio { accent-color: #6366f1 !important; }
.gradio-container .tabs .tab-nav button {
    color: #94a3b8 !important;
    border: none !important;
    font-weight: 500 !important;
}
.gradio-container .tabs .tab-nav button.selected {
    color: #60a5fa !important;
    border-bottom: 2px solid #60a5fa !important;
}
.gradio-container .json-holder { background: rgba(15, 23, 42, 0.6) !important; }
.gradio-container .markdown-text { color: #e2e8f0 !important; }
.gradio-container .markdown-text h1,
.gradio-container .markdown-text h2,
.gradio-container .markdown-text h3 { color: #f1f5f9 !important; }
.gradio-container .markdown-text code {
    background: rgba(99, 102, 241, 0.15) !important;
    color: #a5b4fc !important;
    padding: 2px 6px;
    border-radius: 4px;
}
"""


def _load_products_and_stock() -> tuple[list[str], list[str]]:
    """Return (market product names, stock product names) from a fresh env."""
    env = StoreEnvironment()
    obs = env.reset()
    market_products = [m.product_name for m in obs.market]
    stock_products = [s.product_name for s in obs.stock]
    return market_products, stock_products


def _load_marketing_catalog() -> list[dict[str, Any]]:
    env = StoreEnvironment()
    obs = env.reset()
    return [dump_model(m, exclude_none=True) for m in obs.marketing_actions]


def _format_marketing_label(m: dict[str, Any]) -> str:
    parts = [f"#{m['action_id']} {m['description']}"]
    parts.append(f"cost ₹{m['cost_inr']:,}")
    if m.get("recurring_cost_inr"):
        parts.append(f"recurring ₹{m['recurring_cost_inr']:,}/day")
    parts.append(f"satisfaction {m['customer_satisfaction_delta']:+d}")
    if m.get("order_processing_capacity_delta"):
        parts.append(f"capacity {m['order_processing_capacity_delta']:+d}")
    return " · ".join(parts)


def _render_status_html(message: str, kind: str = "info") -> str:
    return f'<div class="status-banner {kind}">{message}</div>'


def _render_metrics_html(obs: StoreObservation) -> str:
    store = obs.store

    def _card(label: str, value: str, css_class: str = "") -> str:
        vc = f' class="value {css_class}"' if css_class else ' class="value"'
        return (
            f'<div class="metric-card">'
            f'<div class="label">{label}</div>'
            f"<div{vc}>{value}</div>"
            f"</div>"
        )

    profit = store.cumulative_profit_inr
    profit_cls = "positive" if profit >= 0 else "negative"
    reward_val = f"{obs.reward:.4f}" if obs.reward else "—"
    recurring_cost = store.active_recurring_cost_inr
    total_daily_burn = store.current_daily_expense_inr + recurring_cost
    burn_cls = "negative" if total_daily_burn > 0 else ""
    recurring_cls = "negative" if recurring_cost > 0 else ""

    cards = [
        _card("Day", str(store.day)),
        _card("Bank Balance", f"₹{store.bank_balance_inr:,.0f}"),
        _card("Base Burn / Day", f"₹{store.current_daily_expense_inr:,.0f}", burn_cls),
        _card("Recurring / Day", f"₹{recurring_cost:,.0f}", recurring_cls),
        _card("Total Burn / Day", f"₹{total_daily_burn:,.0f}", burn_cls),
        _card("Satisfaction", f"{store.customer_satisfaction:.0f}%"),
        _card("Profit", f"₹{profit:,.0f}", profit_cls),
        _card("Reward", reward_val),
        _card("Capacity", f"{store.order_processing_capacity:.0f}"),
    ]
    return f'<div class="metric-row">{"".join(cards)}</div>'


THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.blue,
    secondary_hue=gr.themes.colors.indigo,
    neutral_hue=gr.themes.colors.slate,
    font=gr.themes.GoogleFont("Inter"),
).set(
    body_background_fill="transparent",
    block_background_fill="rgba(30, 41, 59, 0.7)",
    block_border_width="1px",
    block_border_color="rgba(148, 163, 184, 0.12)",
    block_label_text_color="#cbd5e1",
    block_title_text_color="#e2e8f0",
    input_background_fill="rgba(15, 23, 42, 0.6)",
    input_border_color="rgba(148, 163, 184, 0.2)",
    button_primary_background_fill="linear-gradient(135deg, #3b82f6, #6366f1)",
    button_primary_text_color="#ffffff",
    button_secondary_background_fill="rgba(51, 65, 85, 0.7)",
    button_secondary_text_color="#cbd5e1",
    button_secondary_border_color="rgba(148, 163, 184, 0.2)",
)


def make_demo() -> gr.Blocks:
    """Build the interactive Gradio app."""

    market_products, stock_products = _load_products_and_stock()
    marketing_catalog = _load_marketing_catalog()
    marketing_labels = [_format_marketing_label(m) for m in marketing_catalog]
    marketing_id_by_label = {
        _format_marketing_label(m): m["action_id"] for m in marketing_catalog
    }

    with gr.Blocks(title="Store OpenEnv Playground") as demo:
        env_state = gr.State(value=None)
        obs_state = gr.State(value=None)

        # ── Header ──
        gr.HTML(
            '<div class="app-header">'
            "<h1>Store OpenEnv Playground</h1>"
            "<p>Simulate a neighborhood kirana store — place orders, set prices, activate marketing</p>"
            "</div>"
        )

        # ── Top metrics ──
        metrics_html = gr.HTML("")
        status_html = gr.HTML(_render_status_html("Press Reset to start a fresh episode.", "info"))

        with gr.Row(equal_height=False):
            # ═══════════ LEFT: ACTION PANEL ═══════════
            with gr.Column(scale=4, elem_classes=["panel-card"]):
                gr.Markdown("### Place Orders", elem_classes=["section-title"])
                gr.Markdown(
                    "*Select a product from the market catalog, choose a seller, and specify quantity.*",
                )

                with gr.Group():
                    with gr.Row():
                        order_product = gr.Dropdown(
                            choices=market_products,
                            label="Product",
                            value=market_products[0] if market_products else None,
                            scale=3,
                        )
                        order_seller = gr.Dropdown(
                            choices=SELLERS,
                            label="Seller",
                            value=SELLERS[0],
                            scale=2,
                        )
                        order_qty = gr.Number(label="Qty", value=10, minimum=1, precision=0, scale=1)
                        add_order_btn = gr.Button("+ Add", size="sm", elem_classes=["btn-primary"], scale=1)

                orders_display = gr.Dataframe(
                    headers=["Product", "Seller", "Quantity"],
                    datatype=["str", "str", "number"],
                    value=[],
                    interactive=False,
                    label="Pending Orders",
                    row_count=(0, "dynamic"),
                    column_count=(3, "fixed"),
                )
                clear_orders_btn = gr.Button("Clear Orders", size="sm", elem_classes=["btn-danger"])

                gr.Markdown("---")
                gr.Markdown("### Update Prices", elem_classes=["section-title"])
                gr.Markdown("*Adjust selling prices for products currently in stock.*")

                with gr.Group():
                    with gr.Row():
                        price_product = gr.Dropdown(
                            choices=stock_products,
                            label="Product",
                            value=stock_products[0] if stock_products else None,
                            scale=3,
                        )
                        price_value = gr.Number(label="New Price ₹", value=100, minimum=1, precision=0, scale=2)
                        add_price_btn = gr.Button("+ Add", size="sm", elem_classes=["btn-primary"], scale=1)

                prices_display = gr.Dataframe(
                    headers=["Product", "Price (₹)"],
                    datatype=["str", "number"],
                    value=[],
                    interactive=False,
                    label="Price Changes",
                    row_count=(0, "dynamic"),
                    column_count=(2, "fixed"),
                )
                clear_prices_btn = gr.Button("Clear Prices", size="sm", elem_classes=["btn-danger"])

                gr.Markdown("---")
                gr.Markdown("### Marketing Actions", elem_classes=["section-title"])
                gr.Markdown("*Select marketing or operational upgrades to activate.*")

                marketing_checkboxes = gr.CheckboxGroup(
                    choices=marketing_labels,
                    label="Available Actions",
                    value=[],
                )

                gr.Markdown("---")
                with gr.Row():
                    reset_btn = gr.Button("Reset Episode", variant="secondary", elem_classes=["btn-secondary"])
                    pass_btn = gr.Button("Pass (Skip Turn)", variant="secondary", elem_classes=["btn-secondary"])
                    step_btn = gr.Button("Submit Step", variant="primary", elem_classes=["btn-primary"], scale=2)

            # ═══════════ RIGHT: OBSERVATION PANEL ═══════════
            with gr.Column(scale=5, elem_classes=["panel-card"]):
                gr.Markdown("### Live State", elem_classes=["section-title"])

                with gr.Tabs():
                    with gr.Tab("Store"):
                        store_json = gr.JSON(label="Store snapshot")
                    with gr.Tab("Stock"):
                        stock_json = gr.JSON(label="Current stock levels")
                    with gr.Tab("Pending Orders"):
                        pending_json = gr.JSON(label="In-transit supplier orders")
                    with gr.Tab("Market"):
                        market_json = gr.JSON(label="Market catalog with seller quotes")
                    with gr.Tab("Marketing"):
                        marketing_json = gr.JSON(label="Marketing actions catalog")
                    with gr.Tab("Reward"):
                        reward_json = gr.JSON(label="Reward breakdown")
                    with gr.Tab("Full Observation"):
                        metadata_json = gr.JSON(label="Complete observation data")

        # ── Helpers ──

        def _initial_env() -> tuple[StoreEnvironment, StoreObservation]:
            env = StoreEnvironment()
            obs = env.reset()
            return env, obs

        def _render_obs(obs: StoreObservation, message: str = "", is_error: bool = False) -> list[Any]:
            kind = "error" if is_error else "success"
            if not message:
                message = obs.summary or "Ready."
                kind = "info"
            return [
                _render_metrics_html(obs),
                _render_status_html(message, kind),
                dump_model(obs.store, exclude_none=True),
                [dump_model(item, exclude_none=True) for item in obs.stock],
                [dump_model(item, exclude_none=True) for item in obs.pending_orders],
                [dump_model(item, exclude_none=True) for item in obs.market],
                [dump_model(item, exclude_none=True) for item in obs.marketing_actions],
                obs.metadata.get("reward_breakdown", {}),
                dump_model(obs, exclude_none=True),
            ]

        render_outputs = [
            metrics_html, status_html,
            store_json, stock_json, pending_json, market_json,
            marketing_json, reward_json, metadata_json,
        ]

        def _update_stock_dropdown(env: StoreEnvironment | None, obs: StoreObservation | None):
            if obs is None:
                return gr.update()
            names = [s.product_name for s in obs.stock]
            return gr.update(choices=names, value=names[0] if names else None)

        # ── Order management ──

        def _df_to_list(df: Any) -> list[list]:
            if df is None:
                return []
            try:
                import pandas as pd
                if isinstance(df, pd.DataFrame):
                    return df.values.tolist() if not df.empty else []
            except ImportError:
                pass
            if isinstance(df, list):
                return df
            return []

        def _add_order(current_orders: Any, product: str, seller: str, qty: int | float):
            if not product or not seller:
                return current_orders
            qty = int(qty) if qty else 10
            row = [product, seller, qty]
            updated = _df_to_list(current_orders)
            updated.append(row)
            return updated

        def _clear_orders():
            return []

        add_order_btn.click(
            fn=_add_order,
            inputs=[orders_display, order_product, order_seller, order_qty],
            outputs=[orders_display],
        )
        clear_orders_btn.click(fn=_clear_orders, outputs=[orders_display])

        # ── Price management ──

        def _add_price(current_prices: Any, product: str, price: int | float):
            if not product:
                return current_prices
            price = int(price) if price else 100
            updated = _df_to_list(current_prices)
            for i, row in enumerate(updated):
                if row[0] == product:
                    updated[i] = [product, price]
                    return updated
            updated.append([product, price])
            return updated

        def _clear_prices():
            return []

        add_price_btn.click(
            fn=_add_price,
            inputs=[prices_display, price_product, price_value],
            outputs=[prices_display],
        )
        clear_prices_btn.click(fn=_clear_prices, outputs=[prices_display])

        # ── Reset ──

        def _reset():
            env, obs = _initial_env()
            stock_names = [s.product_name for s in obs.stock]
            active_ids = set(obs.store.active_action_ids)
            available_labels = [
                lbl for lbl, aid in marketing_id_by_label.items()
                if aid not in active_ids
            ]
            return (
                env, obs,
                [],  # orders display
                [],  # prices display
                gr.update(choices=available_labels, value=[]),  # marketing
                gr.update(choices=stock_names, value=stock_names[0] if stock_names else None),
                *_render_obs(obs, "Episode reset. Day 1 loaded.", False),
            )

        reset_outputs = [
            env_state, obs_state,
            orders_display, prices_display,
            marketing_checkboxes, price_product,
            *render_outputs,
        ]

        reset_btn.click(fn=_reset, inputs=[], outputs=reset_outputs)
        demo.load(fn=_reset, inputs=[], outputs=reset_outputs)

        # ── Step ──

        def _step(
            env: StoreEnvironment | None,
            obs: StoreObservation | None,
            orders_data: Any,
            prices_data: Any,
            selected_marketing: list[str],
        ):
            if env is None or obs is None:
                env, obs = _initial_env()

            orders_rows = _df_to_list(orders_data)
            prices_rows = _df_to_list(prices_data)

            try:
                order_requests = [
                    OrderRequest(product=row[0], seller=row[1], quantity=int(row[2]))
                    for row in orders_rows
                    if len(row) >= 3 and row[0]
                ]
                price_updates = [
                    PriceUpdate(product=row[0], price=int(row[1]))
                    for row in prices_rows
                    if len(row) >= 2 and row[0]
                ]
                action_ids = [
                    marketing_id_by_label[lbl]
                    for lbl in (selected_marketing or [])
                    if lbl in marketing_id_by_label
                ]

                action = StoreAction(
                    orders=order_requests,
                    prices=price_updates if price_updates else None,
                    action_ids=action_ids,
                )
                next_obs = env.step(action)
                message = "Step completed successfully."
            except Exception as exc:
                message = f"{type(exc).__name__}: {exc}"
                stock_names = [s.product_name for s in obs.stock]
                return (
                    env, obs,
                    orders_rows, prices_rows,
                    gr.update(), gr.update(),
                    *_render_obs(obs, message, is_error=True),
                )

            stock_names = [s.product_name for s in next_obs.stock]
            active_ids = set(next_obs.store.active_action_ids)
            available_labels = [
                lbl for lbl, aid in marketing_id_by_label.items()
                if aid not in active_ids
            ]
            return (
                env, next_obs,
                [],  # clear orders
                [],  # clear prices
                gr.update(choices=available_labels, value=[]),
                gr.update(choices=stock_names, value=stock_names[0] if stock_names else None),
                *_render_obs(next_obs, message),
            )

        step_outputs = [
            env_state, obs_state,
            orders_display, prices_display,
            marketing_checkboxes, price_product,
            *render_outputs,
        ]

        step_btn.click(
            fn=_step,
            inputs=[env_state, obs_state, orders_display, prices_display, marketing_checkboxes],
            outputs=step_outputs,
        )

        # ── Pass ──

        def _pass_step(env: StoreEnvironment | None, obs: StoreObservation | None):
            return _step(env, obs, [], [], [])

        pass_btn.click(
            fn=_pass_step,
            inputs=[env_state, obs_state],
            outputs=step_outputs,
        )

    return demo


def main(host: str = "127.0.0.1", port: int = 7860, share: bool = False) -> None:
    """Launch the Gradio app."""
    demo = make_demo()
    demo.launch(server_name=host, server_port=port, share=share, css=CSS, theme=THEME)


if __name__ == "__main__":
    main()