import os, re, json
from typing import List, Dict, Any, Optional, Tuple

import numpy as np
import pandas as pd
import streamlit as st
from dotenv import load_dotenv, find_dotenv
from openai import OpenAI
import snowflake.connector as sf
import re

# ---------- Setup ----------
DOTENV_PATH = find_dotenv(usecwd=True) or os.path.join(os.getcwd(), ".env")
load_dotenv(DOTENV_PATH, override=True)

st.set_page_config(page_title="Job Listings (Conversational RAG)", page_icon="💼", layout="wide")
st.title("💼 Adobe's Job Listings Conversational Search")

# Env / config
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
SF_ACCOUNT  = os.getenv("SNOWFLAKE_ACCOUNT")
SF_USER     = os.getenv("SNOWFLAKE_USER")
SF_WAREHOUSE= os.getenv("SNOWFLAKE_WAREHOUSE")
SF_DATABASE = os.getenv("SNOWFLAKE_DATABASE", "JOB_LISTINGS_100")
SF_SCHEMA   = os.getenv("SNOWFLAKE_SCHEMA", "PUBLIC")
SF_ROLE     = os.getenv("SNOWFLAKE_ROLE")
SF_PAT      = os.getenv("SNOWFLAKE_PAT") or os.getenv("SNOWFLAKE_PROGRAMMATIC_ACCESS_TOKEN")

TABLE_NAME = os.getenv("JOB_TABLE", f"{SF_DATABASE}.{SF_SCHEMA}.SAMPLE_JOB_LISTINGS")
TITLE_COL  = "TITLE"
DESC_COL   = "DESCRIPTION"
ID_COL     = os.getenv("ID_COL", "ID")  # if you don't have an ID, we’ll hash later

# --- Brand/Vendor knowledge used by the LLM and SQL filters
VENDOR_KB = {
    # Adobe (partial but high coverage)
    "photoshop":"Adobe","lightroom":"Adobe","illustrator":"Adobe","indesign":"Adobe","premiere pro":"Adobe",
    "after effects":"Adobe","audition":"Adobe","acrobat":"Adobe","xd":"Adobe","animate":"Adobe","bridge":"Adobe",
    "fresco":"Adobe","dimension":"Adobe","substance 3d":"Adobe","firefly":"Adobe","character animator":"Adobe",
    "media encoder":"Adobe","express":"Adobe",
    # Non-Adobe design / UI
    "figma":"Figma","sketch":"Sketch","canva":"Canva","coreldraw":"Corel","affinity designer":"Serif",
    "affinity photo":"Serif","affinity publisher":"Serif","quarkxpress":"Quark","inkscape":"Inkscape","krita":"Krita",
    # Non-Adobe photo/video/audio/3D
    "davinci resolve":"Blackmagic","final cut pro":"Apple","imovie":"Apple","logic pro":"Apple","garageband":"Apple",
    "media composer":"Avid","vegas pro":"MAGIX","capcut":"ByteDance",
    "blender":"Blender","maya":"Autodesk","3ds max":"Autodesk","autocad":"Autodesk","cinema 4d":"Maxon",
    "procreate":"Savage","capture one":"Phase One","gimp":"GIMP","unreal engine":"Epic","unity":"Unity"
}
ADOBE_SET = {k for k,v in VENDOR_KB.items() if v == "Adobe"}
NON_ADOBE_SET = {k for k,v in VENDOR_KB.items() if v != "Adobe"}

import re

def _tool_alt_with_boundaries(names: list[str]) -> str:
    """
    Build a Snowflake-friendly regex that matches any of the given tool names as whole tokens,
    without using look-behind. We:
      - lowercase and escape each name,
      - allow spaces or hyphens between words,
      - wrap with (^|[^A-Za-z0-9])  …  ([^A-Za-z0-9]|$) as boundaries.
    """
    def _escape_tool(n: str) -> str:
        s = re.escape(n.lower())         # escape regex metacharacters
        s = s.replace(r"\ ", r"[\s\-]+") # allow spaces/hyphens between words
        return s

    alts = "|".join(_escape_tool(n) for n in names)
    # No lookarounds; just capture boundary chars in groups 1 and 2.
    return rf"(^|[^A-Za-z0-9])(?:{alts})([^A-Za-z0-9]|$)"

