import os
import sys
import json
from typing import List, Dict, Any

import pandas as pd
import streamlit as st
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain_groq import ChatGroq

# Allow imports from project root
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from api_client import WarehouseAPI


# -----------------------------
# Data loading (Streamlit cache)
# -----------------------------
@st.cache_data(show_spinner=False)
def load_warehouse_data():
    """
    Load all warehouse-related data via the WarehouseAPI.

    Returns:
        inventory_df, order_items_df, product_df,
        spare_parts_df, warehouses_df, orders_df
    """
    api = WarehouseAPI()

    try:
        inventory_df = api.get_inventory()
        order_items_df = api.get_order_items()
        product_df = api.get_products()
        spare_parts_df = api.get_spare_parts()
        warehouses_df = api.get_warehouses()
        orders_df = api.get_orders()
    except Exception as e:
        st.error(f"API load error: {e}")
        empty = pd.DataFrame()
        return empty, empty, empty, empty, empty, empty

    # Ensure datetime fields are parsed
    if not orders_df.empty:
        if "order_date" in orders_df.columns:
            orders_df["order_date"] = pd.to_datetime(
                orders_df["order_date"], errors="coerce"
            )
        if "estimated_delivery" in orders_df.columns:
            orders_df["estimated_delivery"] = pd.to_datetime(
                orders_df["estimated_delivery"], errors="coerce"
            )

    # In case order_items also has order_date
    if not order_items_df.empty and "order_date" in order_items_df.columns:
        order_items_df["order_date"] = pd.to_datetime(
            order_items_df["order_date"], errors="coerce"
        )

    return (
        inventory_df,
        order_items_df,
        product_df,
        spare_parts_df,
        warehouses_df,
        orders_df,
    )


