src/streamlit_app.py

import streamlit as st

import csv
import pandas as pd
import logging
import sys
import os
import re
import shutil
from pathlib import Path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.agents.CleanReponseAgent import ai_format_response
from src.data.FinancialKBIngestor import FinancialKBIngestor



from src.core.FinAgentEngine import FinAgentEngine
from src.agents.ModerationAgent import ModerationAgent

from src.core.settings import get_settings
from src.core.errors import add_error
from src.core.logging_config import configure_logging

logger = logging.getLogger(__name__)
settings = get_settings()

configure_logging()

def reset_chroma_if_configured() -> None:
    """
    If FIN_ASSISTANT_RESET_CHROMA=1, wipe the local Chroma persist directory.

    This clears *both* the Knowledge Base vectors and the semantic cache because they
    share `src/data/.chroma`.
    """
    flag = str(os.getenv("FIN_ASSISTANT_RESET_CHROMA", "") or "").strip().lower()
    if flag not in {"1", "true", "yes", "y"}:
        return

    # Avoid repeated resets on Streamlit reruns within the same session/process.
    try:
        if st.session_state.get("_chroma_reset_done"):
            return
    except Exception:
        pass

    # If any Chroma-backed objects were cached via Streamlit, clear them first so we
    # don't keep stale in-memory clients pointing at a deleted directory.
    try:
        st.cache_resource.clear()
    except Exception:
        pass
    try:
        st.cache_data.clear()
    except Exception:
        pass

    chroma_dir = (Path(__file__).resolve().parents[0] / "data" / ".chroma").resolve()
    if chroma_dir.exists():
        try:
            shutil.rmtree(chroma_dir)
            logger.info("Reset Chroma persist directory: %s", str(chroma_dir))
        except Exception as e:
            logger.exception("Failed to reset Chroma persist directory: %s", e)
            return
    try:
        chroma_dir.mkdir(parents=True, exist_ok=True)
    except Exception:
        pass

    try:
        st.session_state["_chroma_reset_done"] = True
    except Exception:
        pass

def _strip_html(text: str) -> str:
    """
    Best-effort HTML removal for converting rendered assistant messages back into plain text.
    """
    if not text:
        return ""
    # Remove tags and collapse whitespace.
    no_tags = re.sub(r"<[^>]+>", " ", text)
    return re.sub(r"\s+", " ", no_tags).strip()


def build_llm_conversation_history(messages: list[dict]) -> list[dict]:
    """
    Convert UI chat messages into a clean, plain-text conversation history for the LLM.

    - Prefer `msg["raw"]` when present (authoritative plain text).
    - Fall back to stripping HTML from `msg["content"]` (for older sessions).
    """
    out: list[dict] = []
    for m in messages or []:
        if not isinstance(m, dict):
            continue
        role = m.get("role")
        content = m.get("raw")
        if content is None:
            content = _strip_html(str(m.get("content", "")))
        out.append({"role": role, "content": str(content or "")})
    return out

def get_router():
    return FinAgentEngine()

def get_moderation_agent():
    return ModerationAgent()

@st.cache_resource
def get_tax_kb():
    """
    Initialize and ingest the local JSON knowledge base(s) into Chroma only once.
    """
    logger.info("Initializing local JSON Knowledge Bases (Chroma ingest)...")
    ingestor = FinancialKBIngestor(
        json_file=[
            "tax_kb.json",
            "us_insurance_kb.json",
            "us_market_basics.json",
            "us_portfolio_kb.json",
            "us_stock_crypto_kb.json",
        ],
        use_qa_format=True
    )
    ingestor.run()
    logger.info("Local JSON Knowledge Bases initialized.")


BAD_LANGUAGE_RESPONSE = (
    "Please try with different query, like : What is the current price of Apple Stock ? "
    "Explain 401k in simple terms ?"
)

BAD_WORDS_CSV_PATH = str(
    (Path(__file__).resolve().parents[1] / "data" / "bad_words.csv").resolve()
)


def load_bad_words(csv_path: str) -> set[str]:
    """
    Load bad words from a CSV file.
    """
    bad_words: set[str] = set()
    path = Path(csv_path)
    if not path.exists():
        return bad_words

    with path.open("r", encoding="utf-8", newline="") as f:
        reader = csv.reader(f)
        for row in reader:
            if not row:
                continue
            w = (row[0] or "").strip().lower()
            if not w or w.startswith("#") or w == "word":
                continue
            bad_words.add(w)
    return bad_words

def _contains_bad_language(text: str) -> bool:
    bad_words = load_bad_words(BAD_WORDS_CSV_PATH)
    normalized = (text or "").lower()
    tokens = re.findall(r"[a-z]+", normalized)
    return any(t in bad_words for t in tokens)

def sanitize_user_input(user_input: str) -> tuple[bool, str]:
    """
    Returns (is_allowed, response_text_if_blocked).
    """
    if _contains_bad_language(user_input):
        return False, BAD_LANGUAGE_RESPONSE

    moderation_agent = get_moderation_agent()
    moderation_result = moderation_agent.classify(user_input)
    if moderation_result.get("flagged"):
        return False, BAD_LANGUAGE_RESPONSE
    return True, ""



