app.py

import os
import re
import streamlit as st
import google.generativeai as genai


# Backend / Model Logic

GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

if not GOOGLE_API_KEY:
    st.error("Google API Key not found. Add it in Hugging Face → Settings → Secrets.")
    st.stop()

genai.configure(api_key=GOOGLE_API_KEY)

_model = genai.GenerativeModel("gemini-2.5-flash-lite")


st.set_page_config(
    page_title="AI Lawyer — Indian Law Assistant",
    page_icon="⚖️",
    layout="wide",
    initial_sidebar_state="expanded",
)


# ╔══════════════════════════════════════════════════════════════════╗
# ║                       BACKEND / MODEL LOGIC                     ║
# ╚══════════════════════════════════════════════════════════════════╝

SYSTEM_PROMPT = """You are an expert AI lawyer specializing in Indian law.
You have deep knowledge of the Indian Penal Code (IPC), Code of Criminal Procedure (CrPC),
Constitution of India, Civil Procedure Code (CPC), and all major Indian statutes.

Rules:
- Answer clearly and precisely with relevant section numbers where applicable.
- If a question is outside Indian law, politely redirect.
- Structure long answers with headings when needed.
- Always end with a brief disclaimer: "This is for informational purposes only and not a substitute for professional legal advice."
"""

def markdown_to_html(text: str) -> str:
    """Convert markdown response to styled HTML for chat bubble rendering."""
    lines  = text.split("\n")
    html   = []
    in_ul  = False

    for line in lines:
        # ### Heading 3
        if line.startswith("### "):
            if in_ul: html.append("</ul>"); in_ul = False
            heading = re.sub(r"\*\*(.+?)\*\*", r"<strong>\1</strong>", line[4:].strip())
            html.append(f'<h3 class="md-h3">{heading}</h3>')

        # ## Heading 2
        elif line.startswith("## "):
            if in_ul: html.append("</ul>"); in_ul = False
            heading = re.sub(r"\*\*(.+?)\*\*", r"<strong>\1</strong>", line[3:].strip())
            html.append(f'<h2 class="md-h2">{heading}</h2>')

        # # Heading 1
        elif line.startswith("# "):
            if in_ul: html.append("</ul>"); in_ul = False
            heading = re.sub(r"\*\*(.+?)\*\*", r"<strong>\1</strong>", line[2:].strip())
            html.append(f'<h1 class="md-h1">{heading}</h1>')

        # Bullet points * or -
        elif re.match(r"^\s*[\*\-]\s+", line):
            if not in_ul: html.append("<ul class='md-ul'>"); in_ul = True
            item = re.sub(r"^\s*[\*\-]\s+", "", line)
            item = re.sub(r"\*\*(.+?)\*\*", r"<strong>\1</strong>", item)
            item = re.sub(r"\*(.+?)\*",     r"<em>\1</em>",          item)
            html.append(f"<li class='md-li'>{item}</li>")

        # Blank line
        elif line.strip() == "":
            if in_ul: html.append("</ul>"); in_ul = False
            html.append("<div class='md-gap'></div>")

        # Normal paragraph
        else:
            if in_ul: html.append("</ul>"); in_ul = False
            line = re.sub(r"\*\*(.+?)\*\*", r"<strong>\1</strong>", line)
            line = re.sub(r"\*(.+?)\*",     r"<em>\1</em>",          line)
            html.append(f"<p class='md-p'>{line}</p>")

    if in_ul:
        html.append("</ul>")

    return "".join(html)


def ask_llm(chat_history: list) -> str:
    """Send the full conversation history to Gemini and return the reply."""
    messages = []
    for msg in chat_history:
        role = "user" if msg["role"] == "user" else "model"
        messages.append({"role": role, "parts": [msg["content"]]})

    # Prepend system instruction as first user/model exchange
    full_messages = [
        {"role": "user",  "parts": [SYSTEM_PROMPT]},
        {"role": "model", "parts": ["Understood. I am your AI Lawyer specializing in Indian law. How can I assist you today?"]},
    ] + messages

    try:
        response = _model.generate_content(full_messages)
        return response.text
    except Exception as e:
        err = str(e)
        if "429" in err or "quota" in err.lower() or "ResourceExhausted" in err:
            return "⚠️ **API Quota Exceeded** — The Gemini free-tier limit has been reached. Please wait a few minutes and try again."
        return f"⚠️ **An error occurred:** {err}"