def count_adobe(conn) -> int:
    rx_adb = _tool_alt_with_boundaries(sorted(ADOBE_SET))
    sql = f"""
      SELECT COUNT(*)
      FROM {TABLE_NAME}
      WHERE REGEXP_INSTR(LOWER({TITLE_COL}), %s) > 0
         OR REGEXP_INSTR(LOWER({DESC_COL}),  %s) > 0
    """
    with conn.cursor() as cur:
        cur.execute(sql, (rx_adb, rx_adb))
        return int(cur.fetchone()[0] or 0)

def count_non_adobe(conn) -> int:
    rx_non = _tool_alt_with_boundaries(sorted(NON_ADOBE_SET))
    rx_adb = _tool_alt_with_boundaries(sorted(ADOBE_SET))
    sql = f"""
      SELECT COUNT(*)
      FROM {TABLE_NAME}
      WHERE (
        REGEXP_INSTR(LOWER({TITLE_COL}), %s) > 0 OR
        REGEXP_INSTR(LOWER({DESC_COL}),  %s) > 0
      )
      AND NOT (
        REGEXP_INSTR(LOWER({TITLE_COL}), %s) > 0 OR
        REGEXP_INSTR(LOWER({DESC_COL}),  %s) > 0
      )
    """
    with conn.cursor() as cur:
        cur.execute(sql, (rx_non, rx_non, rx_adb, rx_adb))
        return int(cur.fetchone()[0] or 0)


def count_adobe(conn) -> int:
    rx_adb = _word_boundary_regex(sorted(ADOBE_SET))
    sql = f"""
      SELECT COUNT(*)
      FROM {TABLE_NAME}
      WHERE REGEXP_INSTR(LOWER({TITLE_COL}), %s) > 0
         OR REGEXP_INSTR(LOWER({DESC_COL}),  %s) > 0
    """
    with conn.cursor() as cur:
        cur.execute(sql, (rx_adb, rx_adb))
        return int(cur.fetchone()[0] or 0)

def count_non_adobe(conn) -> int:
    rx_non = _word_boundary_regex(sorted(NON_ADOBE_SET))
    rx_adb = _word_boundary_regex(sorted(ADOBE_SET))
    sql = f"""
      SELECT COUNT(*)
      FROM {TABLE_NAME}
      WHERE (
        REGEXP_INSTR(LOWER({TITLE_COL}), %s) > 0 OR
        REGEXP_INSTR(LOWER({DESC_COL}),  %s) > 0
      )
      AND NOT (
        REGEXP_INSTR(LOWER({TITLE_COL}), %s) > 0 OR
        REGEXP_INSTR(LOWER({DESC_COL}),  %s) > 0
      )
    """
    with conn.cursor() as cur:
        cur.execute(sql, (rx_non, rx_non, rx_adb, rx_adb))
        return int(cur.fetchone()[0] or 0)

# Embedding columns (native VECTOR types)
EMB_COL_TITLE = "EMBEDDING_TITLE"
EMB_COL_DESC  = "EMBEDDING_DESCRIPTION"

# Models
EMBED_MODEL = "text-embedding-3-large"  # 3072 dims
CHAT_MODEL  = "gpt-4o-mini"

if not OPENAI_API_KEY:
    st.error("Missing OPENAI_API_KEY in environment/.env")
    st.stop()

client = OpenAI(api_key=OPENAI_API_KEY)