reset_chroma_if_configured()
get_tax_kb()

# --- Page Config ---
st.set_page_config(page_title="AI Finance Assistant", page_icon="💰", layout="wide")

# --- Custom CSS ---
st.markdown("""
<style>
.main { background-color: #0E1117; }
.block-container { padding-top: 2rem; }
.chat-bubble { padding: 12px 16px; border-radius: 12px; margin-bottom: 10px; max-width: 80%; }
.user-bubble { background-color: #2563EB; color: white; margin-left: auto; }
.assistant-bubble { background-color: #1F2937; color: #E5E7EB; margin-right: auto; }
.header { font-size: 28px; font-weight: 700; margin-bottom: 0.5rem; }
.subheader { color: #9CA3AF; margin-bottom: 1.5rem; }
</style>
""", unsafe_allow_html=True)

# --- Session State ---
if "messages" not in st.session_state:
    st.session_state.messages = []

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

if "user_profile" not in st.session_state:
    st.session_state.user_profile = {
        "risk": settings.default_user_profile.risk,
        "experience": settings.default_user_profile.experience,
    }

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

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

# --- Sidebar ---
st.sidebar.title("⚙️ Settings")
st.sidebar.markdown("Customize your experience")

# Load S&P 500 data
def load_sp500_data():
    try:
        file_path = "data/sp500_top100.csv"
        if not os.path.isabs(file_path):
            base_dir = os.path.dirname(os.path.abspath(__file__))
            file_path = os.path.join(base_dir, file_path)
        df = pd.read_csv(file_path)
        return df
    except Exception as e:
        logger.error(f"Error loading S&P 500 data: {e}")
        return pd.DataFrame(columns=["Symbol", "Company Name"])

sp500_df = load_sp500_data()

risk = st.sidebar.selectbox(
    "Risk Profile",
    ["Low", "Moderate", "High"],
    index=["Low", "Moderate", "High"].index(st.session_state.user_profile["risk"].capitalize())
)
experience = st.sidebar.selectbox(
    "Experience",
    ["Beginner", "Intermediate", "Advanced"],
    index=["Beginner", "Intermediate", "Advanced"].index(st.session_state.user_profile["experience"].capitalize())
)

st.session_state.user_profile["risk"] = risk.lower()
st.session_state.user_profile["experience"] = experience.lower()

st.sidebar.markdown("---")
st.sidebar.subheader("👤 Profile")

# Dropdown for portfolio selection
selected_stocks = st.sidebar.multiselect(
    "Select Stocks for Portfolio",
    options=sp500_df["Symbol"].tolist(),
    default=[p["symbol"] for p in st.session_state.portfolio if "symbol" in p],
    max_selections=5
)

# New section for entering quantities
new_portfolio = []
if selected_stocks:
    st.sidebar.markdown("##### Set Quantities")
    for symbol in selected_stocks:
        # Try to find existing quantity in session state
        existing_qty = next((p["quantity"] for p in st.session_state.portfolio if p.get("symbol") == symbol), 1)
        qty = st.sidebar.number_input(f"Shares of {symbol}", min_value=1, value=existing_qty, key=f"qty_{symbol}")
        new_portfolio.append({"symbol": symbol, "quantity": qty})

# Update session state portfolio immediately
st.session_state.portfolio = new_portfolio

# --- Header ---
st.markdown('<div class="header">💰 AI Finance Assistant</div>', unsafe_allow_html=True)
st.markdown('<div class="subheader">Ask about investing, portfolio, market trends, or taxes. Stock prices are sourced from Alpha Vantage and Finnhub.</div>', unsafe_allow_html=True)

# Profile info display
p_risk = st.session_state.user_profile.get("risk")
p_exp = st.session_state.user_profile.get("experience")
p_portfolio_display = [f"{p['symbol']} ({p['quantity']})" for p in st.session_state.portfolio if "symbol" in p]

if p_risk and p_exp and p_portfolio_display:
    st.markdown(f"""
        <div style="margin-bottom: 1rem;">
            <span style="color: #2563EB; font-weight: bold;">Risk Profile: {p_risk.capitalize()}</span> | 
            <span style="color: #10B981; font-weight: bold;">Experience: {p_exp.capitalize()}</span> | 
            <span style="color: #EF4444; font-weight: bold;">Portfolio: {', '.join(p_portfolio_display)}</span>
        </div>
    """, unsafe_allow_html=True)
else:
    st.warning("Please select Risk Profile, Experience and Portfolio from settings")

# --- Status banner (best-effort telemetry) ---
last_errors = st.session_state.get("last_errors") or []
if isinstance(last_errors, list) and last_errors:
    st.warning("Some tools/data sources returned errors. Results may be incomplete.")
    with st.expander("Show details"):
        for e in last_errors[:12]:
            try:
                agent = e.get("agent", "unknown")
                msg = e.get("message", "")
                code = e.get("code", "")
                st.write(f"- {agent} [{code}]: {msg}")
            except Exception:
                continue
    if st.button("Dismiss errors"):
        st.session_state.last_errors = []
        st.rerun()