# ╔══════════════════════════════════════════════════════════════════╗
# ║                          UI / STREAMLIT                         ║
# ╚══════════════════════════════════════════════════════════════════╝

# ── Session State ─────────────────────────────────────────────────
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# ── CSS ───────────────────────────────────────────────────────────
st.markdown("""
<style>
@import url('https://fonts.googleapis.com/css2?family=Playfair+Display:wght@400;700;900&family=DM+Sans:ital,wght@0,300;0,400;0,500;1,300&display=swap');

:root {
    --gold:    #C9A84C;
    --gold-lt: #E8C97A;
    --dark:    #0D0D0D;
    --dark2:   #141414;
    --dark3:   #1C1C1C;
    --dark4:   #222222;
    --text:    #E8E0D0;
    --muted:   #8A8070;
    --border:  rgba(201,168,76,0.18);
}

html, body, [class*="css"] {
    font-family: 'DM Sans', sans-serif;
    background-color: var(--dark);
    color: var(--text);
}
.stApp { background-color: var(--dark); }
#MainMenu, footer, header { visibility: hidden; }
.block-container { padding: 1.8rem 2.5rem 2rem; max-width: 900px; }

/* ── Sidebar ── */
[data-testid="stSidebar"] {
    background: var(--dark2);
    border-right: 1px solid var(--border);
}
[data-testid="stSidebar"] .block-container { padding: 1.5rem 1.2rem; }

/* ── Hero ── */
.hero {
    background: linear-gradient(135deg, #0D0D0D 0%, #1A1205 55%, #0D0D0D 100%);
    border: 1px solid var(--border);
    border-radius: 16px;
    padding: 2rem 2.6rem 1.8rem;
    margin-bottom: 1.8rem;
    position: relative;
    overflow: hidden;
}
.hero::before {
    content: "\2696";
    position: absolute; right: -8px; top: -18px;
    font-size: 13rem; opacity: 0.04; line-height: 1; pointer-events: none;
}
.hero h1 {
    font-family: 'Playfair Display', serif;
    font-size: 2.4rem; font-weight: 900;
    background: linear-gradient(90deg, var(--gold), var(--gold-lt), var(--gold));
    -webkit-background-clip: text; -webkit-text-fill-color: transparent;
    margin: 0 0 0.3rem; line-height: 1.1;
}
.hero .sub  { color: var(--muted); font-size: 0.9rem; font-weight: 300; letter-spacing: 0.03em; margin: 0; }
.hero .pwr  { margin-top: 0.7rem; font-size: 0.68rem; letter-spacing: 0.08em; color: rgba(201,168,76,0.4); font-style: italic; }

/* ── Divider ── */
.gold-divider { border: none; border-top: 1px solid var(--border); margin: 1.3rem 0; }

/* ── Section label ── */
.section-label {
    font-size: 0.67rem; letter-spacing: 0.2em;
    text-transform: uppercase; color: var(--gold);
    font-weight: 500; margin-bottom: 0.5rem;
}

/* ── Chat window ── */
.chat-window {
    background: var(--dark2);
    border: 1px solid var(--border);
    border-radius: 14px;
    padding: 1.4rem 1.6rem;
    margin-bottom: 1.2rem;
    max-height: 520px;
    overflow-y: auto;
}
.chat-window::-webkit-scrollbar { width: 5px; }
.chat-window::-webkit-scrollbar-track { background: transparent; }
.chat-window::-webkit-scrollbar-thumb { background: var(--border); border-radius: 99px; }

/* ── Bubbles ── */
.user-bubble {
    display: flex; justify-content: flex-end; margin: 0.7rem 0;
}
.user-bubble .bub {
    background: linear-gradient(135deg, #A8832A, var(--gold));
    color: #0D0D0D;
    border-radius: 18px 18px 4px 18px;
    padding: 0.75rem 1.1rem;
    max-width: 78%;
    font-size: 0.92rem;
    line-height: 1.6;
    font-weight: 500;
}

.ai-bubble {
    display: flex; justify-content: flex-start;
    align-items: flex-start; gap: 0.7rem;
    margin: 0.7rem 0;
}
.ai-bubble .avatar {
    width: 32px; height: 32px; min-width: 32px;
    background: linear-gradient(135deg, #1A1205, #2A1F08);
    border: 1px solid var(--border);
    border-radius: 50%;
    display: flex; align-items: center; justify-content: center;
    font-size: 0.95rem;
}
.ai-bubble .bub {
    background: var(--dark3);
    border: 1px solid var(--border);
    border-radius: 4px 18px 18px 18px;
    padding: 0.75rem 1.3rem;
    max-width: 78%;
    font-size: 0.92rem;
    line-height: 1.75;
    color: var(--text);
}

/* ── Markdown inside AI bubble ── */
.md-h1 {
    font-family: 'Playfair Display', serif;
    font-size: 1.15rem; font-weight: 700;
    color: var(--gold);
    margin: 1rem 0 0.4rem 0;
    padding-bottom: 0.3rem;
    border-bottom: 1px solid var(--border);
}
.md-h2 {
    font-family: 'Playfair Display', serif;
    font-size: 1.05rem; font-weight: 700;
    color: var(--gold-lt);
    margin: 0.9rem 0 0.35rem 0;
}
.md-h3 {
    font-family: 'Playfair Display', serif;
    font-size: 0.95rem; font-weight: 700;
    color: var(--gold);
    margin: 0.8rem 0 0.25rem 0;
    letter-spacing: 0.02em;
}
.md-p {
    margin: 0.15rem 0;
    line-height: 1.75;
}
.md-gap { height: 0.4rem; }
.md-ul {
    margin: 0.3rem 0 0.3rem 1.3rem;
    padding: 0;
    list-style: none;
}
.md-li {
    position: relative;
    padding-left: 1rem;
    margin: 0.25rem 0;
    line-height: 1.7;
}
.md-li::before {
    content: "›";
    position: absolute; left: 0;
    color: var(--gold); font-weight: 700;
}

/* ── Typing indicator ── */
.typing { display: flex; gap: 4px; padding: 0.4rem 0; align-items: center; }
.typing span {
    width: 7px; height: 7px; background: var(--gold);
    border-radius: 50%; animation: bounce 1.2s infinite;
}
.typing span:nth-child(2) { animation-delay: 0.2s; }
.typing span:nth-child(3) { animation-delay: 0.4s; }
@keyframes bounce { 0%,80%,100%{transform:scale(0.7);opacity:0.4} 40%{transform:scale(1);opacity:1} }

/* ── Empty state ── */
.empty-chat {
    text-align: center; padding: 2.5rem 1rem;
    color: var(--muted); font-size: 0.88rem; line-height: 1.8;
}
.empty-chat .big { font-size: 2.8rem; margin-bottom: 0.6rem; }
.empty-chat .title {
    font-family: 'Playfair Display', serif;
    color: rgba(201,168,76,0.6); font-size: 1.05rem; margin-bottom: 0.4rem;
}

/* ── Inputs ── */
.stTextInput > div > div > input {
    background: var(--dark3) !important; border: 1px solid var(--border) !important;
    border-radius: 12px !important; color: var(--text) !important;
    font-family: 'DM Sans', sans-serif !important; font-size: 0.95rem !important;
    padding: 0.65rem 1rem !important;
}
.stTextInput > div > div > input:focus {
    border-color: var(--gold) !important;
    box-shadow: 0 0 0 2px rgba(201,168,76,0.14) !important;
}

/* ── Buttons ── */
.stButton > button {
    background: linear-gradient(135deg, var(--gold), #A8832A) !important;
    color: #0D0D0D !important; border: none !important;
    border-radius: 10px !important; font-family: 'DM Sans', sans-serif !important;
    font-weight: 600 !important; font-size: 0.86rem !important;
    padding: 0.48rem 1.4rem !important; letter-spacing: 0.02em !important;
    transition: all 0.2s !important;
}
.stButton > button:hover {
    transform: translateY(-1px) !important;
    box-shadow: 0 5px 18px rgba(201,168,76,0.28) !important;
}

/* ── Sidebar elements ── */
.sidebar-heading {
    font-family: 'Playfair Display', serif;
    font-size: 1.15rem; color: var(--gold); margin-bottom: 0.15rem;
}
.sb-credits {
    margin-top: 2rem; padding-top: 1rem;
    border-top: 1px solid rgba(201,168,76,0.12);
    font-size: 0.68rem; color: rgba(138,128,112,0.6);
    line-height: 1.75; text-align: center;
}
.sb-credits .sb-name {
    color: rgba(201,168,76,0.75);
    font-family: 'Playfair Display', serif;
    font-size: 0.82rem;
}

/* ── Footer ── */
.footer {
    margin-top: 2.5rem; padding: 1.2rem 0 0.4rem;
    border-top: 1px solid var(--border); text-align: center;
}
.footer .name {
    font-family: 'Playfair Display', serif;
    font-size: 0.95rem; color: var(--gold); letter-spacing: 0.04em;
}
.footer .org  { font-size: 0.74rem; color: var(--muted); margin-top: 0.2rem; letter-spacing: 0.06em; }
.footer .copy { font-size: 0.62rem; color: rgba(138,128,112,0.4); margin-top: 0.5rem; letter-spacing: 0.1em; text-transform: uppercase; }
</style>
""", unsafe_allow_html=True)