# ---------- Unified Persona ----------
ASSISTANT_PERSONA = """
You are a research assistant for job listings focused on Creative Professionals (“CPros”).
Be concise, factual, and context-aware; use plain language.

Creative Professionals (or CPros) are workers who complete tasks requiring imagination, originality,
and artistic skills. They produce content that is visually innovative, expressive, or aesthetically engaging, such as:
- Photo/image editing (e.g., cropping, filters, compositing)
- Video/audio editing (e.g., trimming/combining clips, effects, music production, motion design)
- Graphic design or digital drawing/painting (e.g., icons, logos, graphics, illustrations)
- Web/app design (e.g., site/app design, flows, wireframes)
- Layout design (e.g., flyers, posters, brochures, magazines)
- 3D design or immersive experiences (e.g., modeling, animation, 3D aggregation)

Examples: Graphic Designers, UX/UI Designers, Game Designers, Video Editors, Photographers, Animators,
Special Effects Artists. By contrast, Chief Executive Officers, Accountants, Truck Drivers, and house Painters
are not digital Creative Professionals.

Sub-groups and synonyms:
- “Photo Pros” → occupations heavy in photography and photo-editing (Photographers, Retouchers).
- “Video Pros” → occupations mainly using video/audio editing software (Video Editors, Sound Designers).
- “Design Pros” → occupations using design software (Illustrator, Canva, Figma, AutoCAD, UX/UI).

Normalization:
- Treat “Creative Professional(s)” and “CPro(s)” as synonyms.
- Treat “Photo Pros”, “Video Pros”, “Design Pros” and common variants as the above categories.
- Prefer TITLE when the user asks about roles; prefer DESCRIPTION when the user asks about skills/tools/tasks.
"""

# --- Brand guide injected into the intent system prompt
BRAND_GUIDE = (
    "Known apps and vendors (lowercase keys): "
    + json.dumps(VENDOR_KB, ensure_ascii=False)
    + "\nAdobe apps: "
    + json.dumps(sorted(ADOBE_SET), ensure_ascii=False)
    + "\nNon-Adobe apps: "
    + json.dumps(sorted(NON_ADOBE_SET), ensure_ascii=False)
)

INTENT_SYS = f"""
{ASSISTANT_PERSONA}

Use this brand guide when classifying queries about 'Adobe' vs 'non-Adobe':
{BRAND_GUIDE}

Task: Analyze the user's question about job listings and return STRICT JSON:
{{
  "intent": "TITLE" | "DESCRIPTION" | "BOTH",
  "detail": "titles" | "detailed",
  "keywords": [string, ...],
  "followup": "NONE" | "FILTER" | "EXPAND" | "RELATED",
  "cpro": boolean,
  "photo_pro": boolean,
  "video_pro": boolean,
  "design_pro": boolean
}}
Rules:
- Set boolean flags true if the query refers to those groups or synonyms.
- If role-centric → "TITLE"; if skill/tool/task-centric → "DESCRIPTION".
- Return ONLY JSON; no extra text.
"""

FOLLOWUP_SYS = f"""
{ASSISTANT_PERSONA}

Task: Given a user's query and top job listings (title + short snippet),
return STRICT JSON:
{{
  "summary": string,                 # 2–4 lines, factual, grounded in results
  "suggestions": [string,string,string]
}}
Style: concise, professional, grounded in evidence; return ONLY JSON.
"""

# ---------- Helpers ----------
def connect_snowflake() -> sf.SnowflakeConnection:
    return sf.connect(
        account=SF_ACCOUNT,
        user=SF_USER,
        password=SF_PAT,
        warehouse=SF_WAREHOUSE,
        database=SF_DATABASE,
        schema=SF_SCHEMA,
        role=SF_ROLE,
        client_session_keep_alive=True,
    )

def embed_texts(texts: List[str]) -> np.ndarray:
    if not texts:
        return np.zeros((0, 3072), dtype=np.float32)
    resp = client.embeddings.create(model=EMBED_MODEL, input=texts)
    vecs = [np.array(d.embedding, dtype=np.float32) for d in resp.data]
    return np.vstack(vecs)

def analyze_query(history: List[Dict[str,str]], user_query: str) -> Dict[str, Any]:
    msgs = [{"role":"system","content":INTENT_SYS}] + history + [{"role":"user","content":user_query}]
    res = client.chat.completions.create(
        model=CHAT_MODEL, messages=msgs, temperature=0.2, response_format={"type":"json_object"}
    )
    try:
        return json.loads(res.choices[0].message.content)
    except Exception:
        return {"intent":"BOTH","detail":"detailed","keywords":[],"followup":"NONE"}

def summarize_and_suggest(user_query: str, results: pd.DataFrame) -> Dict[str, Any]:
    sample = results[[TITLE_COL, DESC_COL]].head(6).to_dict(orient="records") if not results.empty else []
    msgs = [
        {"role":"system","content":FOLLOWUP_SYS},
        {"role":"user","content":json.dumps({"query": user_query, "results": sample})},
    ]
    res = client.chat.completions.create(
        model=CHAT_MODEL, messages=msgs, temperature=0.3, response_format={"type":"json_object"}
    )
    try:
        return json.loads(res.choices[0].message.content)
    except Exception:
        return {"summary": "", "suggestions": []}