class WarehouseChatbot:
    """
    AI-powered chatbot for answering warehouse inventory and order questions
    using LangChain and Groq.
    """

    def __init__(self, api_key: str | None = None):
        """
        Initialize the chatbot with Groq via LangChain.

        Args:
            api_key: Groq API key (or set GROQ_API_KEY env var)
        """
        self.api_key = api_key or os.getenv("GROQ_API_KEY")
        if not self.api_key:
            raise ValueError(
                "GROQ_API_KEY is not set. "
                "Set it in environment or pass api_key to WarehouseChatbot."
            )

        self.llm = ChatGroq(
            model="llama-3.3-70b-versatile",
            groq_api_key=self.api_key,
            temperature=0.3,
            max_tokens=1024,
        )

        # Conversation history
        self.chat_history: List[Dict[str, str]] = []

        # Load data
        (
            self.inventory_df,
            self.order_items_df,
            self.product_df,
            self.spare_parts_df,
            self.warehouses_df,
            self.orders_df,
        ) = load_warehouse_data()

        # Build merged tables
        self.order_data = self._build_order_data()
        self.inventory_product_df = self._build_inventory_product_data()

    # -----------------------------
    # Internal helpers
    # -----------------------------
    def _build_order_data(self) -> pd.DataFrame:
        """
        Merge orders (header) and order_items (line items) using order_id.

        Ensures a 1-to-many relationship and removes duplicated columns
        that were copied into order_items when it was created.

        Returns:
            DataFrame with combined order + line-item fields.
        """
        if (
            self.orders_df is None
            or self.orders_df.empty
            or self.order_items_df is None
            or self.order_items_df.empty
        ):
            return pd.DataFrame()

        # Use unique orders by order_id
        orders = self.orders_df.drop_duplicates(subset=["order_id"])

        # Remove any header-level columns accidentally duplicated in order_items
        duplicate_cols = [
            c for c in self.order_items_df.columns if c in orders.columns and c != "order_id"
        ]
        clean_items = self.order_items_df.drop(columns=duplicate_cols, errors="ignore")

        try:
            merged = clean_items.merge(
                orders,
                how="left",
                on="order_id",
                validate="many_to-one",  # enforce 1 order -> many items
            )
            return merged
        except Exception as e:
            # Fallback to line items only if merge fails
            print(f"[WarehouseChatbot] Error merging orders and order_items: {e}")
            return clean_items

    def _build_inventory_product_data(self) -> pd.DataFrame:
        """
        Merge inventory with product and warehouse metadata.

        - Links inventory rows to product attributes (name, category, etc.)
        - Links inventory to warehouse attributes (name, location, capacity, etc.)

        Returns:
            DataFrame with inventory + product + warehouse fields.
        """
        if (
            self.inventory_df is None
            or self.inventory_df.empty
            or self.product_df is None
            or self.product_df.empty
        ):
            return pd.DataFrame()

        inv = self.inventory_df.copy()
        prod = self.product_df.copy()

        # Make sure keys exist
        has_product_id = "product_id" in inv.columns and "product_id" in prod.columns
        has_sku = "sku" in inv.columns and "sku" in prod.columns

        # Avoid column name clashes
        if "sku" in inv.columns:
            inv = inv.rename(columns={"sku": "inventory_sku"})

        if has_product_id:
            merged = inv.merge(
                prod,
                how="left",
                on="product_id",
                suffixes=("_inv", "_prod"),
            )
        elif has_sku:
            inv["inventory_sku"] = inv["inventory_sku"].astype(str).str.upper()
            prod["sku"] = prod["sku"].astype(str).str.upper()
            merged = inv.merge(
                prod,
                left_on="inventory_sku",
                right_on="sku",
                how="left",
                suffixes=("_inv", "_prod"),
            )
        else:
            # Can't reliably link inventory to product metadata
            merged = inv

        # Attach warehouse metadata if available
        if (
            self.warehouses_df is not None
            and not self.warehouses_df.empty
            and "warehouse_id" in merged.columns
            and "warehouse_id" in self.warehouses_df.columns
        ):
            wh = self.warehouses_df.copy()
            merged = merged.merge(
                wh,
                how="left",
                on="warehouse_id",
                suffixes=("", "_wh"),
            )

        # Create a canonical SKU column:
        # Prefer product.sku, fall back to inventory_sku
        if "sku" in merged.columns and "inventory_sku" in merged.columns:
            merged["sku_canonical"] = (
                merged["sku"].fillna(merged["inventory_sku"]).astype(str).str.upper()
            )
        elif "sku" in merged.columns:
            merged["sku_canonical"] = merged["sku"].astype(str).str.upper()
        elif "inventory_sku" in merged.columns:
            merged["sku_canonical"] = merged["inventory_sku"].astype(str).str.upper()

        return merged


    # -----------------------------
    # Context builders
    # -----------------------------
    def get_data_context(self) -> str:
        """
        Generate a high-level summary of available data
        to give the model structure and grounding.
        """
        context_parts: list[str] = []

        # Orders header summary
        if self.orders_df is not None and not self.orders_df.empty:
            total_orders = len(self.orders_df)
            date_min = (
                self.orders_df["order_date"].min().date()
                if "order_date" in self.orders_df.columns
                and pd.api.types.is_datetime64_any_dtype(self.orders_df["order_date"])
                else "unknown"
            )
            date_max = (
                self.orders_df["order_date"].max().date()
                if "order_date" in self.orders_df.columns
                and pd.api.types.is_datetime64_any_dtype(self.orders_df["order_date"])
                else "unknown"
            )
            total_revenue = (
                float(self.orders_df["order_total"].sum())
                if "order_total" in self.orders_df.columns
                else 0.0
            )
            unique_customers = (
                self.orders_df["customer_id"].nunique()
                if "customer_id" in self.orders_df.columns
                else 0
            )
            context_parts.append(
                f"Orders: {total_orders} total orders "
                f"(from {date_min} to {date_max}), "
                f"total revenue ≈ ${total_revenue:,.2f}, "
                f"{unique_customers} unique customers."
            )

            if "assigned_warehouse_id" in self.orders_df.columns:
                warehouses_count = self.orders_df["assigned_warehouse_id"].nunique()
                context_parts.append(
                    f"Orders shipped from {warehouses_count} warehouses "
                    f"(field: assigned_warehouse_id)."
                )

        # Order items summary
        if self.order_items_df is not None and not self.order_items_df.empty:
            total_lines = len(self.order_items_df)
            unique_skus = (
                self.order_items_df["sku"].nunique()
                if "sku" in self.order_items_df.columns
                else 0
            )
            context_parts.append(
                f"Order Items: {total_lines} line items across all orders, "
                f"{unique_skus} unique SKUs."
            )

            if "sku" in self.order_items_df.columns and "qty" in self.order_items_df.columns:
                top_skus = (
                    self.order_items_df.groupby("sku")["qty"]
                    .sum()
                    .nlargest(5)
                )
                top_sku_str = ", ".join(
                    f"{sku} ({qty} units)"
                    for sku, qty in top_skus.to_dict().items()
                )
                context_parts.append(
                    f"Top 5 SKUs by ordered quantity: {top_sku_str}."
                )

        # Inventory summary (if present)
        if self.inventory_df is not None and not self.inventory_df.empty:
            context_parts.append(
                f"Inventory: {len(self.inventory_df)} inventory records "
                f"across {self.inventory_df['warehouse_id'].nunique() if 'warehouse_id' in self.inventory_df.columns else '?'} warehouses."
            )

        # Product catalog summary
        if self.product_df is not None and not self.product_df.empty:
            context_parts.append(
                f"Products: {len(self.product_df)} SKUs with attributes like "
                f"product_name, category, base_price, model_number, description, weight, "
                f"dimensions, brand, and status."
            )

        # Warehouses summary
        if self.warehouses_df is not None and not self.warehouses_df.empty:
            context_parts.append(
                f"Warehouses: {len(self.warehouses_df)} total with metadata such as "
                f"warehouse_name, city, state, country, capacity, operating_hours, and manager."
            )

        # Spare parts summary (high-level)
        if self.spare_parts_df is not None and not self.spare_parts_df.empty:
            context_parts.append(
                f"Spare Parts: {len(self.spare_parts_df)} spare part records linked to products."
            )

        return "\n".join(context_parts)

    # -----------------------------
    # Query analysis
    # -----------------------------
    def analyze_query(self, query: str) -> Dict[str, Any]:
        """
        Analyze the query to extract relevant data slices and simple stats.

        Returns:
            {
              'query_type': str | None,
              'data': Any (often a list/dict for the model),
              'insights': List[str]
            }
        """
        analysis: Dict[str, Any] = {
            "query_type": None,
            "data": None,
            "insights": [],
        }

        q = query.lower()

        # Helper flags
        mentions_orders = any(w in q for w in ["order", "orders", "sales"])
        mentions_inventory = any(w in q for w in ["inventory", "stock", "quantity", "on hand"])
        mentions_sku = "sku" in q or "product" in q
        mentions_customer = "customer" in q
        mentions_warehouse = "warehouse" in q
        mentions_revenue = any(w in q for w in ["revenue", "order value", "aov", "sales $"])
        mentions_status = any(w in q for w in ["status", "shipped", "pending", "open", "late"])

        # 1) Order-level questions (orders / sales)
        if mentions_orders and not mentions_inventory and self.orders_df is not None and not self.orders_df.empty:
            analysis["query_type"] = "orders"

            # Recent 20 orders for context
            orders_sorted = self.orders_df.copy()
            if "order_date" in orders_sorted.columns:
                orders_sorted = orders_sorted.sort_values("order_date")

            analysis["data"] = orders_sorted.tail(20).to_dict("records")
            analysis["insights"].append(f"Total orders: {len(self.orders_df)}")

            if "order_total" in self.orders_df.columns:
                total_rev = float(self.orders_df["order_total"].sum())
                avg_order = float(self.orders_df["order_total"].mean())
                analysis["insights"].append(
                    f"Total revenue ≈ ${total_rev:,.2f}, "
                    f"average order value ≈ ${avg_order:,.2f}."
                )

        # 2) Inventory / stock questions (global)
        elif mentions_inventory and not mentions_warehouse and self.inventory_df is not None and not self.inventory_df.empty:
            analysis["query_type"] = "inventory"
            # Only pass a small sample to the model
            analysis["data"] = self.inventory_df.head(50).to_dict("records")
            analysis["insights"].append(
                f"Tracking {len(self.inventory_df)} inventory records."
            )

        # 3) Warehouse inventory questions: "what's inside warehouse X"
        elif mentions_warehouse and self.inventory_product_df is not None and not self.inventory_product_df.empty and not mentions_orders:
            analysis["query_type"] = "warehouse_inventory"

            # Try to extract a warehouse_id token from the query
            # You can support IDs like "WH-001" or "W1" by just grabbing
            # any token that appears in the warehouse_id column.
            warehouse_id_candidate = None
            if "warehouse_id" in self.inventory_product_df.columns:
                candidates = set(
                    str(wid) for wid in self.inventory_product_df["warehouse_id"].dropna().unique()
                )
                for word in query.replace(",", " ").split():
                    cleaned = word.strip(" ,.;:()[]{}")
                    if cleaned in candidates:
                        warehouse_id_candidate = cleaned
                        break

            df = self.inventory_product_df

            if warehouse_id_candidate is not None:
                wh_rows = df[df["warehouse_id"] == warehouse_id_candidate]

                if not wh_rows.empty:
                    # Aggregate view
                    total_units = (
                        wh_rows["q_on_hand"].sum()
                        if "q_on_hand" in wh_rows.columns
                        else None
                    )
                    unique_skus = wh_rows["sku_canonical"].nunique() if "sku_canonical" in wh_rows.columns else None

                    # Provide a sample of rows for the model
                    analysis["data"] = wh_rows.head(50).to_dict("records")

                    insight_parts = [
                        f"Warehouse {warehouse_id_candidate} contains "
                        f"{unique_skus} unique SKUs" if unique_skus is not None else ""
                    ]
                    if total_units is not None:
                        insight_parts.append(f"with ~{int(total_units)} total on-hand units.")

                    # Warehouse metadata if present
                    wh_name = wh_rows["warehouse_name"].iloc[0] if "warehouse_name" in wh_rows.columns else None
                    wh_city = wh_rows["city"].iloc[0] if "city" in wh_rows.columns else None
                    wh_state = wh_rows["state"].iloc[0] if "state" in wh_rows.columns else None

                    if wh_name:
                        insight_parts.append(f"Warehouse name: {wh_name}.")
                    if wh_city or wh_state:
                        loc_str = ", ".join([p for p in [wh_city, wh_state] if p])
                        insight_parts.append(f"Location: {loc_str}.")

                    analysis["insights"].append(" ".join(p for p in insight_parts if p))
                else:
                    analysis["insights"].append(
                        f"No inventory records found for warehouse_id={warehouse_id_candidate}."
                    )
            else:
                # No specific warehouse ID: summarize all warehouses
                if "warehouse_id" in df.columns and "sku_canonical" in df.columns:
                    sku_counts = (
                        df.groupby("warehouse_id")["sku_canonical"]
                        .nunique()
                        .sort_values(ascending=False)
                        .to_dict()
                    )
                    unit_counts = (
                        df.groupby("warehouse_id")["q_on_hand"]
                        .sum()
                        .sort_values(ascending=False)
                        .to_dict()
                        if "q_on_hand" in df.columns
                        else {}
                    )
                    analysis["data"] = {
                        "warehouse_sku_counts": sku_counts,
                        "warehouse_unit_counts": unit_counts,
                    }
                    analysis["insights"].append(
                        "Summarized unique SKUs and total on-hand units per warehouse."
                    )

        # 4) SKU-specific questions (inventory + product metadata, plus optional order history)
        elif mentions_sku:
            analysis["query_type"] = "sku"

            # Try to extract an explicit SKU token like 'SKU-1234'
            sku_candidate = None
            tokens = query.replace(",", " ").split()
            for i, word in enumerate(tokens):
                w = word.strip(" ,.;:()[]{}").upper()
                if w.startswith("SKU-") or w.startswith("SKU_"):
                    sku_candidate = w
                    break
                # Handle patterns like "sku 1234"
                if w == "SKU" and i + 1 < len(tokens):
                    sku_candidate = tokens[i + 1].strip(" ,.;:()[]{}").upper()
                    break

            # If we still don't have a candidate, fall back to nothing and
            # let the model answer conceptually.
            if sku_candidate and self.inventory_product_df is not None and not self.inventory_product_df.empty:
                df = self.inventory_product_df
                mask = df["sku_canonical"].astype(str).str.upper() == sku_candidate
                sku_rows = df[mask]

                if not sku_rows.empty:
                    # Aggregate metrics
                    total_on_hand = sku_rows["q_on_hand"].sum() if "q_on_hand" in sku_rows.columns else None
                    total_reserved = sku_rows["q_reserved"].sum() if "q_reserved" in sku_rows.columns else None
                    n_wh = sku_rows["warehouse_id"].nunique() if "warehouse_id" in sku_rows.columns else None

                    # Sample for the model
                    analysis["data"] = sku_rows.head(50).to_dict("records")

                    product_name = sku_rows["product_name"].iloc[0] if "product_name" in sku_rows.columns else None
                    category = sku_rows["category"].iloc[0] if "category" in sku_rows.columns else None
                    model_number = sku_rows["model_number"].iloc[0] if "model_number" in sku_rows.columns else None

                    insight_bits = [f"SKU {sku_candidate}"]

                    if product_name:
                        insight_bits.append(f"({product_name})")
                    if category:
                        insight_bits.append(f"in category '{category}'")
                    if model_number:
                        insight_bits.append(f"model {model_number}")

                    if n_wh is not None:
                        insight_bits.append(f"is stored in {n_wh} warehouse(s)")
                    if total_on_hand is not None:
                        insight_bits.append(f"with ~{int(total_on_hand)} units on hand")
                    if total_reserved is not None and int(total_reserved) > 0:
                        insight_bits.append(f"and {int(total_reserved)} units reserved")

                    analysis["insights"].append(" ".join(insight_bits) + ".")

                    # Optional: basic order history for this SKU
                    if self.order_data is not None and not self.order_data.empty and "sku" in self.order_data.columns:
                        od = self.order_data.copy()
                        od["sku"] = od["sku"].astype(str).str.upper()
                        od_sku = od[od["sku"] == sku_candidate]
                        if not od_sku.empty and "qty" in od_sku.columns:
                            total_ordered = int(od_sku["qty"].sum())
                            order_count = od_sku["order_id"].nunique()
                            analysis["insights"].append(
                                f"In historical orders, {total_ordered} units of {sku_candidate} "
                                f"were ordered across {order_count} orders."
                            )

        # 5) Customer-focused questions
        elif mentions_customer and self.orders_df is not None and not self.orders_df.empty:
            analysis["query_type"] = "customer"

            if "customer_id" in self.orders_df.columns and "order_total" in self.orders_df.columns:
                top_customers = (
                    self.orders_df.groupby("customer_id")["order_total"]
                    .sum()
                    .nlargest(5)
                    .to_dict()
                )
                analysis["data"] = {"top_customers": top_customers}
                analysis["insights"].append(
                    "Top customers by revenue (customer_id): "
                    + ", ".join(f"{cid} (${val:,.2f})" for cid, val in top_customers.items())
                )

        # 6) Warehouse shipping / revenue questions (orders by warehouse)
        elif mentions_warehouse and mentions_orders and self.orders_df is not None and not self.orders_df.empty:
            analysis["query_type"] = "warehouse_orders"

            if "assigned_warehouse_id" in self.orders_df.columns and "order_total" in self.orders_df.columns:
                by_wh = (
                    self.orders_df.groupby("assigned_warehouse_id")["order_total"]
                    .sum()
                    .sort_values(ascending=False)
                    .to_dict()
                )
                analysis["data"] = {"warehouse_revenue": by_wh}
                analysis["insights"].append(
                    "Orders grouped by assigned_warehouse_id with total revenue."
                )

        # 7) Revenue / value questions
        elif mentions_revenue and self.orders_df is not None and not self.orders_df.empty:
            analysis["query_type"] = "revenue"

            if "order_date" in self.orders_df.columns and "order_total" in self.orders_df.columns:
                by_day = (
                    self.orders_df
                    .set_index("order_date")["order_total"]
                    .resample("D").sum().tail(30)
                )
                analysis["data"] = by_day.reset_index().to_dict("records")
                analysis["insights"].append(
                    "Last 30 days revenue by day based on order_total."
                )

        # 8) Status / fulfillment questions
        elif mentions_status and self.orders_df is not None and not self.orders_df.empty:
            analysis["query_type"] = "status"

            if "status" in self.orders_df.columns:
                status_counts = self.orders_df["status"].value_counts().to_dict()
                analysis["data"] = {"status_counts": status_counts}
                analysis["insights"].append(
                    "Order status distribution: "
                    + ", ".join(f"{s}: {c}" for s, c in status_counts.items())
                )

        # Fallback: provide a small sample of merged order_data if available
        elif self.order_data is not None and not self.order_data.empty:
            analysis["query_type"] = "mixed"
            analysis["data"] = self.order_data.head(30).to_dict("records")
            analysis["insights"].append(
                "Provided a mixed sample of orders + order_items for flexible reasoning."
            )

        return analysis
    # -----------------------------
    # Chat interface
    # -----------------------------
    def chat(self, user_message: str) -> str:
        """
        Process a user message and return an AI response.

        Args:
            user_message: The user's question.

        Returns:
            AI-generated response string.
        """
        # Analyze the query for structure and data slices
        analysis = self.analyze_query(user_message)

        # Build system prompt with warehouse context
        system_prompt = f"""You are an intelligent warehouse management assistant.
You help warehouse managers understand their inventory, orders, customers,
products, spare parts, and warehouse operations.

You have access to:
- Order header data (orders table) including order_id, order_date, customer info,
  order_total, items (list of SKUs), assigned_warehouse_id, status, payment_method,
  shipping_address, estimated_delivery.
- Order line-item data (order_items table) with one row per SKU per order.
- A merged view `order_data` that combines order headers and line-items via order_id.
- Inventory records by SKU and warehouse (inventory table).
- Product catalog data (products table) including product_name, sku, category,
  base_price, model_number, description, weight, dimensions, brand, and status.
- Warehouse metadata (warehouses table) such as warehouse_name, city, state, country,
  capacity, operating_hours, and manager.
- A merged table `inventory_product_df` linking inventory rows with product and
  warehouse metadata via product_id/sku and warehouse_id.

Important semantic rules:
- The SKU in the inventory table and the SKU in the products table refer to
  the same product; `sku_canonical` is the unified SKU identifier.
- For questions like "what's inside warehouse X", use the inventory + product +
  warehouse metadata to describe the contents (SKUs, product names, categories,
  quantities, and warehouse location details).
- For SKU questions, combine product metadata (name, category, model, etc.)
  with inventory distribution across warehouses and, when useful, basic order history.

Warehouse Data Summary:
{self.get_data_context()}

Query Analysis:
Type: {analysis['query_type']}
Insights: {', '.join(analysis['insights']) if analysis['insights'] else 'None'}

Guidelines:
- Be concise but informative.
- Use specific numbers and time ranges when they are provided in the context.
- Clearly distinguish between number of orders and number of line items.
- If the question is about revenue or values, use order_total when available.
- For warehouse content questions, summarize both SKUs and quantities and
  mention key product metadata (names/categories) when possible.
- If exact data is not in the provided snippet, say so and reason approximately.
- Prefer bullet points for lists or recommendations.
- You can suggest follow-up analyses or metrics the user might care about.

Current Query Data (sampled rows or aggregates):
{str(analysis['data'])[:1200] if analysis['data'] is not None else 'No specific data extracted; answer conceptually based on your warehouse knowledge.'}
"""

        # Build message chain
        messages: List[Any] = [SystemMessage(content=system_prompt)]

        # Add recent chat history for context (last 5 turns / 10 messages)
        for msg in self.chat_history[-10:]:
            if msg["role"] == "user":
                messages.append(HumanMessage(content=msg["content"]))
            else:
                messages.append(AIMessage(content=msg["content"]))

        # Add current user message
        messages.append(HumanMessage(content=user_message))

        # Get response from Groq LLM
        response = self.llm.invoke(messages)

        # Update chat history
        self.chat_history.append({"role": "user", "content": user_message})
        self.chat_history.append({"role": "assistant", "content": response.content})

        return response.content

    def clear_history(self):
        """Clear conversation history."""
        self.chat_history = []


# Example usage for local testing (non-Streamlit)
if __name__ == "__main__":
    chatbot = WarehouseChatbot()  # Requires GROQ_API_KEY env var

    example_queries = [
        "How many total orders do we have and what's the total revenue?",
        "What are our top selling SKUs?",
        "Which warehouses are handling the most order value?",
        "Tell me about customer trends.",
        "What is inside warehouse W-001?",
        "Give me details for SKU-12345, including where it's stored.",
    ]

    for q in example_queries:
        print(f"\n🧑 User: {q}")
        ans = chatbot.chat(q)
        print(f"🤖 Assistant: {ans}\n")
