import streamlit as st
import pandas as pd
import sys, os
from pathlib import Path
import altair as alt

# Backend API
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from api_client import WarehouseAPI

# Add project root for AI modules
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

# Import AI modules
try:
    from forecasting.smart_suggestions import get_smart_suggestions
    from utils.suggestions_panel import display_smart_suggestions
    imports_successful = True
except ImportError as e:
    imports_successful = False
    import_error = str(e)

st.set_page_config(
    page_title="AI Warehouse Management",
    page_icon="",
    layout="wide"
)

# ================================
# DATA LOADING
# ================================
@st.cache_data(show_spinner=False)
def load_data():
    api = WarehouseAPI()
    try:
        inventory = api.get_inventory()
        orders = api.get_orders()
        order_items = api.get_order_items()
        product = api.get_products()
        spare_parts = api.get_spare_parts()
        warehouses = api.get_warehouses()

        # Ensure datetime conversion
        if "order_date" in orders.columns:
            orders["order_date"] = pd.to_datetime(orders["order_date"], errors="coerce")

        if "estimated_delivery" in orders.columns:
            orders["estimated_delivery"] = pd.to_datetime(
                orders["estimated_delivery"], errors="coerce"
            )

        if "order_date" in order_items.columns:
            order_items["order_date"] = pd.to_datetime(
                order_items["order_date"], errors="coerce"
            )

        return inventory, orders, order_items, product, spare_parts, warehouses

    except Exception as e:
        st.error(f"API Error: {e}")
        return [pd.DataFrame()] * 6


# Load data
inventory, orders, order_items, product, spare_parts, warehouses = load_data()

# Store in session for reuse
st.session_state["inventory"] = inventory
st.session_state["orders"] = orders
st.session_state["order_items"] = order_items
st.session_state["product"] = product
st.session_state["spare_parts"] = spare_parts
st.session_state["warehouses"] = warehouses

# Create merged order_data for AI
if not orders.empty and not order_items.empty:
    duplicate_cols = [
        c for c in order_items.columns if c in orders.columns and c != "order_id"
    ]
    order_items_clean = order_items.drop(columns=duplicate_cols, errors="ignore")

    order_data = order_items_clean.merge(
        orders.drop_duplicates(subset=["order_id"]),
        how="left",
        on="order_id",
    )
    st.session_state["order_data"] = order_data
else:
    st.session_state["order_data"] = pd.DataFrame()


# ================================
# PAGE HEADER
# ================================
st.title("AI-Driven Warehouse Stock Management")
st.caption("Executive overview with KPIs, category analytics, and operational performance.")

today = pd.Timestamp.today()


def merge_inventory_product(inventory, product):
    if "product_id" in inventory.columns and "product_id" in product.columns:
        return inventory.merge(product, on="product_id", how="left")
    return inventory


# ================================
# KPI CARDS
# ================================
total_inventory = (
    inventory["q_on_hand"].sum()
    if "q_on_hand" in inventory.columns and not inventory.empty
    else 0
)

processing_orders = 0
if "status" in orders.columns and not orders.empty:
    processing_orders = len(
        orders[orders["status"].astype(str).str.strip().str.lower() == "processing"]
    )

total_products = (
    product["product_id"].nunique()
    if "product_id" in product.columns and not product.empty
    else 0
)

total_warehouses = (
    inventory["warehouse_id"].nunique()
    if "warehouse_id" in inventory.columns and not inventory.empty
    else 0
)

col1, col2, col3, col4 = st.columns(4)
col1.metric("Total Inventory", f"{total_inventory:,}")
col2.metric("Processing Orders", processing_orders)
col3.metric("Total Products", total_products)
col4.metric("Total Warehouses", total_warehouses)


# ================================
# AI SMART SUGGESTIONS PANEL
# ================================
st.divider()
st.markdown("### 🤖 AI Smart Suggestions")
st.markdown("*AI-powered insights updated in real-time*")

