import os
from collections import Counter
from typing import List, Dict, Tuple

import gradio as gr
import pandas as pd

try:
    from huggingface_hub import InferenceClient
except ImportError:
    InferenceClient = None

# -----------------------------------------------------
# DATA LOADING (same logic as before, no UI changes)
# -----------------------------------------------------
DATA_CANDIDATES = [
    "Backlog_updated.xlsx",
    "Backlog_updated.xls",
    "data/Backlog_updated.xlsx",
    "data/Backlog_updated.xls",
]

SO_COL = "SALES_ORDER_NO"
CUSTOMER_COL = "CUSTOMER NAME"
WORTH_COL = "EXTENDED_RESALE"
DATE_COL = "ATP_DATE_RAW"

BLOCK_COLUMNS = ["CREDIT BLOCK", "DELIVERY BLOCK", "PRICING BLOCK", "SHIP DEBIT BLOCK"]
BLOCK_MAP = {
    "CREDIT BLOCK": "Credit",
    "DELIVERY BLOCK": "Delivery",
    "PRICING BLOCK": "Pricing",
    "SHIP DEBIT BLOCK": "Ship Debit",
}
BLOCKS_ORDER = ["Pricing", "Credit", "Delivery", "Ship Debit"]


def find_data_file() -> str:
    for p in DATA_CANDIDATES:
        if os.path.exists(p):
            return p
    raise FileNotFoundError(f"Could not find backlog file in {DATA_CANDIDATES}")


def load_raw_df(path: str) -> pd.DataFrame:
    df = pd.read_excel(path)
    df.columns = df.columns.str.strip().str.upper()
    needed = [SO_COL, CUSTOMER_COL, WORTH_COL, DATE_COL, *BLOCK_COLUMNS]
    missing = [c for c in needed if c not in df.columns]
    if missing:
        raise KeyError(f"Missing required columns in Excel: {missing}")
    return df


DATA_PATH = find_data_file()
RAW_DF = load_raw_df(DATA_PATH)

# -----------------------------------------------------
# ANALYTICS HELPERS (same computations as Flask app)
# -----------------------------------------------------
def fmt_money(x: float) -> str:
    return "${:,.0f}".format(float(x or 0))


def block_summary_df() -> pd.DataFrame:
    tmp = RAW_DF.copy()
    tmp["BLOCK_TYPE"] = tmp[BLOCK_COLUMNS].apply(
        lambda row: " + ".join(
            [BLOCK_MAP[c] for c in BLOCK_COLUMNS if str(row[c]).upper() == "X"]
        )
        or "No Block",
        axis=1,
    )
    tmp = tmp[tmp["BLOCK_TYPE"] != "No Block"]
    if tmp.empty:
        return pd.DataFrame(columns=["BLOCK_TYPE", WORTH_COL])
    agg = (
        tmp.groupby("BLOCK_TYPE", as_index=False)[WORTH_COL]
        .sum()
        .sort_values(WORTH_COL, ascending=False)
    )
    return agg


def actionable_focus_df():
    bs = block_summary_df()
    if bs.empty:
        return pd.DataFrame(), "No actionable focus area data."
    total_val = bs[WORTH_COL].sum()
    bs["SHARE_%"] = (bs[WORTH_COL] / total_val) * 100
    top2 = bs.head(2)
    blocks = ", ".join(top2["BLOCK_TYPE"])
    insight = f"{blocks} contribute {top2['SHARE_%'].sum():.1f}% of total blocked value."
    return top2, insight


def multiple_block_df() -> pd.DataFrame:
    tmp = RAW_DF.copy()
    tmp["BLOCK_COUNT"] = tmp[BLOCK_COLUMNS].apply(
        lambda row: sum(str(x).upper() == "X" for x in row), axis=1
    )
    multiblocks = tmp[tmp["BLOCK_COUNT"] > 1][
        [SO_COL, CUSTOMER_COL, WORTH_COL, "BLOCK_COUNT"]
    ].copy()
    return multiblocks