def is_count_query(q: str) -> bool:
    COUNT_PATTERNS = (r"\bhow many\b", r"\bcount\b", r"\bnumber of\b")
    ql = q.lower()
    return any(re.search(p, ql) for p in COUNT_PATTERNS)

def primary_keyword(q: str, llm_keywords: List[str]) -> Optional[str]:
    return (llm_keywords[0] if llm_keywords else None)

# --- Heuristic CPro count (method #2) ---
def run_count_cpro_heuristic(conn) -> int:
    inc_terms = [
        "photoshop","lightroom","illustrator","indesign","figma","sketch","canva","autocad",
        "after effects","premiere","davinci","resolve","blender","maya","3d","animation",
        "ux","ui","graphic","designer","retoucher","photograph","video","motion","editor",
        "visual","art","layout","brand","creative","content creator"
    ]
    exc_terms = [
        "accountant","driver","warehouse","receptionist","cashier","cook","barista","janitor","ceo","chief executive"
    ]
    inc = " OR ".join([f"TITLE ILIKE '%{w}%'" for w in inc_terms] + [f"DESCRIPTION ILIKE '%{w}%'" for w in inc_terms])
    exc = " OR ".join([f"TITLE ILIKE '%{w}%'" for w in exc_terms] + [f"DESCRIPTION ILIKE '%{w}%'" for w in exc_terms])
    sql = f"SELECT COUNT(*) FROM {TABLE_NAME} WHERE ({inc}) AND NOT ({exc})"
    with conn.cursor() as cur:
        cur.execute(sql)
        row = cur.fetchone()
    return int(row[0] or 0)

def run_count_titles(conn, keyword: Optional[str], use_desc: bool) -> int:
    txt_col = DESC_COL if use_desc else TITLE_COL
    with conn.cursor() as cur:
        if keyword:
            safe = f"%{keyword}%"
            cur.execute(f"SELECT COUNT(DISTINCT {TITLE_COL}) FROM {TABLE_NAME} WHERE {txt_col} ILIKE %s", (safe,))
        else:
            cur.execute(f"SELECT COUNT(DISTINCT {TITLE_COL}) FROM {TABLE_NAME}")
        row = cur.fetchone()
    return int(row[0] or 0)

# --- Vector search (Snowflake native → fallback) ---
def _sql_try_vector_distance(cur, qvec: np.ndarray, top_k: int, use_desc: bool, ilike_kw: Optional[str]) -> Optional[pd.DataFrame]:
    emb_col = EMB_COL_DESC if use_desc else EMB_COL_TITLE
    qlist = ",".join(str(x) for x in qvec.astype(float).tolist())
    where_parts = []
    if ilike_kw:
        safe = ilike_kw.replace("'", "''")
        where_parts.append(f"({TITLE_COL} ILIKE '%{safe}%' OR {DESC_COL} ILIKE '%{safe}%')")
    where_sql = ("WHERE " + " AND ".join(where_parts)) if where_parts else ""
    sql = f"""
        SELECT {ID_COL}, {TITLE_COL}, {DESC_COL},
               VECTOR_DISTANCE({emb_col}, TO_VECTOR(ARRAY_CONSTRUCT({qlist})), 'COSINE') AS score
        FROM {TABLE_NAME}
        {where_sql}
        ORDER BY score ASC
        LIMIT {top_k}
    """
    try:
        cur.execute(sql); rows = cur.fetchall()
        cols = [c[0] for c in cur.description]
        return pd.DataFrame(rows, columns=cols)
    except Exception:
        return None

