src/streamlit_app.py

import os
import streamlit as st
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_groq import ChatGroq

st.set_page_config(
    page_title="Scheme Saathi — Indian Govt Schemes",
    page_icon="🏛️",
    layout="centered"
)

# ── Custom CSS ────────────────────────────────────────────────────────────────
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=Sora:wght@400;600;700&family=Inter:wght@400;500&display=swap');

/* ── Base ── */
html, body, [class*="css"] {
    font-family: 'Inter', sans-serif;
    background-color: #0F1117;
    color: #E8E8E8;
}

/* ── Hide Streamlit chrome ── */
#MainMenu, footer, header {visibility: hidden;}
.block-container {padding: 2rem 1.5rem 4rem;}

/* ── Hero ── */
.hero {
    text-align: center;
    padding: 3rem 1rem 2rem;
}
.hero-badge {
    display: inline-block;
    background: linear-gradient(135deg, #FF6B35, #FF8C42);
    color: white;
    font-family: 'Sora', sans-serif;
    font-size: 0.7rem;
    font-weight: 700;
    letter-spacing: 0.15em;
    text-transform: uppercase;
    padding: 0.35rem 1rem;
    border-radius: 2rem;
    margin-bottom: 1.2rem;
}
.hero-title {
    font-family: 'Sora', sans-serif;
    font-size: clamp(2rem, 5vw, 3.2rem);
    font-weight: 700;
    line-height: 1.15;
    color: #FFFFFF;
    margin: 0 0 0.8rem;
}
.hero-title span {
    background: linear-gradient(135deg, #FF6B35, #FFB347);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
}
.hero-sub {
    font-size: 1rem;
    color: #8A8A9A;
    max-width: 480px;
    margin: 0 auto 2rem;
    line-height: 1.6;
}

/* ── Stats row ── */
.stats-row {
    display: flex;
    justify-content: center;
    gap: 2.5rem;
    margin-bottom: 2.5rem;
    flex-wrap: wrap;
}
.stat {
    text-align: center;
}
.stat-num {
    font-family: 'Sora', sans-serif;
    font-size: 1.6rem;
    font-weight: 700;
    color: #FF6B35;
    line-height: 1;
}
.stat-label {
    font-size: 0.72rem;
    color: #5A5A6A;
    text-transform: uppercase;
    letter-spacing: 0.08em;
    margin-top: 0.2rem;
}

/* ── Divider ── */
.divider {
    border: none;
    border-top: 1px solid #1E1E2E;
    margin: 0 0 2rem;
}

/* ── Suggestion chips ── */
.chips-label {
    font-size: 0.75rem;
    color: #5A5A6A;
    text-transform: uppercase;
    letter-spacing: 0.1em;
    margin-bottom: 0.6rem;
}

/* ── Chat messages ── */
.msg-user {
    background: #1A1A2E;
    border: 1px solid #2A2A3E;
    border-radius: 1rem 1rem 0.2rem 1rem;
    padding: 0.9rem 1.2rem;
    margin: 0.5rem 0 0.5rem 3rem;
    font-size: 0.95rem;
    line-height: 1.6;
    color: #E8E8E8;
}
.msg-bot {
    background: #141420;
    border: 1px solid #FF6B3520;
    border-left: 3px solid #FF6B35;
    border-radius: 0.2rem 1rem 1rem 1rem;
    padding: 0.9rem 1.2rem;
    margin: 0.5rem 3rem 0.5rem 0;
    font-size: 0.95rem;
    line-height: 1.7;
    color: #D8D8E8;
}
.msg-label {
    font-size: 0.65rem;
    font-weight: 600;
    letter-spacing: 0.1em;
    text-transform: uppercase;
    margin-bottom: 0.3rem;
}
.label-user { color: #5A5A7A; }
.label-bot  { color: #FF6B35; }

/* ── Input area ── */
.stTextInput > div > div > input {
    background: #1A1A2E !important;
    border: 1px solid #2A2A4E !important;
    border-radius: 0.75rem !important;
    color: #E8E8E8 !important;
    padding: 0.85rem 1.1rem !important;
    font-size: 0.95rem !important;
    font-family: 'Inter', sans-serif !important;
}
.stTextInput > div > div > input:focus {
    border-color: #FF6B35 !important;
    box-shadow: 0 0 0 2px #FF6B3530 !important;
}
.stTextInput > div > div > input::placeholder {
    color: #3A3A5A !important;
}

/* ── Button ── */
.stButton > button {
    background: linear-gradient(135deg, #FF6B35, #FF8C42) !important;
    color: white !important;
    border: none !important;
    border-radius: 0.75rem !important;
    padding: 0.7rem 1.8rem !important;
    font-family: 'Sora', sans-serif !important;
    font-weight: 600 !important;
    font-size: 0.9rem !important;
    cursor: pointer !important;
    transition: opacity 0.2s !important;
    width: 100% !important;
}
.stButton > button:hover {
    opacity: 0.88 !important;
}

/* ── Source tag ── */
.source-tag {
    display: inline-block;
    background: #FF6B3515;
    border: 1px solid #FF6B3530;
    color: #FF8C42;
    font-size: 0.7rem;
    padding: 0.2rem 0.6rem;
    border-radius: 0.3rem;
    margin-top: 0.6rem;
}

/* ── Empty state ── */
.empty-state {
    text-align: center;
    padding: 3rem 1rem;
    color: #3A3A5A;
}
.empty-icon { font-size: 3rem; margin-bottom: 0.5rem; }
.empty-text { font-size: 0.9rem; }
</style>
""", unsafe_allow_html=True)


# ── Load resources ────────────────────────────────────────────────────────────
@st.cache_resource
def load_resources():
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    db = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
    llm = ChatGroq(
        api_key=os.environ["GROQ_API_KEY"],
        model="llama-3.3-70b-versatile",
        temperature=0.0
    )
    return db, llm

db, llm = load_resources()

# ── Session state ─────────────────────────────────────────────────────────────
if "messages" not in st.session_state:
    st.session_state.messages = []


# ── Hero section ──────────────────────────────────────────────────────────────
st.markdown("""
<div class="hero">
    <div class="hero-badge">🇮🇳 Powered by myScheme.gov.in</div>
    <h1 class="hero-title">Your <span>Scheme Saathi</span></h1>
    <p class="hero-sub">Ask anything about Indian government schemes — eligibility, benefits, documents, and how to apply.</p>
    <div class="stats-row">
        <div class="stat">
            <div class="stat-num">4418</div>
            <div class="stat-label">Schemes indexed</div>
        </div>
        <div class="stat">
            <div class="stat-num">RAG</div>
            <div class="stat-label">AI powered</div>
        </div>
        <div class="stat">
            <div class="stat-num">Free</div>
            <div class="stat-label">No login needed</div>
        </div>
    </div>
</div>
<hr class="divider">
""", unsafe_allow_html=True)


# ── Suggestion chips ──────────────────────────────────────────────────────────
suggestions = [
    "What schemes are available for women empowerment?",
    "Which government schemes help students get education loans?",
    "What schemes exist for senior citizens?",
    "Which schemes help farmers?",
]

st.markdown('<div class="chips-label">Try asking</div>', unsafe_allow_html=True)
cols = st.columns(2)
for i, s in enumerate(suggestions):
    with cols[i % 2]:
        if st.button(s, key=f"chip_{i}"):
            st.session_state.messages.append({"role": "user", "content": s})
            results_with_scores = db.similarity_search_with_relevance_scores(s, k=12)
            results = [doc for doc, score in results_with_scores if score > 0.3]
            if not results:
                results = [doc for doc, score in results_with_scores[:3]]
            context = "\n\n".join([r.page_content for r in results])
            prompt  = f"""You are Scheme Saathi, a friendly and knowledgeable assistant helping Indian citizens discover and understand government schemes.

COMMON ABBREVIATIONS (always expand these):
- PMUY = Pradhan Mantri Ujjwala Yojana
- PMAY = Pradhan Mantri Awas Yojana
- BBBP = Beti Bachao Beti Padhao
- PMJDY = Pradhan Mantri Jan Dhan Yojana
- PMKSY / PM Kisan = Pradhan Mantri Kisan Samman Nidhi
- APY = Atal Pension Yojana
- PMSBY = Pradhan Mantri Suraksha Bima Yojana
- PMJJBY = Pradhan Mantri Jeevan Jyoti Bima Yojana
- PMMVY = Pradhan Mantri Matru Vandana Yojana
- PMMY / MUDRA = Pradhan Mantri MUDRA Yojana
- MGNREGS / MGNREGA = Mahatma Gandhi National Rural Employment Guarantee Scheme
- NSAP = National Social Assistance Programme
- PMKVY = Pradhan Mantri Kaushal Vikas Yojana
- PMFBY = Pradhan Mantri Fasal Bima Yojana
- PMGSY = Pradhan Mantri Gram Sadak Yojana

INSTRUCTIONS FOR ANSWERING:

1. DOCUMENTS: When asked about required documents, papers, certificates, or what to bring —
   - Treat it the same as asking "how to apply" and search thoroughly
   - Check: "Required Documents" section, prerequisites in Application Process, KYC requirements, Aadhaar/bank details mentioned anywhere, documents listed in eligibility criteria
   - If documents are listed inside the Application Process section, present them clearly as "Required Documents"

2. ELIGIBILITY: When asked who can apply or who is eligible —
   - Look for: age limits, income limits, category (SC/ST/OBC/General/EWS), occupation, gender, marital status, state residency, BPL status
   - List all conditions clearly

3. BENEFITS: When asked about benefits, amount, subsidy, pension, or what you get —
   - Look for: financial amounts in Rs/₹, services provided, subsidies, in-kind benefits (cylinders, houses, equipment, food)
   - Always mention the exact amount if available

4. APPLICATION: When asked how to apply, where to apply, or application process —
   - Look for: Mode (Online/Offline/Both), official URL links, step-by-step process, office or agency to visit
   - Always include the URL if available in context

5. VAGUE QUERIES: If someone asks "schemes for farmers" or "schemes for women" or "schemes for students" —
   - List ALL relevant schemes found in context
   - Give scheme name + one-line description + key benefit for each

6. STATE SPECIFIC: If someone mentions a specific state —
   - Prioritize schemes from that state
   - Also mention relevant central government schemes available to all citizens of that state

7. PARTIAL NAMES: If the exact scheme name is not found but a related scheme exists —
   - Provide the related scheme's information
   - Mention: "This may be related to what you are looking for"

8. ABBREVIATION QUERIES: If someone asks only using an abbreviation like "What is PMUY?" or "Tell me about APY" —
   - Expand the abbreviation first
   - Then provide full details: description, eligibility, benefits, and how to apply

9. KYC / BANK QUERIES: If someone asks about KYC documents or bank account requirements —
   - Look for Aadhaar-based KYC options, offline KYC, virtual ID, bank account seeding
   - These are often in the Application Process section

STRICT RULES:
- Answer ONLY using the context provided below
- Never make up or assume information not present in the context
- If genuinely not found say: "I don't have specific information about this in my database. Please check myscheme.gov.in for the latest details."
- Keep answers clear, structured, and easy for a common citizen to understand
- Use bullet points or numbered lists when listing documents, steps, or multiple schemes
- Always be helpful and encouraging — many users may be first-time applicants

CONTEXT:
{context}

QUESTION: {s}

ANSWER:"""
            answer = llm.invoke(prompt).content
            st.session_state.messages.append({"role": "assistant", "content": answer})
            st.rerun()

st.markdown("<br>", unsafe_allow_html=True)


# ── Chat history ──────────────────────────────────────────────────────────────
if not st.session_state.messages:
    st.markdown("""
    <div class="empty-state">
        <div class="empty-icon">🏛️</div>
        <div class="empty-text">Ask a question above or type below to get started.</div>
    </div>
    """, unsafe_allow_html=True)
else:
    for msg in st.session_state.messages:
        if msg["role"] == "user":
            st.markdown(f"""
            <div class="msg-label label-user">You</div>
            <div class="msg-user">{msg["content"]}</div>
            """, unsafe_allow_html=True)
        else:
            st.markdown(f"""
            <div class="msg-label label-bot">Scheme Saathi</div>
            <div class="msg-bot">{msg["content"]}
                <div class="source-tag">📚 Source: myScheme.gov.in</div>
            </div>
            """, unsafe_allow_html=True)

st.markdown("<br>", unsafe_allow_html=True)

# ── Input ─────────────────────────────────────────────────────────────────────
with st.form("chat_form", clear_on_submit=True):
    col1, col2 = st.columns([5, 1])
    with col1:
        question = st.text_input(
            "",
            placeholder="Ask about any government scheme...",
            label_visibility="collapsed"
        )
    with col2:
        submitted = st.form_submit_button("Ask →")

if submitted and question.strip():
    st.session_state.messages.append({"role": "user", "content": question})

    results_with_scores = db.similarity_search_with_relevance_scores(question, k=12)
    results = [doc for doc, score in results_with_scores if score > 0.3]
    if not results:
        results = [doc for doc, score in results_with_scores[:3]]
    context = "\n\n".join([r.page_content for r in results])

    prompt  = f"""You are Scheme Saathi, a friendly and knowledgeable assistant helping Indian citizens discover and understand government schemes.

COMMON ABBREVIATIONS (always expand these):
- PMUY = Pradhan Mantri Ujjwala Yojana
- PMAY = Pradhan Mantri Awas Yojana
- BBBP = Beti Bachao Beti Padhao
- PMJDY = Pradhan Mantri Jan Dhan Yojana
- PMKSY / PM Kisan = Pradhan Mantri Kisan Samman Nidhi
- APY = Atal Pension Yojana
- PMSBY = Pradhan Mantri Suraksha Bima Yojana
- PMJJBY = Pradhan Mantri Jeevan Jyoti Bima Yojana
- PMMVY = Pradhan Mantri Matru Vandana Yojana
- PMMY / MUDRA = Pradhan Mantri MUDRA Yojana
- MGNREGS / MGNREGA = Mahatma Gandhi National Rural Employment Guarantee Scheme
- NSAP = National Social Assistance Programme
- PMKVY = Pradhan Mantri Kaushal Vikas Yojana
- PMFBY = Pradhan Mantri Fasal Bima Yojana
- PMGSY = Pradhan Mantri Gram Sadak Yojana

INSTRUCTIONS FOR ANSWERING:

1. DOCUMENTS: When asked about required documents, papers, certificates, or what to bring —
   - Treat it the same as asking "how to apply" and search thoroughly
   - Check: "Required Documents" section, prerequisites in Application Process, KYC requirements, Aadhaar/bank details mentioned anywhere, documents listed in eligibility criteria
   - If documents are listed inside the Application Process section, present them clearly as "Required Documents"

2. ELIGIBILITY: When asked who can apply or who is eligible —
   - Look for: age limits, income limits, category (SC/ST/OBC/General/EWS), occupation, gender, marital status, state residency, BPL status
   - List all conditions clearly

3. BENEFITS: When asked about benefits, amount, subsidy, pension, or what you get —
   - Look for: financial amounts in Rs/₹, services provided, subsidies, in-kind benefits (cylinders, houses, equipment, food)
   - Always mention the exact amount if available

4. APPLICATION: When asked how to apply, where to apply, or application process —
   - Look for: Mode (Online/Offline/Both), official URL links, step-by-step process, office or agency to visit
   - Always include the URL if available in context

5. VAGUE QUERIES: If someone asks "schemes for farmers" or "schemes for women" or "schemes for students" —
   - List ALL relevant schemes found in context
   - Give scheme name + one-line description + key benefit for each

6. STATE SPECIFIC: If someone mentions a specific state —
   - Prioritize schemes from that state
   - Also mention relevant central government schemes available to all citizens of that state

7. PARTIAL NAMES: If the exact scheme name is not found but a related scheme exists —
   - Provide the related scheme's information
   - Mention: "This may be related to what you are looking for"

8. ABBREVIATION QUERIES: If someone asks only using an abbreviation like "What is PMUY?" or "Tell me about APY" —
   - Expand the abbreviation first
   - Then provide full details: description, eligibility, benefits, and how to apply

9. KYC / BANK QUERIES: If someone asks about KYC documents or bank account requirements —
   - Look for Aadhaar-based KYC options, offline KYC, virtual ID, bank account seeding
   - These are often in the Application Process section

STRICT RULES:
- Answer ONLY using the context provided below
- Never make up or assume information not present in the context
- If genuinely not found say: "I don't have specific information about this in my database. Please check myscheme.gov.in for the latest details."
- Keep answers clear, structured, and easy for a common citizen to understand
- Use bullet points or numbered lists when listing documents, steps, or multiple schemes
- Always be helpful and encouraging — many users may be first-time applicants

CONTEXT:
{context}

QUESTION: {question}

ANSWER:"""

    with st.spinner("Finding the right scheme info..."):
        answer = llm.invoke(prompt).content

    st.session_state.messages.append({"role": "assistant", "content": answer})
    st.rerun()

# ── Footer ────────────────────────────────────────────────────────────────────
st.markdown("""
<br><br>
<div style="text-align:center; color:#2A2A3A; font-size:0.72rem;">
    Data sourced from myScheme.gov.in · Built with Streamlit + LangChain + Groq
</div>
""", unsafe_allow_html=True)