# --- Layout ---
main_col, history_col = st.columns([3, 1])

# --- Chat UI ---
with main_col:
    for msg in st.session_state.messages:
        if msg["role"] == "user":
            st.markdown(f'<div class="chat-bubble user-bubble">{msg["content"]}</div>', unsafe_allow_html=True)
        else:
            st.markdown(f'<div class="chat-bubble assistant-bubble">{msg["content"]}</div>', unsafe_allow_html=True)

    user_input = st.chat_input("Ask something like 'Explain ETFs'")

# --- Backend ---
def call_backend(query, state):
    router = get_router()
    try:
        result_state = router.invoke(query, initial_state=state)
        return result_state
    except Exception as e:
        result_state = {"response": f"An error occurred during processing: {e}", "errors": []}
        add_error(
            result_state,  # type: ignore[arg-type]
            code="backend_error",
            message=str(e),
            agent="streamlit_app",
        )
        return result_state

# --- Handle Input ---
if user_input:
    allowed, blocked_response = sanitize_user_input(user_input)
    if not allowed:
        st.session_state.messages.append({"role": "user", "content": user_input, "raw": user_input})
        st.markdown(f'<div class="chat-bubble user-bubble">{user_input}</div>', unsafe_allow_html=True)
        st.session_state.messages.append({"role": "assistant", "content": blocked_response, "raw": blocked_response})
        st.markdown(f'<div class="chat-bubble assistant-bubble">{blocked_response}</div>', unsafe_allow_html=True)
        st.rerun()

    st.session_state.messages.append({"role": "user", "content": user_input, "raw": user_input})

    with st.spinner("Working on it..."):
        graph_state = {
            "user_query": user_input,
            "conversation_history": build_llm_conversation_history(st.session_state.messages),
            "user_profile": st.session_state.user_profile,
            "portfolio": st.session_state.portfolio,
            "errors": [],
        }
        result_state = call_backend(user_input, graph_state)
        response_raw = result_state.get("response", "No response was generated by the agents.")
        st.session_state.last_errors = result_state.get("errors") or []

        # Update persistent portfolio from graph state (if present)
        if "portfolio" in result_state:
            st.session_state.portfolio = result_state["portfolio"]

        response_rendered = ai_format_response(response_raw)
        st.session_state.messages.append(
            {"role": "assistant", "content": response_rendered, "raw": response_raw}
        )

    if user_input not in [h["query"] for h in st.session_state.history]:
        item = {"query": user_input, "pinned": False}
        st.session_state.history.insert(0, item)
        st.session_state.history = st.session_state.history[:50]

    st.rerun()

# --- History Panel ---
with history_col:
    st.markdown("### 🕘 Search History")

    col1, col2 = st.columns(2)
    with col1:
        if st.button("🧹 Clear All"):
            st.session_state.history = []
            st.session_state.messages = []
            st.rerun()
    with col2:
        show_pinned = st.checkbox("⭐ Pinned only")

    history_sorted = sorted(st.session_state.history, key=lambda x: (not x.get("pinned", False)))

    if show_pinned:
        history_sorted = [h for h in history_sorted if h.get("pinned")]

    if history_sorted:
        for idx, item in enumerate(history_sorted):
            q = item["query"]
            pinned = item.get("pinned", False)

            c1, c2, c3 = st.columns([6,1,1])

            with c1:
                if st.button(q, key=f"run_{idx}"):
                    st.session_state.messages.append({"role": "user", "content": q, "raw": q})
                    with st.spinner("Working on it..."):
                        graph_state = {
                            "user_query": q,
                            "conversation_history": build_llm_conversation_history(st.session_state.messages),
                            "user_profile": st.session_state.user_profile,
                            "portfolio": st.session_state.portfolio,
                            "errors": [],
                        }
                        result_state = call_backend(q, graph_state)
                        st.session_state.last_errors = result_state.get("errors") or []
                        resp_raw = result_state.get("response", "No response was generated by the agents.")
                        resp_rendered = ai_format_response(resp_raw)
                        st.session_state.messages.append(
                            {"role": "assistant", "content": resp_rendered, "raw": resp_raw}
                        )
                    st.rerun()

            with c2:
                icon = "⭐" if pinned else "☆"
                if st.button(icon, key=f"pin_{idx}"):
                    for i, h in enumerate(st.session_state.history):
                        if h["query"] == q:
                            st.session_state.history[i]["pinned"] = not h.get("pinned", False)
                            break
                    st.rerun()

            with c3:
                if st.button("🗑️", key=f"del_{idx}"):
                    # remove from history
                    st.session_state.history = [h for h in st.session_state.history if h["query"] != q]

                    # remove related messages (user + assistant)
                    new_messages = []
                    skip_next = False
                    for i, msg in enumerate(st.session_state.messages):
                        if msg["role"] == "user" and msg["content"] == q:
                            skip_next = True
                            continue
                        if skip_next:
                            skip_next = False
                            continue
                        new_messages.append(msg)

                    st.session_state.messages = new_messages

                    st.rerun()
    else:
        st.write("No history yet")