def _sql_try_cosine_similarity(cur, qvec: np.ndarray, top_k: int, use_desc: bool, ilike_kw: Optional[str]) -> Optional[pd.DataFrame]:
    emb_col = EMB_COL_DESC if use_desc else EMB_COL_TITLE
    qlist = ",".join(str(x) for x in qvec.astype(float).tolist())
    where_parts = []
    if ilike_kw:
        safe = ilike_kw.replace("'", "''")
        where_parts.append(f"({TITLE_COL} ILIKE '%{safe}%' OR {DESC_COL} ILIKE '%{safe}%')")
    where_sql = ("WHERE " + " AND ".join(where_parts)) if where_parts else ""
    sql = f"""
        SELECT {ID_COL}, {TITLE_COL}, {DESC_COL},
               VECTOR_COSINE_SIMILARITY({emb_col}, TO_VECTOR(ARRAY_CONSTRUCT({qlist}))) AS score
        FROM {TABLE_NAME}
        {where_sql}
        ORDER BY score DESC
        LIMIT {top_k}
    """
    try:
        cur.execute(sql); rows = cur.fetchall()
        cols = [c[0] for c in cur.description]
        return pd.DataFrame(rows, columns=cols)
    except Exception:
        return None

def _py_fallback_cosine(cur, qvec: np.ndarray, top_k: int, use_desc: bool, ilike_kw: Optional[str]) -> pd.DataFrame:
    emb_col = EMB_COL_DESC if use_desc else EMB_COL_TITLE
    where_parts = []
    if ilike_kw:
        safe = ilike_kw.replace("'", "''")
        where_parts.append(f"({TITLE_COL} ILIKE '%{safe}%' OR {DESC_COL} ILIKE '%{safe}%')")
    where_sql = ("WHERE " + " AND ".join(where_parts)) if where_parts else ""
    sql = f"""
        SELECT {ID_COL}, {TITLE_COL}, {DESC_COL}, {emb_col}
        FROM {TABLE_NAME}
        {where_sql}
        LIMIT {max(top_k*10, 500)}
    """
    cur.execute(sql); rows = cur.fetchall()
    cols = [c[0] for c in cur.description]
    df = pd.DataFrame(rows, columns=cols)
    if df.empty:
        return df.assign(score=np.nan)
    # ensure numpy arrays
    df[emb_col] = df[emb_col].apply(lambda v: np.array(list(v), dtype=np.float32))
    mat = np.vstack(df[emb_col].values)
    qn = qvec.reshape(1, -1); qn = qn / (np.linalg.norm(qn, axis=1, keepdims=True) + 1e-9)
    mn = mat / (np.linalg.norm(mat, axis=1, keepdims=True) + 1e-9)
    sims = (mn @ qn.T).ravel()
    df = df.drop(columns=[emb_col]); df["score"] = sims
    return df.sort_values("score", ascending=False).head(top_k).reset_index(drop=True)

def vector_search(conn, query_text: str, top_k: int, use_desc: bool, ilike_kw: Optional[str]) -> Tuple[pd.DataFrame, str]:
    qvec = embed_texts([query_text])[0]
    with conn.cursor() as cur:
        df = _sql_try_vector_distance(cur, qvec, top_k, use_desc, ilike_kw)
        if df is not None and not df.empty:
            return df, "VECTOR_DISTANCE(COSINE)"
        df = _sql_try_cosine_similarity(cur, qvec, top_k, use_desc, ilike_kw)
        if df is not None and not df.empty:
            return df, "VECTOR_COSINE_SIMILARITY"
        df = _py_fallback_cosine(cur, qvec, top_k, use_desc, ilike_kw)
        return df, "python_cosine_fallback"

