src/streamlit_app.py

"""
Amazon RAG Chatbot — Streamlit UI
Run with: streamlit run app.py
"""

import streamlit as st
import pandas as pd
from sqlalchemy import create_engine
from rag_core import (
    LocalRAG,
    build_clean_views,
    ensure_llm_up,
    DB_PATH,
    LLM_MODEL,
    HF_TOKEN,
)

st.set_page_config(
    page_title="Amazon Analytics Chatbot",
    page_icon="📦",
    layout="wide",
    initial_sidebar_state="expanded",
)

st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap');

    .stApp { font-family: 'Inter', sans-serif; background: #0d1117; }
    .block-container { padding-top: 0.5rem !important; max-width: 1100px; }

    .main-header {
        background: linear-gradient(135deg, #161b22 0%, #21262d 100%);
        border: 1px solid #30363d; border-radius: 16px;
        padding: 1.8rem 2rem; margin-bottom: 1.5rem;
        display: flex; align-items: center; gap: 1.2rem;
    }
    .main-header .logo { font-size: 2.8rem; line-height: 1; }
    .main-header h1 { font-weight: 700; font-size: 1.7rem; color: #f0f6fc; margin: 0; letter-spacing: -0.5px; }
    .main-header p  { color: #8b949e; font-size: 0.9rem; margin: 0.2rem 0 0 0; }

    .chat-wrap { display: flex; flex-direction: column; gap: 0.8rem; margin-bottom: 1rem; }

    .msg-user {
        align-self: flex-end;
        background: linear-gradient(135deg, #1f6feb 0%, #388bfd 100%);
        color: #fff; border-radius: 18px 18px 4px 18px;
        padding: 0.9rem 1.2rem; max-width: 75%;
        font-size: 0.92rem; line-height: 1.55;
        box-shadow: 0 2px 12px rgba(31,111,235,0.3);
    }
    .msg-user .lbl { font-size: 0.72rem; font-weight: 600; opacity: 0.8; margin-bottom: 0.3rem; text-transform: uppercase; letter-spacing: 0.8px; }

    .msg-bot {
        align-self: flex-start; background: #161b22;
        border: 1px solid #30363d; border-left: 3px solid #3fb950;
        color: #c9d1d9; border-radius: 4px 18px 18px 18px;
        padding: 0.9rem 1.2rem; max-width: 85%;
        font-size: 0.92rem; line-height: 1.6;
    }
    .msg-bot .lbl { font-size: 0.72rem; font-weight: 600; color: #3fb950; margin-bottom: 0.3rem; text-transform: uppercase; letter-spacing: 0.8px; }
    .msg-bot.err  { border-left-color: #f85149; }
    .msg-bot.err .lbl { color: #f85149; }

    .badge { display: inline-block; font-family: 'JetBrains Mono', monospace; font-size: 0.65rem; font-weight: 600; padding: 0.15rem 0.5rem; border-radius: 4px; margin-bottom: 0.5rem; text-transform: uppercase; letter-spacing: 1px; }
    .badge-sql   { background: #1f6feb; color: #fff; }
    .badge-rag   { background: transparent; color: #58a6ff; border: 1px solid #58a6ff; }
    .badge-error { background: #f85149; color: #fff; }

    .sql-block { background: #0d1117; border: 1px solid #30363d; border-radius: 8px; padding: 0.75rem 1rem; margin-top: 0.6rem; font-family: 'JetBrains Mono', monospace; font-size: 0.75rem; color: #8b949e; overflow-x: auto; white-space: pre-wrap; word-break: break-all; }

    section[data-testid="stSidebar"] { background: #0d1117; border-right: 1px solid #21262d; }
    section[data-testid="stSidebar"] h3 { color: #f0f6fc !important; font-size: 0.95rem !important; }
    section[data-testid="stSidebar"] p, section[data-testid="stSidebar"] li { color: #8b949e !important; font-size: 0.83rem !important; }

    .pill { display: inline-flex; align-items: center; gap: 0.3rem; padding: 0.3rem 0.8rem; border-radius: 99px; font-size: 0.8rem; font-weight: 500; }
    .pill-green { background: #12261e; color: #3fb950; border: 1px solid #238636; }
    .pill-red   { background: #2d1b1b; color: #f85149; border: 1px solid #da3633; }

    .stButton > button { font-size: 0.81rem; border-radius: 8px; background: #161b22; border: 1px solid #30363d; color: #c9d1d9; transition: all 0.15s; width: 100%; }
    .stButton > button:hover { background: #21262d; border-color: #58a6ff; color: #f0f6fc; }

    .empty-state { text-align: center; padding: 4rem 2rem; color: #6e7681; }
    .empty-state .icon { font-size: 3rem; margin-bottom: 0.8rem; }
    .empty-state h3 { color: #8b949e; font-weight: 500; margin-bottom: 0.5rem; }
    .empty-state p  { font-size: 0.88rem; }

    hr { border-color: #21262d !important; }
    .stDataFrame { border-radius: 10px; overflow: hidden; }
</style>
""", unsafe_allow_html=True)


@st.cache_resource(show_spinner="Initializing RAG engine...")
def init_bot():
    engine = create_engine(f"sqlite:///{DB_PATH}", echo=False)
    build_clean_views(engine)
    bot = LocalRAG(engine)
    return bot, engine


# ── Sidebar ──
with st.sidebar:
    st.markdown("### System Status")
    llm_ok = ensure_llm_up()
    if llm_ok:
        st.markdown(f'<div class="pill pill-green">&#9679; LLM ({LLM_MODEL.split("/")[-1]}) Online</div>', unsafe_allow_html=True)
    else:
        st.markdown(f'<div class="pill pill-red">&#9679; LLM ({LLM_MODEL.split("/")[-1]}) Offline</div>', unsafe_allow_html=True)
        if not HF_TOKEN:
            st.warning("RAG mode unavailable.\n\n**HF_TOKEN** is not set. Add it in\nSpace Settings → Secrets.")
        else:
            st.warning("RAG mode unavailable.\nLLM API is unreachable. Check your HF_TOKEN or model quota.")

    st.markdown("<br>", unsafe_allow_html=True)
    st.markdown("### Data Coverage")
    st.markdown("""
- **Orders** — revenue, units, AOV, ASP, profit
- **Business Reports** — sessions, page views, buy box, conversion
- **Advertising** — spend, impressions, clicks, ACOS, ROAS
- **Keepa** — ratings, reviews, sales rank
- **Filters** — year, quarter, month, last N days/weeks/months
- **Dimensions** — country, city, state, channel, ASIN, SKU
    """)

    st.markdown("---")
    st.markdown("### Example Questions")
    examples = [
        "Total revenue in 2023 Q1",
        "Monthly sessions trend last 30 days",
        "Total ad spend in 2024",
        "2024 H1 B2B revenue by state",
        "Average buy box percentage",
        "Top search terms by spend",
    ]
    for eq in examples:
        if st.button(eq, key=f"ex_{eq}", use_container_width=True):
            st.session_state["pending_question"] = eq

    st.markdown("---")
    if st.button("Clear conversation", use_container_width=True, type="secondary"):
        st.session_state["messages"] = []
        st.rerun()

    st.markdown('<p style="color:#484f58;font-size:0.72rem;text-align:center;margin-top:1rem;">Amazon RAG v3 · SQL + Semantic Search</p>', unsafe_allow_html=True)


# ── Header ──
st.markdown("""
<div class="main-header">
    <div class="logo">&#128230;</div>
    <div>
        <h1>Amazon Analytics Chatbot</h1>
        <p>Ask questions about your data in plain English — powered by SQL &amp; semantic search</p>
    </div>
</div>
""", unsafe_allow_html=True)

bot, engine = init_bot()

if "messages" not in st.session_state:
    st.session_state["messages"] = []

# ── Chat history ──
if not st.session_state["messages"]:
    st.markdown("""
    <div class="empty-state">
        <div class="icon">&#128172;</div>
        <h3>Start a conversation</h3>
        <p>Type a question below or pick an example from the sidebar.<br>
        Try: <em>"What was total revenue in 2023?"</em></p>
    </div>
    """, unsafe_allow_html=True)
else:
    st.markdown('<div class="chat-wrap">', unsafe_allow_html=True)
    for msg in st.session_state["messages"]:
        if msg["role"] == "user":
            st.markdown(f'<div class="msg-user"><div class="lbl">You</div>{msg["content"]}</div>', unsafe_allow_html=True)
        else:
            st.markdown(f'<div class="msg-bot {msg.get("css_class","")}" ><div class="lbl">Assistant</div>{msg["content"]}</div>', unsafe_allow_html=True)
            if msg.get("dataframe") is not None:
                st.dataframe(pd.DataFrame(msg["dataframe"]), use_container_width=True, hide_index=True)
    st.markdown('</div>', unsafe_allow_html=True)

# ── Input ──
pending    = st.session_state.pop("pending_question", None)
user_input = st.chat_input("Ask about your Amazon data...")
question   = pending or user_input

if question:
    st.session_state["messages"].append({"role": "user", "content": question})
    st.markdown(f'<div class="msg-user"><div class="lbl">You</div>{question}</div>', unsafe_allow_html=True)

    with st.spinner("Processing..."):
        result = bot.answer(question, k=6)

    mode          = result.get("mode", "unknown")
    response_html = ""
    css_class     = ""
    df_rows       = None

    if mode == "template_sql":
        rows          = result.get("result", [])
        response_html = (
            f'<span class="badge badge-sql">SQL</span>&nbsp;{result["answer"]}'
            f'<details><summary style="font-size:0.8rem;color:#6e7681;margin-top:0.5rem;cursor:pointer;">View SQL query</summary>'
            f'<div class="sql-block">{result["sql"]}</div></details>'
        )
        if len(rows) > 1:
            df_rows = rows

    elif mode == "template_sql_error":
        css_class     = "err"
        response_html = (
            f'<span class="badge badge-error">SQL Error</span>&nbsp;<strong>{result.get("error","Unknown error")}</strong>'
        )
        if result.get("sql"):
            response_html += f'<details><summary style="font-size:0.8rem;color:#6e7681;margin-top:0.5rem;cursor:pointer;">View SQL</summary><div class="sql-block">{result["sql"]}</div></details>'

    elif mode == "rag":
        response_html = f'<span class="badge badge-rag">Semantic</span>&nbsp;{result["answer"]}'
    else:
        response_html = f"Unknown mode: {mode}"

    st.markdown(f'<div class="msg-bot {css_class}"><div class="lbl">Assistant</div>{response_html}</div>', unsafe_allow_html=True)

    if df_rows:
        st.dataframe(pd.DataFrame(df_rows), use_container_width=True, hide_index=True)

    st.session_state["messages"].append({
        "role": "bot", "content": response_html,
        "css_class": css_class, "dataframe": df_rows,
    })