src/streamlit_app.py

"""
Handbook Generator — Streamlit UI with ADK agent.
All RAG operations are synchronous — no event loop issues.
"""
import asyncio
import time
from pathlib import Path

import streamlit as st
from google.adk.sessions import InMemorySessionService
from google.adk.runners import Runner
from google.genai import types

from config import OPENAI_API_KEY, UPLOADS_DIR, _DATA_ROOT
from rag import index_pdf, reset_index
from handbook_generator import build_handbook
from agent import root_agent

HANDBOOK_EXPORT_PATH = _DATA_ROOT / "handbook_export.md"

# ────────────────────────────────────────────────
APP_NAME = "handbook_app"
session_service = InMemorySessionService()
runner = Runner(
    agent=root_agent,
    app_name=APP_NAME,
    session_service=session_service,
)

# ────────────────────────────────────────────────
st.set_page_config(page_title="Handbook Generator", page_icon="📖", layout="wide")

st.markdown("""
    <style>
        .stChatMessage { margin-bottom: 1.1rem !important; border-radius: 16px !important; }
        .stChatInput > div > div { border-radius: 24px !important; padding: 0.5rem 1rem; }
    </style>
""", unsafe_allow_html=True)


def ensure_api_key():
    if not OPENAI_API_KEY:
        st.error("OPENAI_API_KEY is not set. Create a .env with OPENAI_API_KEY=sk-...")
        return False
    return True


if not ensure_api_key():
    st.stop()

st.title("📖 Handbook Generator")
st.caption("Upload PDFs → Chat with ADK agent (RAG) → Generate 20k-word handbook")

for key, value in {"messages": [], "user_id": "default_user"}.items():
    if key not in st.session_state:
        st.session_state[key] = value

tab1, tab2, tab3 = st.tabs(["Upload PDFs", "Chat", "Generate Handbook"])

# ── Tab 1: Upload (synchronous — no async needed) ────────────────
def extract_text_from_pdf(file_bytes):
    pdf_reader = PyPDF2.PdfReader(BytesIO(file_bytes))
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text() or ""
    return text
with tab1:
    st.subheader("Upload and index PDFs")
    files = st.file_uploader("Choose PDF files", type=["pdf"], accept_multiple_files=True)
    if st.button("Index PDFs"):
        if not files:
            st.warning("Select at least one PDF.")
        else:
            reset_index()
            results = []
            for f in files:
                dest = UPLOADS_DIR / f.name
                dest.write_bytes(f.getvalue())
                try:
                    n = index_pdf(dest, source_name=f.name)
                    results.append(f"✅ {f.name}: {n} chunks indexed")
                except Exception as e:
                    results.append(f"❌ {f.name}: Error — {e}")
            st.success("\n".join(results))

# ── Tab 2: Chat ───────────────────────────────────────────────────
with tab2:
    st.subheader("Chat (ADK agent + RAG tool)")

    for msg in st.session_state.messages:
        role = "user" if msg["role"] == "user" else "assistant"
        avatar = "👤" if role == "user" else "🤖"
        with st.chat_message(role, avatar=avatar):
            st.markdown(msg["content"])

    user_input = st.chat_input("Ask about your uploaded documents...")

    if user_input:
        st.session_state.messages.append({"role": "user", "content": user_input})
        with st.chat_message("user", avatar="👤"):
            st.markdown(user_input)

        with st.chat_message("assistant", avatar="🤖"):
            placeholder = st.empty()
            placeholder.markdown("▋ Thinking…")

        user_id = st.session_state.user_id
        session_id = f"{user_id}_session"

        # ADK agent is async, run it properly
        async def run_agent():
            try:
                session = await session_service.get_session(
                    app_name=APP_NAME, user_id=user_id, session_id=session_id,
                )
                if not session:
                    await session_service.create_session(
                        app_name=APP_NAME, user_id=user_id, session_id=session_id,
                    )

                user_content = types.Content(
                    role="user",
                    parts=[types.Part.from_text(text=user_input)],
                )

                response_text = ""
                async for event in runner.run_async(
                    user_id=user_id, session_id=session_id, new_message=user_content,
                ):
                    if event.is_final_response():
                        if event.content and event.content.parts:
                            response_text = event.content.parts[0].text
                        break

                return response_text or "(No response generated)"
            except Exception as exc:
                return f"**Error occurred:** {str(exc)}"

        try:
            response = asyncio.run(run_agent())
        except RuntimeError:
            # Fallback if event loop already running
            import concurrent.futures
            with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
                response = pool.submit(asyncio.run, run_agent()).result(timeout=120)

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

# ── Tab 3: Handbook (synchronous) ─────────────────────────────────
with tab3:
    st.subheader("Generate 20k-word handbook")
    topic = st.text_input(
        "Handbook topic", placeholder="e.g. Retrieval-Augmented Generation",
    )
    if st.button("Generate handbook"):
        if not (topic and topic.strip()):
            st.warning("Enter a topic.")
        else:
            status_placeholder = st.empty()
            progress_msgs: list[str] = []

            def on_progress(msg):
                progress_msgs.append(msg)
                status_placeholder.text("\n".join(progress_msgs))

            try:
                full_md = build_handbook(topic.strip(), on_progress=on_progress)
                status_placeholder.success("Generation complete.")
                st.markdown(full_md)
                HANDBOOK_EXPORT_PATH.write_text(full_md, encoding="utf-8")
                st.download_button(
                    "Download as Markdown", data=full_md,
                    file_name="handbook.md", mime="text/markdown",
                )
            except Exception as e:
                status_placeholder.error(str(e))