# ---------- UI (sidebar/options) ----------
with st.sidebar:
    st.subheader("Options")

    focus = st.selectbox(
        "Focus",
        ["AUTO","TITLE","DESCRIPTION"],
        index=0,
        help="AUTO uses the intent planner; otherwise force TITLE or DESCRIPTION."
    )

    ilike_kw = st.text_input(
        "Optional keyword filter",
        help="Restrict results to rows whose TITLE or DESCRIPTION contains this word/phrase."
    )

    @st.cache_data(show_spinner=False, ttl=120)
    def get_max_k(focus_opt: str, ilike_opt: str) -> int:
        """
        Count rows that have embeddings available for the active focus,
        optionally applying an ILIKE filter. This anchors the Top-K slider.
        """
        emb_col = {
            "TITLE": "EMBEDDING_TITLE",
            "DESCRIPTION": "EMBEDDING_DESCRIPTION",
        }.get(focus_opt, None)  # AUTO -> None (count rows that have either)

        where = []
        params = []

        if ilike_opt:
            where.append("(TITLE ILIKE %s OR DESCRIPTION ILIKE %s)")
            like = f"%{ilike_opt}%"
            params.extend([like, like])

        if emb_col:
            where.append(f"{emb_col} IS NOT NULL")
        else:
            # AUTO: allow either column to be present
            where.append("(EMBEDDING_TITLE IS NOT NULL OR EMBEDDING_DESCRIPTION IS NOT NULL)")

        where_sql = "WHERE " + " AND ".join(where) if where else ""

        sql = f"SELECT COUNT(*) FROM {TABLE_NAME} {where_sql}"

        try:
            conn = connect_snowflake()
            with conn.cursor() as cur:
                cur.execute(sql, tuple(params) if params else None)
                n = cur.fetchone()[0] or 0
        finally:
            try: conn.close()
            except Exception: pass

        return max(int(n), 1)  # avoid 0 for slider bounds

    max_k = get_max_k(focus, ilike_kw)

    # Make slider max + default dynamic
    step = 10 if max_k > 100 else 1
    top_k = st.slider(
        "Top-K",
        min_value=1,
        max_value=max_k,
        value=max_k,        # default = max (your request)
        step=step,
        help="How many similar job postings to fetch per query (limited to rows with embeddings and optional filter)."
    )

    show_table = st.checkbox("Show results table", value=False)
    show_debug = st.checkbox("Show debug info (engine/focus)", value=False)
    if st.button("🧹 Clear conversation / context"):
        st.session_state.history = []
        st.session_state.last_results = None
        st.session_state.last_query = None
        st.session_state.last_intent = {}
        st.success("Conversation context cleared.")


# ---------- Conversation state ----------
if "history" not in st.session_state: st.session_state.history = []
if "last_results" not in st.session_state: st.session_state.last_results = None
if "last_query" not in st.session_state: st.session_state.last_query = None
if "last_intent" not in st.session_state: st.session_state.last_intent = {}

# Render conversation so far
for m in st.session_state.history:
    st.chat_message(m["role"]).markdown(m["content"])

