# app.py
# Inventory Management Assistant – HuggingFace / Gradio app

import os
from typing import Dict, List, Tuple

import pandas as pd
import gradio as gr

# ----------------------------- Data Load ----------------------------- #

DATA_FILES: Dict[str, List[str]] = {
    # --- exact files you have in /data/ ---
    "backlog": [
        "data/Backlog_updated.xlsx",
    ],
    "inventory": [
        "data/Inventory_updated.xlsx",
    ],
    "billing": [
        "data/Billing_updated.xlsx",
    ],
    "lead_time": [
        "data/Lead time_updated.xlsx",          # note the space
    ],
    "purchase_orders": [
        "data/Purchase Orders_updated.xlsx",    # note the space
    ],
    "safety_stock": [
        "data/safety_stock_updated.xlsx",
    ],
}


def _load_first_existing(path_candidates: List[str]) -> pd.DataFrame:
    """Try all candidate paths and return the first one that exists."""
    for p in path_candidates:
        if os.path.exists(p):
            if p.lower().endswith(".csv"):
                return pd.read_csv(p)
            else:
                return pd.read_excel(p)
    raise FileNotFoundError(f"None of these files were found: {path_candidates}")


def load_all_data() -> Dict[str, pd.DataFrame]:
    data = {}
    missing = []
    for key, paths in DATA_FILES.items():
        try:
            df = _load_first_existing(paths)
            data[key] = df
        except FileNotFoundError:
            missing.append(key)

    if missing:
        # Fail loudly so HF logs show which logical tables are missing
        raise RuntimeError(f"Data load error – missing logical tables: {missing}")
    return data


DATA = load_all_data()
INV = DATA["inventory"]
BACKLOG = DATA["backlog"]

# ----------------------------- Helper functions ----------------------------- #

def _pick(colnames: List[str], candidates: List[str]):
    for c in candidates:
        if c in colnames:
            return c
    return None


def _basic_inventory_view(df: pd.DataFrame, top_n: int = 20) -> pd.DataFrame:
    """Return a light, generic view that won't break if columns differ."""
    cols = df.columns.tolist()

    mat_col = _pick(cols, ["SAP_MATERIAL_NO", "MATERIAL", "MATERIAL_NO"])
    desc_col = _pick(cols, ["MATERIAL_DESCRIPTION", "MAT_DESC", "DESCRIPTION"])
    plant_col = _pick(cols, ["PLANT", "LOCATION", "SITE"])
    qty_col = _pick(cols, ["QOH", "QTY", "UNRESTRICTED_STOCK", "TOTAL_STOCK"])
    age_col = _pick(cols, ["AGE_DAYS", "DAYS_ON_HAND", "DAYS_COVER"])

    selected = [c for c in [mat_col, desc_col, plant_col, qty_col, age_col] if c]
    if not selected:
        return df.head(top_n)
    return df[selected].head(top_n)


# ----------------------------- Business Logic ----------------------------- #

def get_fast_moving_materials(top_n: int = 25) -> pd.DataFrame:
    """Very simple heuristic: lowest days cover / age, then highest demand/qty."""
    df = INV.copy()
    cols = df.columns.tolist()

    days_cover_col = _pick(cols, ["DAYS_COVER", "AGE_DAYS", "DAYS_ON_HAND"])
    demand_col = _pick(cols, ["AVG_DAILY_DEMAND", "DEMAND_PER_DAY", "ISSUES_PER_DAY"])
    qty_col = _pick(cols, ["QOH", "QTY", "UNRESTRICTED_STOCK", "TOTAL_STOCK"])

    if days_cover_col:
        df = df.sort_values(by=days_cover_col, ascending=True)
    elif demand_col:
        df = df.sort_values(by=demand_col, ascending=False)
    elif qty_col:
        df = df.sort_values(by=qty_col, ascending=False)

    return _basic_inventory_view(df, top_n=top_n)


def get_dead_stock(top_n: int = 25) -> pd.DataFrame:
    """Heuristic: highest age / lowest movement."""
    df = INV.copy()
    cols = df.columns.tolist()

    age_col = _pick(cols, ["AGE_DAYS", "DAYS_ON_HAND", "DAYS_SINCE_MOVEMENT"])
    if age_col:
        df = df.sort_values(by=age_col, ascending=False)
    else:
        # fallback: just low-qty materials
        qty_col = _pick(cols, ["QOH", "QTY", "UNRESTRICTED_STOCK", "TOTAL_STOCK"])
        if qty_col:
            df = df.sort_values(by=qty_col, ascending=True)

    return _basic_inventory_view(df, top_n=top_n)