if not imports_successful:
    st.error(f"Import Error: {import_error}")
    st.info("Ensure forecasting/smart_suggestions.py and utils/suggestions_panel.py exist")
elif not order_items.empty:
    try:
        with st.spinner("Generating AI insights..."):
            recommendations = get_smart_suggestions(order_items, inventory, top_n=3)

        if recommendations:
            col1, col2, col3 = st.columns(3) 
            cols = [col1, col2, col3]
            for i, rec in enumerate(recommendations[:3]):
                with cols[i]:
                    display_smart_suggestions([rec])
        else:
            st.info("No urgent AI recommendations at this time.")

    except Exception as e:
        st.error(f"Error generating suggestions: {str(e)}")
else:
    st.warning("Need order_items and inventory data for AI insights.")


# ================================
# LOW STOCK ALERTS
# ================================
st.divider()
st.subheader("Low Stock Alerts")

if not inventory.empty and {
    "q_on_hand", "safety_stock", "sku", "warehouse_id"
}.issubset(inventory.columns):

    low_stock = inventory[
        (inventory["q_on_hand"] <= inventory["safety_stock"])
        | (inventory["q_on_hand"] < 20)
    ].copy()

    if not low_stock.empty:
        low_stock["criticality"] = (
            (low_stock["safety_stock"] - low_stock["q_on_hand"])
            / low_stock["safety_stock"].replace(0, 1)
        )
        low_stock = low_stock.sort_values("criticality", ascending=False)

        display_cols = [
            "sku",
            "product_id",
            "warehouse_id",
            "q_on_hand",
            "safety_stock",
            "reorder_point",
        ]
        available_cols = [c for c in display_cols if c in low_stock.columns]

        st.dataframe(
            low_stock[available_cols].head(10),
            width="stretch",
            hide_index=True,
        )
        st.caption(f"{len(low_stock)} items are below safety stock levels.")
    else:
        st.success("All items are adequately stocked.")
else:
    st.warning("Inventory data incomplete — missing q_on_hand/safety_stock/sku/warehouse_id.")


# ================================
# ORDER STATUS GAUGES
# ================================
st.divider()
st.subheader("Order Status KPI")