prompt = st.chat_input("Ask about job titles, tools, or tasks…")
# ---------- Single render path ----------
if prompt:
    st.session_state.history.append({"role":"user","content":prompt})

    # Intent
    strategy = analyze_query(st.session_state.history[:-1], prompt)
    st.session_state.last_intent = strategy

    llm_intent = (strategy.get("intent") or "BOTH").upper()
    use_desc = {
        "TITLE": False,
        "DESCRIPTION": True,
        "BOTH": (focus == "DESCRIPTION"),
    }.get(llm_intent, (focus == "DESCRIPTION"))
    if focus == "TITLE": use_desc = False
    if focus == "DESCRIPTION": use_desc = True

    # Derived flags/keywords (compute BEFORE branching)
    kw = primary_keyword(prompt, strategy.get("keywords", []) or [])
    cpro_flag   = bool(strategy.get("cpro"))
    photo_flag  = bool(strategy.get("photo_pro"))
    video_flag  = bool(strategy.get("video_pro"))
    design_flag = bool(strategy.get("design_pro"))

    do_count = is_count_query(prompt)
    ql = prompt.lower()
    is_non_adobe_q = bool(re.search(r"\bnon[-\s]?adobe\b", ql))
    is_adobe_q     = ("adobe" in ql) and not is_non_adobe_q

    results: Optional[pd.DataFrame] = None
    mode_used = ""

    try:
        conn = connect_snowflake()

        if do_count:
            if is_non_adobe_q:
                n = count_non_adobe(conn)
                results = pd.DataFrame({"Metric":["Listings with non-Adobe apps"], "Value":[n]})
                mode_used = "COUNT non-Adobe (regex)"
            elif is_adobe_q:
                n = count_adobe(conn)
                results = pd.DataFrame({"Metric":["Listings with Adobe apps"], "Value":[n]})
                mode_used = "COUNT Adobe (regex)"
            elif cpro_flag:
                n = run_count_cpro_heuristic(conn)
                results = pd.DataFrame({"Metric":["CPro jobs (heuristic)"], "Value":[n]})
                mode_used = "COUNT heuristic (CPro)"
            else:
                n = run_count_titles(conn, keyword=kw or ilike_kw, use_desc=use_desc)
                results = pd.DataFrame({"Metric":["Distinct job titles"], "Value":[n]})
                mode_used = f"COUNT on {'DESCRIPTION' if use_desc else 'TITLE'} (ILIKE)"

        else:
            results, mode_used = vector_search(conn, prompt, top_k=top_k, use_desc=use_desc, ilike_kw=ilike_kw or kw)

    except Exception as e:
        st.error(f"Snowflake error: {e}")
        results = pd.DataFrame()
        mode_used = "error"
    finally:
        try: conn.close()
        except Exception: pass

    # Safety: ensure DataFrame
    if results is None or not isinstance(results, pd.DataFrame):
        results = st.session_state.get("last_results")
        if results is None or not isinstance(results, pd.DataFrame):
            results = pd.DataFrame()

    # Persist for follow-ups
    st.session_state.last_query = prompt
    st.session_state.last_results = results

    # ---- Render (once) ----
    if do_count:
        n = int(results.loc[0, "Value"]) if not results.empty else 0
        # Make the count message contextual
        if is_non_adobe_q:
            msg = f"I found **{n}** listings that include non-Adobe apps."
        elif is_adobe_q:
            msg = f"I found **{n}** listings that include Adobe apps."
        elif cpro_flag:
            msg = f"I found **{n}** listings likely targeting Creative Professionals."
        else:
            msg = f"I found **{n}** distinct job titles matching your request."
        st.chat_message("assistant").markdown(msg)
        st.session_state.history.append({"role":"assistant","content":msg})

        if show_table and not results.empty:
            st.dataframe(results, use_container_width=True, hide_index=True)
        if show_debug:
            st.caption(f"Debug: mode={mode_used}, focus={'DESCRIPTION' if use_desc else 'TITLE'}")
        st.stop()

    else:
        if results.empty:
            msg = "I couldn’t find any relevant job titles for that—try rephrasing or adding a keyword."
            st.chat_message("assistant").markdown(msg)
            st.session_state.history.append({"role":"assistant","content":msg})
            st.stop()

        example_titles = (
            results[TITLE_COL].dropna().astype(str).head(8).tolist()
            if TITLE_COL in results.columns else []
        )
        try:
            sug = summarize_and_suggest(prompt, results)
            summary = (sug.get("summary") or "").strip() or "Here are example job titles:"
        except Exception:
            summary = "Here are example job titles:"
        bullet_lines = "\n".join([f"- {t}" for t in example_titles])
        msg = f"{summary}\n\n**Examples:**\n{bullet_lines}" if example_titles else summary

        st.chat_message("assistant").markdown(msg)
        st.session_state.history.append({"role":"assistant","content":msg})

        if show_table:
            show_cols = [c for c in [TITLE_COL, DESC_COL, "score"] if c in results.columns]
            st.dataframe(results[show_cols], use_container_width=True, hide_index=True)
        if show_debug:
            st.caption(f"Debug: mode={mode_used}, focus={'DESCRIPTION' if use_desc else 'TITLE'}")
        st.stop()

with st.expander("ℹ️ Setup Notes"):
    st.markdown(
        """
        **Environment variables (.env)**  
        - `OPENAI_API_KEY`  
        - `SNOWFLAKE_ACCOUNT`, `SNOWFLAKE_USER`, `SNOWFLAKE_WAREHOUSE`, `SNOWFLAKE_DATABASE`, `SNOWFLAKE_SCHEMA`, `SNOWFLAKE_ROLE`, `SNOWFLAKE_PAT`  

        **Table & columns**  
        - Table: `JOB_LISTINGS_100.PUBLIC.SAMPLE_JOB_LISTINGS`  
        - Text: `TITLE`, `DESCRIPTION`  
        - Vectors: `EMBEDDING_TITLE`, `EMBEDDING_DESCRIPTION` (Snowflake VECTOR)

        **How it works**  
        - Conversational: history kept in `st.session_state.history`.  
        - Intent planner detects **count** queries; otherwise runs **semantic search**.  
        - Uses native Snowflake vector similarity when available; falls back to Python cosine.
        """
    )