def get_reallocation_opportunities(top_n: int = 25) -> pd.DataFrame:
    """
    Simple cross-plant reallocation view:
    - Uses inventory + backlog.
    - Marks surplus/shortage per material/plant.
    """
    inv = INV.copy()
    bl = BACKLOG.copy()

    inv_cols = inv.columns.tolist()
    bl_cols = bl.columns.tolist()

    mat_col_i = _pick(inv_cols, ["SAP_MATERIAL_NO", "MATERIAL", "MATERIAL_NO"])
    plant_col_i = _pick(inv_cols, ["PLANT", "LOCATION", "SITE"])
    qty_col_i = _pick(inv_cols, ["QOH", "QTY", "UNRESTRICTED_STOCK", "TOTAL_STOCK"])

    mat_col_b = _pick(bl_cols, ["SAP_MATERIAL_NO", "MATERIAL", "MATERIAL_NO"])
    plant_col_b = _pick(bl_cols, ["PLANT", "LOCATION", "SITE"])
    demand_col_b = _pick(bl_cols, ["OPEN_QTY", "DEMAND_QTY", "BACKLOG_QTY"])

    required = [mat_col_i, plant_col_i, qty_col_i, mat_col_b, plant_col_b, demand_col_b]
    if any(c is None for c in required):
        # If columns don't line up yet, just show generic message.
        return pd.DataFrame(
            {
                "Message": [
                    "Reallocation logic needs aligned columns in inventory & backlog.",
                    f"Inventory columns: {inv_cols}",
                    f"Backlog columns: {bl_cols}",
                ]
            }
        )

    inv_agg = (
        inv.groupby([mat_col_i, plant_col_i])[qty_col_i]
        .sum()
        .reset_index()
        .rename(columns={qty_col_i: "QOH"})
    )

    bl_agg = (
        bl.groupby([mat_col_b, plant_col_b])[demand_col_b]
        .sum()
        .reset_index()
        .rename(columns={mat_col_b: mat_col_i, plant_col_b: plant_col_i, demand_col_b: "DEMAND"})
    )

    merged = inv_agg.merge(bl_agg, on=[mat_col_i, plant_col_i], how="outer").fillna(0)
    merged["NET"] = merged["QOH"] - merged["DEMAND"]

    # Mark surplus / shortage
    merged["STATUS"] = merged["NET"].apply(
        lambda x: "Surplus" if x > 0 else ("Shortage" if x < 0 else "Balanced")
    )

    # Keep only materials which have at least one surplus and one shortage plant
    mat_status = (
        merged.groupby(mat_col_i)["STATUS"]
        .agg(lambda s: set(s))
        .reset_index()
        .rename(columns={"STATUS": "STATUS_SET"})
    )
    interesting_mats = mat_status[
        mat_status["STATUS_SET"].apply(lambda s: {"Surplus", "Shortage"}.issubset(s))
    ][mat_col_i]

    out = merged[merged[mat_col_i].isin(interesting_mats)]
    out = out.sort_values(by=[mat_col_i, "STATUS", "NET"])
    return out.head(top_n * 4)  # multiple rows per material


def get_risk_recommendations(top_n: int = 25) -> pd.DataFrame:
    """
    Very simple 'at-risk' view:
    - Net = demand – stock; positive = shortage.
    """
    inv = INV.copy()
    bl = BACKLOG.copy()

    inv_cols = inv.columns.tolist()
    bl_cols = bl.columns.tolist()

    mat_col_i = _pick(inv_cols, ["SAP_MATERIAL_NO", "MATERIAL", "MATERIAL_NO"])
    plant_col_i = _pick(inv_cols, ["PLANT", "LOCATION", "SITE"])
    qty_col_i = _pick(inv_cols, ["QOH", "QTY", "UNRESTRICTED_STOCK", "TOTAL_STOCK"])

    mat_col_b = _pick(bl_cols, ["SAP_MATERIAL_NO", "MATERIAL", "MATERIAL_NO"])
    plant_col_b = _pick(bl_cols, ["PLANT", "LOCATION", "SITE"])
    demand_col_b = _pick(bl_cols, ["OPEN_QTY", "DEMAND_QTY", "BACKLOG_QTY"])

    required = [mat_col_i, plant_col_i, qty_col_i, mat_col_b, plant_col_b, demand_col_b]
    if any(c is None for c in required):
        return pd.DataFrame(
            {
                "Message": [
                    "Risk recommendations need aligned inventory & backlog columns.",
                    f"Inventory columns: {inv_cols}",
                    f"Backlog columns: {bl_cols}",
                ]
            }
        )

    inv_agg = (
        inv.groupby([mat_col_i, plant_col_i])[qty_col_i]
        .sum()
        .reset_index()
        .rename(columns={qty_col_i: "QOH"})
    )
    bl_agg = (
        bl.groupby([mat_col_b, plant_col_b])[demand_col_b]
        .sum()
        .reset_index()
        .rename(columns={mat_col_b: mat_col_i, plant_col_b: plant_col_i, demand_col_b: "DEMAND"})
    )

    merged = inv_agg.merge(bl_agg, on=[mat_col_i, plant_col_i], how="outer").fillna(0)
    merged["SHORTAGE"] = merged["DEMAND"] - merged["QOH"]
    merged = merged[merged["SHORTAGE"] > 0]

    cols_out = [mat_col_i, plant_col_i, "QOH", "DEMAND", "SHORTAGE"]
    return merged[cols_out].sort_values("SHORTAGE", ascending=False).head(top_n)