required_cols = {"status", "order_date", "estimated_delivery"}
if not orders.empty and required_cols.issubset(orders.columns):

    orders_copy = orders.dropna(subset=["order_date", "estimated_delivery"]).copy()

    delivered_on_time = orders_copy[
        (orders_copy["status"].str.lower() == "delivered")
        & (orders_copy["estimated_delivery"] <= today)
    ]

    delayed = orders_copy[
        (orders_copy["status"].str.lower() != "delivered")
        & (orders_copy["estimated_delivery"] < today)
    ]

    in_progress = orders_copy[
        (orders_copy["status"].str.lower() != "delivered")
        & (orders_copy["estimated_delivery"] >= today)
    ]

    total_orders = len(orders_copy)
    delivered_pct = (len(delivered_on_time) / total_orders * 100) if total_orders else 0
    delayed_pct = (len(delayed) / total_orders * 100) if total_orders else 0
    in_progress_pct = (len(in_progress) / total_orders * 100) if total_orders else 0

    # Gauge chart function
    def make_full_gauge(value, label, color):
        base = pd.DataFrame(
            {"value": [value, 100 - value], "category": [label, "Remaining"]}
        )
        chart = (
            alt.Chart(base)
            .mark_arc(innerRadius=60)
            .encode(
                theta=alt.Theta("value:Q"),
                color=alt.Color(
                    "category:N",
                    scale=alt.Scale(
                        domain=[label, "Remaining"],
                        range=[color, "#E0E0E0"],
                    ),
                    legend=None,
                ),
            )
            .properties(width=200, height=200)
        )

        value_text = (
            alt.Chart(pd.DataFrame({"text": [f"{value:.1f}%"]}))
            .mark_text(fontSize=22, fontWeight="bold", color=color)
            .encode(text="text:N")
        )

        label_text = (
            alt.Chart(pd.DataFrame({"text": [label]}))
            .mark_text(fontSize=14, dy=40, color="#FFFFFF")
            .encode(text="text:N")
        )

        return chart + value_text + label_text

    # Display the gauges
    col1, col2, col3 = st.columns(3)
    col1.altair_chart(make_full_gauge(delivered_pct, "Delivered", "#2ecc71"))
    col2.altair_chart(make_full_gauge(delayed_pct, "Delayed", "#f39c12"))
    col3.altair_chart(make_full_gauge(in_progress_pct, "In Progress", "#3498db"))

    # ================================
    # AVERAGE LEAD TIME CHART
    # ================================
    # st.divider()
    # st.subheader("Average Lead Time Over Past 4 Months")

    past_4_months = today - pd.DateOffset(months=4)
    delivered_orders = orders_copy[
        (orders_copy["status"].str.lower() == "delivered")
        & (orders_copy["order_date"] >= past_4_months)
    ].copy()

    if not delivered_orders.empty:
        delivered_orders["lead_time_days"] = (
            delivered_orders["estimated_delivery"]
            - delivered_orders["order_date"]
        ).dt.days
        delivered_orders["order_week"] = (
            delivered_orders["order_date"].dt.to_period("W").dt.start_time
        )

        weekly_lead = (
            delivered_orders.groupby("order_week")["lead_time_days"]
            .mean()
            .reset_index()
        )

        line_chart = (
            alt.Chart(weekly_lead)
            .mark_line(point=True)
            .encode(
                x=alt.X("order_week:T", title="Week"),
                y=alt.Y("lead_time_days:Q", title="Average Lead Time (days)"),
                tooltip=[
                    alt.Tooltip("order_week:T", title="Week"),
                    alt.Tooltip("lead_time_days:Q", title="Avg Lead Time"),
                ],
            )
            .properties(width=700, height=300)
        )
        #st.altair_chart(line_chart, width='stretch')
    else:
        st.info("No delivered orders in the last 4 months.")

else:
    st.warning("Required columns missing for Order KPI metrics.")


# ================================
# CATEGORY ANALYSIS
# ================================
st.divider()
st.subheader("Products by Category")

if not inventory.empty and not product.empty:

    merged_prod = merge_inventory_product(inventory, product)

    if {"category", "q_on_hand"}.issubset(merged_prod.columns):

        categories = merged_prod["category"].dropna().unique()

        if len(categories) > 0:
            selected_category = st.selectbox(
                "Select Product Category", sorted(categories)
            )

            cat_df = merged_prod[merged_prod["category"] == selected_category]

            if "product_name" in cat_df.columns:
                product_sums = (
                    cat_df.groupby("product_name")["q_on_hand"]
                    .sum()
                    .reset_index()
                )

                bar_chart = (
                    alt.Chart(product_sums)
                    .mark_bar()
                    .encode(
                        x=alt.X("product_name:N", sort="-y", title="Product Name"),
                        y=alt.Y("q_on_hand:Q", title="Quantity On Hand"),
                        color= alt.Color("product_name:N", legend=None),
                        tooltip=["product_name", "q_on_hand"],
                    )
                )
                #st.altair_chart(bar_chart, width='stretch')
                #st.altair_chart(bar_chart, width='stretch')
            else:
                st.warning("Missing 'product_name' column in product dataset.")
        else:
            st.info("No product categories found.")
    else:
        st.warning("Missing 'category' or 'q_on_hand' columns.")

else:
    st.warning("Missing inventory or product data.")

col1, col2 = st.columns(2)
with col1: 
    st.subheader("Inventory by Product")
    st.altair_chart(bar_chart, width = 'stretch')
with col2: 
    st.subheader("Average Lead Time (Last 4 Months)")
    st.altair_chart(line_chart, width = 'stretch')

