import streamlit as st
import sys
import os

# Ensure the app can find the local modules
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
if BASE_DIR not in sys.path:
    sys.path.insert(0, BASE_DIR)

from unified_legal_rag import ask

st.set_page_config(
    page_title="LegalAIapex",
    page_icon="⚖️",
    layout="wide"
)

st.title("⚖️ LegalAIapex: Unified Statute + Judgment RAG")
st.markdown("""
Welcome to the Unified Legal RAG system. This system dynamically routes your query to **Statutes** (IPC, BNS, CrPC, BNSS, IEA, BSA) and **Supreme Court Judgments**. It features a dual-layer verification system to ensure zero hallucinations.
""")

with st.sidebar:
    st.header("⚙️ Settings")
    intent_mode = st.selectbox(
        "Response Format",
        options=[
            "⚖️ General Legal Research (Default)",
            "✨ Auto-Detect Format (AI decides)",
            "📖 Analyze My Story (Legal Advice)",
            "📄 Brief Case/Statute Summary",
            "📚 In-Depth Case Study",
            "🔄 Compare Laws/Cases"
        ],
        index=0,
        help="Choose the structure of the answer. By default, it provides a clean Citation Table."
    )

INTENT_MAP = {
    "⚖️ General Legal Research (Default)": "LEGAL_RESEARCH",
    "✨ Auto-Detect Format (AI decides)": "AUTO",
    "📖 Analyze My Story (Legal Advice)": "STORY_EVALUATION",
    "📄 Brief Case/Statute Summary": "CASE_SUMMARY",
    "📚 In-Depth Case Study": "COMPREHENSIVE_CASE_STUDY",
    "🔄 Compare Laws/Cases": "CASE_COMPARISON"
}
force_intent = INTENT_MAP[intent_mode]

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    if message["role"] != "system":
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

# React to user input
if prompt := st.chat_input("Ask a legal question... (e.g. 'What is the law on murder under BNS 103?')"):
    # Display user message in chat message container
    st.chat_message("user").markdown(prompt)

    with st.chat_message("assistant"):
        status_text = st.empty()
        status_text.text("Retrieving legal evidence and generating response...")

        with st.spinner("Searching the legal database..."):
            try:
                # The prompt is added to history after, so we pass history up to this point
                result = ask(prompt, st.session_state.messages, force_intent=force_intent)

                # Clear status text
                status_text.empty()

                # Display the main answer (which uses the markdown table format you requested)
                st.markdown(result.get("answer", "No answer generated."))

                # Add an expander for the "behind-the-scenes" metadata
                with st.expander("🔍 Retrieval & Verification Details", expanded=True):
                    # Route, Intent & Speed
                    st.caption(
                        f"**Intent:** {result.get('intent', 'LEGAL_RESEARCH')} | **Route Taken:** {result.get('route')} | **Speed:** {result.get('elapsed_seconds')}s")

                    st.divider()

                    # Verification Block
                    v = result.get("verification", {})
                    if v:
                        grounded = v.get("grounded", None)
                        if grounded:
                            st.success("✅ **FULLY GROUNDED:** All citations perfectly match retrieved evidence.")
                        else:
                            st.warning(
                                "⚠️ **HALLUCINATED CITATIONS DETECTED:** The LLM cited sections/cases not found anywhere in the local database.")

                        if v.get("unverified_sections"):
                            st.write(f"🛑 **Hallucinated Sections (NOT in DB):** {', '.join(v['unverified_sections'])}")
                        if v.get("retrieval_miss_sections"):
                            st.info(
                                f"🔄 **Retrieval Miss (in DB, not retrieved this query):** {', '.join(v['retrieval_miss_sections'])} — These sections exist in the database but were not surfaced by the search engine for this query. The LLM cited them from its training memory.")
                        if v.get("citations_confirmed_via_live_lookup"):
                            st.info(
                                f"🟢 **Confirmed via Bharat-Courts (Live):** {', '.join(v['citations_confirmed_via_live_lookup'])}")
                        if v.get("citations_likely_fabricated"):
                            st.error(f"❌ **Likely Fabricated:** {', '.join(v['citations_likely_fabricated'])}")
                        if v.get("citations_could_not_verify"):
                            st.write(f"❓ **Could not verify:** {', '.join(v['citations_could_not_verify'])}")

                    st.divider()

                    # Evidence Block
                    st.write("**Top Evidence Used for Context:**")
                    evidence_list = result.get("evidence", [])
                    if evidence_list:
                        for ev in evidence_list:
                            type_icon = "📜" if ev['type'] == 'statute' else "🏛️"
                            st.write(f"{type_icon} `[{ev['type']}]` **Score:** {ev['score']:+.3f} — {ev['label']}")
                    else:
                        st.write("No evidence retrieved.")

            except Exception as e:
                status_text.empty()
                st.error(f"An error occurred: {str(e)}")
                result = {"answer": "Error generating response."}

    # Add user message and assistant message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    st.session_state.messages.append({"role": "assistant", "content": result.get("answer", "")})