# ── SIDEBAR ───────────────────────────────────────────────────────
with st.sidebar:
    st.markdown('<p class="sidebar-heading">⚖️ AI Lawyer</p>', unsafe_allow_html=True)
    st.markdown('<p style="color:#8A8070;font-size:0.76rem;margin-top:-6px;">Indian Law Assistant</p>', unsafe_allow_html=True)
    st.markdown('<hr class="gold-divider">', unsafe_allow_html=True)

    st.markdown('<p class="section-label">About</p>', unsafe_allow_html=True)
    st.markdown("""
    <div style="background:#1C1C1C;border:1px solid rgba(201,168,76,0.15);border-radius:10px;padding:1rem 1.1rem;font-size:0.82rem;color:#8A8070;line-height:1.7;">
        Ask any question related to <span style="color:#C9A84C;">Indian law</span> — IPC, CrPC, Constitution, CPC, and more.<br><br>
        The AI will provide precise answers with relevant <span style="color:#C9A84C;">section references</span>.
    </div>
    """, unsafe_allow_html=True)

    st.markdown('<hr class="gold-divider">', unsafe_allow_html=True)
    st.markdown('<p class="section-label">Topics Covered</p>', unsafe_allow_html=True)

    topics = ["⚖️ Indian Penal Code", "📜 Constitution of India", "🔍 CrPC", "📋 Civil Procedure Code", "🏛️ Contract Act", "👨‍👩‍👧 Family Law", "🏘️ Property Law", "💼 Labour Law"]
    for t in topics:
        st.markdown(f'<div style="font-size:0.8rem;color:#8A8070;padding:3px 0;">{t}</div>', unsafe_allow_html=True)

    st.markdown('<hr class="gold-divider">', unsafe_allow_html=True)

    if st.session_state.chat_history:
        msg_count = len([m for m in st.session_state.chat_history if m["role"] == "user"])
        st.markdown(f'<div style="font-size:0.75rem;color:#8A8070;">💬 <span style="color:#C9A84C;">{msg_count}</span> question{"s" if msg_count!=1 else ""} asked</div>', unsafe_allow_html=True)
        st.markdown("")
        if st.button("🗑 Clear Chat", use_container_width=True):
            st.session_state.chat_history = []
            st.rerun()

    st.markdown("""
    <div class="sb-credits">
        Designed &amp; Developed by<br>
        <span class="sb-name">Yedeedya Injeti</span><br>
        Innomatics Research Labs
    </div>
    """, unsafe_allow_html=True)