def search_inventory(query: str) -> pd.DataFrame:
    """Very light search by material number / description / plant."""
    if not query:
        return _basic_inventory_view(INV, top_n=25)

    df = INV.copy()
    cols = df.columns.tolist()
    mat_col = _pick(cols, ["SAP_MATERIAL_NO", "MATERIAL", "MATERIAL_NO"])
    desc_col = _pick(cols, ["MATERIAL_DESCRIPTION", "MAT_DESC", "DESCRIPTION"])
    plant_col = _pick(cols, ["PLANT", "LOCATION", "SITE"])

    mask = pd.Series([False] * len(df))
    if mat_col:
        mask |= df[mat_col].astype(str).str.contains(query, case=False, na=False)
    if desc_col:
        mask |= df[desc_col].astype(str).str.contains(query, case=False, na=False)
    if plant_col:
        mask |= df[plant_col].astype(str).str.contains(query, case=False, na=False)

    results = df[mask]
    if results.empty:
        return pd.DataFrame({"Message": [f"No inventory rows found for '{query}'"]})
    return _basic_inventory_view(results, top_n=50)


# ----------------------------- Gradio Callbacks ----------------------------- #

def handle_tile(tile: str, history: List[Tuple[str, str]]):
    if history is None:
        history = []

    if tile == "fast":
        user_msg = "Show me fast moving materials."
        df = get_fast_moving_materials()
        assistant_msg = "Here are the current fast-moving materials based on days cover / age."
    elif tile == "reallocate":
        user_msg = "Show stock reallocation possibilities."
        df = get_reallocation_opportunities()
        assistant_msg = "These materials have surplus at some plants and shortages at others."
    elif tile == "risk":
        user_msg = "Show inventory risk recommendations."
        df = get_risk_recommendations()
        assistant_msg = "These materials have net shortages based on backlog vs available stock."
    elif tile == "dead":
        user_msg = "Show dead / slow-moving stock."
        df = get_dead_stock()
        assistant_msg = "These materials appear to be slow-moving or dead stock."
    else:
        user_msg = "Unknown action."
        df = pd.DataFrame({"Message": ["Unknown tile clicked."]})
        assistant_msg = "I couldn't identify that tile."

    history = history + [(("user"), user_msg), (("assistant"), assistant_msg)]
    return history, df


def handle_search(message: str, history: List[Tuple[str, str]]):
    if history is None:
        history = []

    history = history + [("user", message)]
    df = search_inventory(message)
    assistant_msg = "Here is what I found in inventory for your search."
    history = history + [("assistant", assistant_msg)]
    return "", history, df


# ----------------------------- UI Layout ----------------------------- #

CUSTOM_CSS = """
.gradio-container {font-family: 'Segoe UI', system-ui, -apple-system, BlinkMacSystemFont, sans-serif;}
#header-bar {background-color: #002b5c; color: white; padding: 10px 16px; font-size: 20px; font-weight: 600;}
.tile-row button {height: 60px; font-size: 16px; font-weight: 600;}
#faq-bar {background-color: #003f87; color: white; padding: 8px 16px; margin-top: 8px;
          border-radius: 8px; font-size: 15px; font-weight: 500;}
"""

with gr.Blocks(css=CUSTOM_CSS, title="Inventory Management Assistant") as demo:
    gr.HTML('<div id="header-bar">Inventory Assistant</div>')

    with gr.Row(elem_id="tile-row"):
        btn_fast = gr.Button("Fast Moving Materials")
        btn_reallocate = gr.Button("Stock Reallocation")
        btn_risk = gr.Button("Risk Recommendations")
        btn_dead = gr.Button("Dead Stock Materials")

    gr.HTML(
        '<div id="faq-bar">💡 FAQ: Where are we at risk on inventory, and where can we reallocate stock?</div>'
    )

    chatbot = gr.Chatbot(label="Inventory Assistant", height=260)
    results_table = gr.Dataframe(
        headers=[],
        datatype="auto",
        label="Results",
        interactive=False,
        visible=True,
        wrap=True,
        height=260,
    )

    with gr.Row():
        txt = gr.Textbox(
            placeholder="Ask about materials, plants or inventory…",
            show_label=False,
            scale=5,
        )
        btn_search = gr.Button("Search", scale=1)

    # Wire the tiles
    btn_fast.click(
        fn=lambda h: handle_tile("fast", h),
        inputs=chatbot,
        outputs=[chatbot, results_table],
    )
    btn_reallocate.click(
        fn=lambda h: handle_tile("reallocate", h),
        inputs=chatbot,
        outputs=[chatbot, results_table],
    )
    btn_risk.click(
        fn=lambda h: handle_tile("risk", h),
        inputs=chatbot,
        outputs=[chatbot, results_table],
    )
    btn_dead.click(
        fn=lambda h: handle_tile("dead", h),
        inputs=chatbot,
        outputs=[chatbot, results_table],
    )

    # Wire the search bar
    btn_search.click(
        fn=handle_search,
        inputs=[txt, chatbot],
        outputs=[txt, chatbot, results_table],
    )
    txt.submit(
        fn=handle_search,
        inputs=[txt, chatbot],
        outputs=[txt, chatbot, results_table],
    )

if __name__ == "__main__":
    demo.launch()