import os
import streamlit as st
from rag import retrieve, ask_llm

# -----------------------------
# Streamlit Page Config
# -----------------------------
st.set_page_config(
    page_title="Only The Truth – Astrology Tutor",
    page_icon="🪐",
    layout="wide"
)

st.title("🪐 Only The Truth – Astrology Tutor")
st.caption("Vedic Astrology • PDF-Based Knowledge • RAG Powered")

# -----------------------------
# User Input
# -----------------------------
query = st.text_input(
    "Ask your astrology question:",
    placeholder="Example: Explain how twins are analyzed using D60 chart"
)

# -----------------------------
# Process Query
# -----------------------------
if query:
    with st.spinner("🔍 Searching ancient wisdom..."):
        contexts = retrieve(query)

    with st.spinner("🧠 Interpreting charts..."):
        answer = ask_llm(query, contexts)

    # -----------------------------
    # Answer Section
    # -----------------------------
    st.subheader("🪐 Answer")
    st.markdown(answer)

    # -----------------------------
    # Image Display Logic (STRICT)
    # -----------------------------
    st.subheader("📘 Reference Diagrams")

    IMAGE_DIR = "data/images"
    image_found = False

    query_lower = query.lower()
    KEYWORDS = [
        "chart", "diagram", "lagna", "horoscope",
        "d60", "sashtyamsa", "divisional", "birth chart"
    ]

    for c in contexts:
        text_lower = c["text"].lower()

        # Only show images if BOTH query & chunk indicate diagram relevance
        if not any(k in query_lower for k in KEYWORDS):
            continue

        if not any(k in text_lower for k in KEYWORDS):
            continue

        for img in c.get("images", []):
            img_path = os.path.join(IMAGE_DIR, img)

            if os.path.exists(img_path):
                st.image(
                    img_path,
                    caption=f"{c['source']} — page {c['page']}",
                    use_container_width=True
                )
                image_found = True

    if not image_found:
        st.info("ℹ No relevant diagrams found in the reference material.")

# -----------------------------
# Footer
# -----------------------------
st.markdown("---")
st.caption("Built with FAISS • SentenceTransformers • Groq • Streamlit")