# ── HERO ──────────────────────────────────────────────────────────
st.markdown("""
<div class="hero">
    <h1>AI Lawyer</h1>
    <p class="sub">Your intelligent legal assistant for Indian law &nbsp;&bull;&nbsp; IPC &bull; CrPC &bull; Constitution &amp; more</p>
    <p class="pwr">Powered by RAG + LLM</p>
</div>
""", unsafe_allow_html=True)


# ── QUICK SUGGESTIONS (shown only when chat is empty) ─────────────
suggestions = [
    "Punishment for murder under IPC?",
    "What is Article 21 of the Constitution?",
    "Define culpable homicide.",
    "What are bailable offences?",
    "Right to self-defence in India?",
    "What is the limitation period for filing a civil suit?",
]

if not st.session_state.chat_history:
    st.markdown('<p class="section-label">Suggested Questions</p>', unsafe_allow_html=True)
    col1, col2, col3 = st.columns(3)
    chosen = None
    for i, sug in enumerate(suggestions):
        col = [col1, col2, col3][i % 3]
        with col:
            if st.button(sug, key=f"sug_{i}", use_container_width=True):
                chosen = sug
    st.markdown("")


# ── CHAT WINDOW ───────────────────────────────────────────────────
if not st.session_state.chat_history:
    st.markdown("""
    <div class="chat-window">
        <div class="empty-chat">
            <div class="big">⚖️</div>
            <div class="title">How can I assist you today?</div>
            Ask me anything about Indian law — I'll provide accurate answers<br>with relevant section references and case principles.
        </div>
    </div>
    """, unsafe_allow_html=True)