# Fiscal year helper: FY runs Apr–Mar, FY = year it ENDS (like we discussed)
def _cb_fiscal_year_quarter(dt: pd.Timestamp, start_month: int = 4) -> Tuple[int, int]:
    if pd.isna(dt):
        return 0, 0
    m, y = dt.month, dt.year
    shifted = ((m - start_month) % 12) + 1
    fq = ((shifted - 1) // 3) + 1
    fy = y if m < start_month else y + 1
    return fy, fq


def quarter_trends_df() -> pd.DataFrame:
    temp = RAW_DF.copy()
    temp[DATE_COL] = pd.to_datetime(temp[DATE_COL], errors="coerce")

    recs = []
    for b in BLOCK_COLUMNS:
        subset = temp[temp[b].astype(str).str.upper().eq("X")].dropna(subset=[DATE_COL])
        if subset.empty:
            continue
        fy_fq = subset[DATE_COL].apply(
            lambda d: _cb_fiscal_year_quarter(d, start_month=4)
        )
        subset = subset.assign(FY=[t[0] for t in fy_fq], FQ=[t[1] for t in fy_fq])

        grouped = subset.groupby(["FY", "FQ"], as_index=False)[WORTH_COL].sum()
        grouped["BLOCK_TYPE"] = BLOCK_MAP[b]
        grouped["QUARTER"] = (
            grouped["FY"].astype(int).astype(str)
            + "Q"
            + grouped["FQ"].astype(int).astype(str)
        )
        recs.append(grouped[["QUARTER", "BLOCK_TYPE", WORTH_COL, "FY", "FQ"]])

    if not recs:
        return pd.DataFrame(columns=["QUARTER", "BLOCK_TYPE", WORTH_COL])
    qdf = pd.concat(recs, ignore_index=True)
    qdf = qdf.sort_values(["FY", "FQ"]).drop(columns=["FY", "FQ"])
    return qdf


# -----------------------------------------------------
# CHAT LOGIC (same rules as earlier Flask chatbot)
# -----------------------------------------------------
def analytics_reply(user_q: str) -> str:
    q = user_q.lower().strip()
    reply = "I can help with blocked orders, customers, and quarterly trends."

    if "customer" in q:
        cust_summary = (
            RAW_DF.groupby(CUSTOMER_COL)[WORTH_COL]
            .sum()
            .sort_values(ascending=False)
            .reset_index()
        )
        if cust_summary.empty:
            return "I couldn't find any customers with backlog in the file."
        top = cust_summary.iloc[0]
        cust_name, value = top[CUSTOMER_COL], top[WORTH_COL]
        block_counts = (
            RAW_DF[RAW_DF[CUSTOMER_COL] == cust_name][BLOCK_COLUMNS]
            .apply(
                lambda row: [
                    BLOCK_MAP[c]
                    for c in BLOCK_COLUMNS
                    if str(row[c]).upper() == "X"
                ],
                axis=1,
            )
            .explode()
            .value_counts()
        )
        main_block = block_counts.index[0] if not block_counts.empty else "general"
        reply = (
            f"{cust_name} currently holds the highest backlog value ({fmt_money(value)}), "
            f"with **{main_block}** as the main issue. Coordinate with that team to resolve it."
        )

    elif "multiple" in q or "two blocks" in q or "more than one" in q:
        mb = multiple_block_df()
        if mb.empty:
            reply = "No orders have multiple active blocks right now."
        else:
            reply = (
                f"{mb[SO_COL].nunique()} sales orders have **multiple active blocks**. "
                f"Prioritize clearing these to restore smoother flow in the backlog."
            )

    elif "quarter" in q or "trend" in q:
        qdf = quarter_trends_df()
        if qdf.empty:
            reply = "No quarterly trend data available in the file."
        else:
            # Pick current FY quarter if present, else last available
            today = pd.Timestamp.today()
            fy_end, fq = _cb_fiscal_year_quarter(today, start_month=4)
            current_key = f"{fy_end}Q{fq}"

            available = set(qdf["QUARTER"])
            if current_key in available:
                q_sel = current_key
            else:
                q_sel = sorted(
                    qdf["QUARTER"].unique(),
                    key=lambda x: (int(x[:4]), int(x[-1])),
                )[-1]

            q_now_df = qdf[qdf["QUARTER"] == q_sel]
            if not q_now_df.empty:
                top_row = q_now_df.sort_values(WORTH_COL, ascending=False).iloc[0]
                fy_end = int(q_sel[:4])
                fq_now = int(q_sel[-1])
                reply = (
                    f"In **Q{fq_now} FY{fy_end}**, the **{top_row['BLOCK_TYPE']}** block shows the "
                    f"highest exposure at **{fmt_money(top_row[WORTH_COL])}**. "
                    "Prepare mitigation plans with the responsible owners."
                )
            else:
                reply = f"No data found for fiscal quarter {current_key}."

    elif "block" in q:
        bs = block_summary_df()
        if bs.empty:
            reply = "I couldn't find any blocked orders in the file."
        else:
            top = bs.iloc[0]
            reply = (
                f"**{top['BLOCK_TYPE']}** block represents the largest backlog exposure "
                f"at **{fmt_money(top[WORTH_COL])}**. "
                "Engage the respective team to address these issues first."
            )

    else:
        # Generic fallback
        reply = (
            "I can help you with:\n"
            "- Which block has the highest exposure\n"
            "- Which customers face backlog\n"
            "- Orders with multiple blocks\n"
            "- Quarterly trend of blocked exposure\n\n"
            "Try asking: **'Quarterly trend?'** or **'Which block has highest exposure?'**"
        )

    return reply


def maybe_rewrite_with_llm(text: str) -> str:
    """
    Optional LLM rewriter. If HF_TOKEN and HF_MODEL_ID are set and
    huggingface_hub is installed, we rewrite the analytics insight.
    """
    hf_token = os.getenv("HF_TOKEN")
    model_id = os.getenv("HF_MODEL_ID", "").strip()

    if not hf_token or not model_id or InferenceClient is None:
        return text  # analytics-only mode

    try:
        client = InferenceClient(model=model_id, token=hf_token)
        prompt = (
            "You are an operations assistant. Rewrite the following insight for a VP of Operations. "
            "Use at most 3 short bullet points. Be concise and actionable.\n\n"
            f"Insight:\n{text}\n\nBullets:\n-"
        )
        resp = client.text_generation(
            prompt,
            max_new_tokens=180,
            temperature=0.3,
            do_sample=True,
        )
        return resp.strip()
    except Exception:
        # If anything fails, fall back to the original
        return text


def chat_fn(message: str, history: List[Tuple[str, str]]) -> str:
    base = analytics_reply(message)
    final = maybe_rewrite_with_llm(base)
    return final


# -----------------------------------------------------
# GRADIO UI – Tiles + Chat (no Flask)
# -----------------------------------------------------
CUSTOM_CSS = """
.chatbot {max-height: 480px;}
.gradio-container {font-family: 'Segoe UI', system-ui, sans-serif;}
.tiles-row button {
    border-radius: 10px !important;
    border: 1px solid #e5e7eb !important;
    background: #ffffff !important;
    font-weight: 600 !important;
    color: #0A2048 !important;
}
.tiles-row button:hover {
    background: #f1f5f9 !important;
}
"""

with gr.Blocks(css=CUSTOM_CSS, title="Blocks AI Assistant") as demo:
    gr.Markdown(
        "## Blocks AI Assistant  \n"
        "Hybrid **AI + backlog analytics** over your blocked orders dataset."
    )

    with gr.Row(elem_classes="tiles-row"):
        btn_summary = gr.Button("Overall Block Summary")
        btn_focus = gr.Button("Actionable Focus Area")
    with gr.Row(elem_classes="tiles-row"):
        btn_multi = gr.Button("Multiple Block Orders")
        btn_quarter = gr.Button("Quarterly Trend")
        btn_faq = gr.Button("FAQ: What can you do?")

    chatbot = gr.Chatbot(
        label="Conversation",
        type="messages",
        height=420,
    )
    msg = gr.Textbox(
        placeholder="Ask about blocks, customers, or quarterly trends…",
        show_label=False,
    )
    clear_btn = gr.Button("Clear Chat")

    def handle_user_message(user_message, history):
        reply = chat_fn(user_message, history or [])
        (history or []).append({"role": "user", "content": user_message})
        history.append({"role": "assistant", "content": reply})
        return "", history

    msg.submit(handle_user_message, [msg, chatbot], [msg, chatbot])
    clear_btn.click(lambda: ([], ""), outputs=[chatbot, msg])

    # Tile → pre-defined questions wired into same logic
    def ask_predefined(prompt, history):
        reply = chat_fn(prompt, history or [])
        (history or []).append({"role": "user", "content": prompt})
        history.append({"role": "assistant", "content": reply})
        return history

    btn_summary.click(
        ask_predefined,
        inputs=[chatbot],
        outputs=[chatbot],
        _js=None,
        kwargs={"prompt": "Which block has highest exposure?"},
    )
    btn_focus.click(
        ask_predefined,
        inputs=[chatbot],
        outputs=[chatbot],
        kwargs={"prompt": "Which block should we focus on?"},
    )
    btn_multi.click(
        ask_predefined,
        inputs=[chatbot],
        outputs=[chatbot],
        kwargs={"prompt": "Orders with multiple blocks?"},
    )
    btn_quarter.click(
        ask_predefined,
        inputs=[chatbot],
        outputs=[chatbot],
        kwargs={"prompt": "Quarterly trend?"},
    )
    btn_faq.click(
        ask_predefined,
        inputs=[chatbot],
        outputs=[chatbot],
        kwargs={"prompt": "What can you do?"},
    )

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