else:
    bubbles_html = '<div class="chat-window">'
    for msg in st.session_state.chat_history:
        if msg["role"] == "user":
            bubbles_html += f"""
            <div class="user-bubble">
                <div class="bub">{msg["content"]}</div>
            </div>"""
        else:
            # ✅ Use markdown_to_html instead of plain .replace("\n", "<br>")
            content = markdown_to_html(msg["content"])
            bubbles_html += f"""
            <div class="ai-bubble">
                <div class="avatar">⚖️</div>
                <div class="bub">{content}</div>
            </div>"""
    bubbles_html += "</div>"
    st.markdown(bubbles_html, unsafe_allow_html=True)


# ── INPUT BAR ─────────────────────────────────────────────────────
col_inp, col_send = st.columns([6, 1])
with col_inp:
    question = st.text_input(
        "input", label_visibility="collapsed",
        placeholder="Ask your legal question…",
        key="chat_input",
    )
with col_send:
    send_btn = st.button("Send ⚖️", use_container_width=True)


# ── PROCESS ───────────────────────────────────────────────────────
final_q = None
if "chosen" in dir() and chosen:
    final_q = chosen
elif send_btn and question.strip():
    final_q = question.strip()

if final_q:
    st.session_state.chat_history.append({"role": "user", "content": final_q})

    with st.spinner(""):
        answer = ask_llm(st.session_state.chat_history)

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


# ── FOOTER ────────────────────────────────────────────────────────
st.markdown("""
<div class="footer">
    <div class="name">&#10022; Designed &amp; Developed by &nbsp;Yedeedya Injeti&nbsp; &#10022;</div>
    <div class="org">Innomatics Research Labs</div>
    <div class="copy">&copy; 2026 &nbsp;&middot;&nbsp; AI Lawyer &nbsp;&middot;&nbsp; Indian Law Assistant</div>
</div>
""", unsafe_allow_html=True)