app.py

from __future__ import annotations

import os
import streamlit as st
from dotenv import load_dotenv

from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, BaseMessage

from orchestrator.settings import Settings
from orchestrator.factories import get_llm
from orchestrator.sql_agent import sql_answer
from orchestrator.graph_agent import graph_answer
from orchestrator.tools import run_tools_once
from orchestrator.graphs import build_router_graph, build_tools_agent_graph

load_dotenv()

st.set_page_config(page_title="Multi-Agent Orchestrator (LangGraph)", page_icon="🧭", layout="wide")


def _dict_messages_to_lc(messages: list[dict]) -> list[BaseMessage]:
    out: list[BaseMessage] = []
    for m in messages:
        role = m.get("role")
        content = m.get("content", "")
        if role == "user":
            out.append(HumanMessage(content=content))
        else:
            out.append(AIMessage(content=content))
    return out


def _extract_tool_names_from_messages(messages: list[BaseMessage]) -> list[str]:
    names: list[str] = []
    for m in messages:
        if isinstance(m, AIMessage):
            tool_calls = getattr(m, "tool_calls", None) or []
            for tc in tool_calls:
                if isinstance(tc, dict):
                    n = tc.get("name")
                else:
                    n = getattr(tc, "name", None)
                if n:
                    names.append(str(n))
    deduped: list[str] = []
    for n in names:
        if n not in deduped:
            deduped.append(n)
    return deduped


def _rewrite_followup_to_standalone(settings: Settings, chat_messages: list[dict], question: str) -> str:
    """
    Used in the *direct* SQL/Graph pages to make follow-ups work better.
    Router graph already does this internally.
    """
    user_count = sum(1 for m in chat_messages if m.get("role") == "user")
    if user_count <= 1:
        return question

    llm = get_llm(settings, temperature=0)

    # Build a short transcript
    recent = chat_messages[-12:]
    lines = []
    for m in recent:
        if m.get("role") == "user":
            lines.append(f"User: {m.get('content','')}")
        else:
            lines.append(f"Assistant: {m.get('content','')}")
    transcript = "\n".join(lines)

    prompt = (
        "Rewrite the user's latest question into a standalone question.\n"
        "Do NOT answer the question.\n\n"
        f"Conversation:\n{transcript}\n\n"
        f"Latest user question:\n{question}\n\n"
        "Standalone question:"
    )

    msg = llm.invoke(
        [
            SystemMessage(content="You rewrite follow-up questions into standalone questions."),
            HumanMessage(content=prompt),
        ]
    )
    rewritten = (msg.content or "").strip()
    return rewritten or question


# --- Sidebar ---
st.sidebar.title("🧭 Multi-Agent Orchestrator")

page = st.sidebar.radio(
    "Navigation",
    ["Router Chat", "SQL Agent", "Graph Agent", "Tools Agent", "Settings"],
    index=0,
)

# Runtime settings overrides (UI -> env-like)
st.sidebar.subheader("Model")
# llm_model = st.sidebar.text_input("LLM_MODEL (Groq)", value=os.getenv("LLM_MODEL", "llama-3.1-8b-instant"))
MODEL_OPTIONS = [
    "llama-3.1-8b-instant",
    "meta-llama/llama-4-maverick-17b-128e-instruct",
    "meta-llama/llama-4-scout-17b-16e-instruct",
    "moonshotai/kimi-k2-instruct-0905",
    "openai/gpt-oss-120b",
    "qwen/qwen3-32b",
]

default_model = os.getenv("LLM_MODEL", "meta-llama/llama-4-maverick-17b-128e-instruct")
if default_model not in MODEL_OPTIONS:
    MODEL_OPTIONS.insert(0, default_model)

llm_model = st.sidebar.selectbox("LLM_MODEL", MODEL_OPTIONS, index=MODEL_OPTIONS.index(default_model))

st.sidebar.subheader("SQL (SQLite)")
sqlite_path = st.sidebar.text_input("SQLITE_PATH", value=os.getenv("SQLITE_PATH", "student.db"))

st.sidebar.subheader("Neo4j (Graph DB)")
neo4j_uri = st.sidebar.text_input("NEO4J_URI", value=os.getenv("NEO4J_URI", ""))
neo4j_username = st.sidebar.text_input("NEO4J_USERNAME", value=os.getenv("NEO4J_USERNAME", ""))
neo4j_password = st.sidebar.text_input("NEO4J_PASSWORD", value=os.getenv("NEO4J_PASSWORD", ""), type="password")

st.sidebar.subheader("UI")
show_routing = st.sidebar.checkbox("Show routed agent", value=True)
show_tools_used = st.sidebar.checkbox("Show tools used", value=True)

settings = Settings(
    groq_api_key=os.getenv("GROQ_API_KEY", ""),
    llm_model=llm_model,
    sqlite_path=sqlite_path,
    neo4j_uri=neo4j_uri,
    neo4j_username=neo4j_username,
    neo4j_password=neo4j_password,
    wiki_doc_content_chars_max=int(os.getenv("WIKI_DOC_CHARS", "2000")),
    debug=os.getenv("DEBUG", "0") in ("1", "true", "True"),
)


@st.cache_resource
def _router_graph_cached(model: str):
    s = Settings(
        groq_api_key=settings.groq_api_key,
        llm_model=model,
        sqlite_path=settings.sqlite_path,
        neo4j_uri=settings.neo4j_uri,
        neo4j_username=settings.neo4j_username,
        neo4j_password=settings.neo4j_password,
        wiki_doc_content_chars_max=settings.wiki_doc_content_chars_max,
        debug=settings.debug,
    )
    return build_router_graph(s)


@st.cache_resource
def _tools_graph_cached(model: str):
    s = Settings(
        groq_api_key=settings.groq_api_key,
        llm_model=model,
        sqlite_path=settings.sqlite_path,
        neo4j_uri=settings.neo4j_uri,
        neo4j_username=settings.neo4j_username,
        neo4j_password=settings.neo4j_password,
        wiki_doc_content_chars_max=settings.wiki_doc_content_chars_max,
        debug=settings.debug,
    )
    return build_tools_agent_graph(s)


# --- Pages ---
if page == "Router Chat":
    st.title("🧭 Router Chat (LangGraph)")
    st.write("Multi-turn chat. The router chooses SQL / Graph / Tools / General automatically.")

    if "router_messages" not in st.session_state:
        st.session_state.router_messages = [
            {"role": "assistant", "content": "Hi! Ask a question — I will route it to the right agent."}
        ]

    c1, c2 = st.columns([1, 4])
    with c1:
        if st.button("Reset chat", key="reset_router"):
            st.session_state.router_messages = [
                {"role": "assistant", "content": "Chat reset. Ask a question!"}
            ]
            st.rerun()

    for m in st.session_state.router_messages:
        with st.chat_message(m["role"]):
            meta = m.get("meta") or {}
            if m["role"] == "assistant" and show_routing and meta.get("route"):
                st.caption(f"🧭 Routed to: `{meta['route']} agent`")
            if m["role"] == "assistant" and show_tools_used and meta.get("tools_used"):
                tools_line = ", ".join([f"`{t}`" for t in meta["tools_used"]])
                st.caption(f"🧰 Tools used: {tools_line}")
            st.write(m["content"])

    prompt = st.chat_input("Ask a question...", key="router_chat_input")
    if prompt:
        st.session_state.router_messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.write(prompt)

        try:
            with st.chat_message("assistant"):
                route_slot = st.empty()
                tools_slot = st.empty()
                answer_slot = st.empty()

                with st.spinner("Thinking..."):
                    graph = _router_graph_cached(settings.llm_model)
                    msgs = _dict_messages_to_lc(st.session_state.router_messages)

                    out = graph.invoke({"messages": msgs})
                    out_msgs = out.get("messages", []) or []

                    last_ai = next((mm for mm in reversed(out_msgs) if isinstance(mm, AIMessage)), None)
                    answer = last_ai.content if last_ai else "(no answer)"

                    dbg = out.get("debug", {}) or {}
                    route = out.get("route") or dbg.get("router_label") or dbg.get("routed_to") or "general"
                    tools_used = dbg.get("tools_used") or []

                # Update same bubble (no jump)
                if show_routing:
                    route_slot.caption(f"🧭 Routed to: `{route}` agent")
                if show_tools_used and tools_used:
                    tools_slot.caption("🧰 Tools used: " + ", ".join([f"`{t}`" for t in tools_used]))
                answer_slot.write(answer)

            # Append to chat history AFTER we have final answer
            st.session_state.router_messages.append(
                {"role": "assistant", "content": answer, "meta": {"route": route, "tools_used": tools_used}}
            )

            with st.expander("Debug (route + steps)"):
                st.write(out.get("debug", {}))
                st.write("Messages produced:", len(out_msgs))

        except Exception as e:
            st.error(str(e))

elif page == "SQL Agent":
    st.title("🧮 SQL Agent (Chat)")
    st.write("Multi-turn SQL chat. Good for follow-ups like “now filter by …”")

    # --- Intro: what the DB contains ---
    with st.expander("📌 What's in the SQL database?", expanded=False):
        st.markdown(
            """
            The database contains information about **students, courses, enrollments, and attendance**.

            - **students**: student_id, name, program, section, year
            - **courses**: course_id, course_code, course_name, department, credits
            - **enrollments**: student-course enrollment per semester with score and grade
            - **attendance**: per-class attendance for each student in each course and semester (present = 1/0)
            - **view**: student_performance (avg_score, num_A grades, num_courses per student per semester)

            Use this chat for analytics questions like rankings, averages, cohorts, and time/semester filtering.
            """
        )

    # --- Session init ---
    if "sql_messages" not in st.session_state:
        st.session_state.sql_messages = [
            {"role": "assistant", "content": "Ask a question about the student analytics database, or try an example below."}
        ]

    # --- Reset ---
    c1, _ = st.columns([1, 5])
    with c1:
        if st.button("Reset chat", key="reset_sql"):
            st.session_state.sql_messages = [{"role": "assistant", "content": "Chat reset. Ask a SQL question!"}]
            st.rerun()

    # --- Example queries (auto-run) ---
    st.subheader("⚡ Try an example")
    e1, e2, e3 = st.columns(3)

    if e1.button("🏆 Top students (2025-Fall)", use_container_width=True):
        st.session_state.sql_demo_query = (
            "Show the top 10 students by average score in semester 2025-Fall. "
            "Use the student_performance view. Return name, program, avg_score, num_courses, and num_A."
        )

    if e2.button("📉 Lowest scoring course (2025-Fall)", use_container_width=True):
        st.session_state.sql_demo_query = (
            "In 2025-Fall, which course has the lowest average score? "
            "Return course_code, course_name, department, and avg_score."
        )

    if e3.button("🧾 Attendance < 70% (2025-Fall)", use_container_width=True):
        st.session_state.sql_demo_query = (
            "For semester 2025-Fall, show students whose overall attendance is below 70%. "
            "Compute attendance_percent as 100 * AVG(present). "
            "Return student name, program, attendance_percent, and total_classes."
        )

    demo_query = st.session_state.pop("sql_demo_query", None)

    # --- Render chat history ---
    for m in st.session_state.sql_messages:
        st.chat_message(m["role"]).write(m["content"])

    # --- Input (manual OR demo) ---
    prompt = st.chat_input("Ask a SQL question...", key="sql_chat_input")
    user_query = prompt or demo_query

    if user_query:
        st.session_state.sql_messages.append({"role": "user", "content": user_query})
        st.chat_message("user").write(user_query)

        try:
            # Create assistant bubble immediately (prevents flicker)
            with st.chat_message("assistant"):
                answer_slot = st.empty()

                with st.spinner("Thinking..."):
                    standalone = _rewrite_followup_to_standalone(
                        settings,
                        st.session_state.sql_messages,
                        user_query,
                    )
                    out = sql_answer(settings, standalone)
                    answer = str(out.get("answer", ""))

                answer_slot.write(answer)

            # Append to history AFTER we have the final answer
            st.session_state.sql_messages.append({"role": "assistant", "content": answer})

            with st.expander("Debug"):
                st.write("Standalone question:", standalone)
                st.json(out)

        except Exception as e:
            st.error(str(e))

elif page == "Graph Agent":
    st.title("🕸️ Graph Agent (Chat)")
    st.write("Multi-turn Cypher/Q&A chat over Neo4j.")

    # --- Explain what graph contains ---
    with st.expander("📌 What's in the Neo4j database?", expanded=False):
        st.markdown(
            """
            **Theme:** Hollywood movies.

            **Nodes**
            - `Movie`: title, tagline, released (year)
            - `Person`: name, born (year)

            **Relationships**
            - `(:Person)-[:ACTED_IN]->(:Movie)`
            - `(:Person)-[:DIRECTED]->(:Movie)`
            - `(:Person)-[:PRODUCED]->(:Movie)`

            **Examples you can ask about**
            - Movies: “The Matrix”, “Top Gun”, “Jerry Maguire”
            - People: “Tom Cruise”, “Keanu Reeves”, “Tom Hanks”
            """
        )

    with st.expander("🧠 Why Neo4j (graph DB) vs Web Search?", expanded=False):
        st.markdown(
            """
            **Neo4j is best for relationship-heavy questions** where you want exact, structured answers:
            - “Who co-starred with Tom Cruise the most?”
            - “Find actors who worked with both Tom Cruise and Tom Hanks.”
            - “Show movies connected to *The Matrix* via shared actors.”

            **Web search is best for open-world facts** (news, definitions, anything outside your dataset).  
            So: Web search = broad; Neo4j = deep structured relationships inside your graph.
            """
        )

    # --- Session init ---
    if "graph_messages" not in st.session_state:
        st.session_state.graph_messages = [
            {"role": "assistant", "content": "Ask a question about the Neo4j movies graph, or try an example below."}
        ]

    # --- Reset button ---
    c1, _ = st.columns([1, 5])
    with c1:
        if st.button("Reset chat", key="reset_graph"):
            st.session_state.graph_messages = [
                {"role": "assistant", "content": "Chat reset. Ask a graph question!"}
            ]
            st.rerun()

    # --- Example queries (auto-run) ---
    st.subheader("⚡ Try an example")
    e1, e2, e3 = st.columns(3)

    if e1.button("🎭 Similar to The Matrix (shared actors)", use_container_width=True):
        st.session_state.graph_demo_query = (
            "Find movies that share at least 2 actors with The Matrix. "
            "Return the movie titles and how many actors are shared."
        )

    if e2.button("🧭 Shortest path: Tom Hanks ↔ Tom Cruise", use_container_width=True):
        st.session_state.graph_demo_query = (
            "Show the shortest connection between Tom Hanks and Tom Cruise."
        )

    if e3.button("🎬 Recommend like Cast Away", use_container_width=True):
        st.session_state.graph_demo_query = (
            "Recommend movies like Cast Away based on shared actor and director, and also name them."
        )

    demo_query = st.session_state.pop("graph_demo_query", None)

    # --- Render chat history ---
    for m in st.session_state.graph_messages:
        st.chat_message(m["role"]).write(m["content"])

    # --- Input (manual OR demo) ---
    prompt = st.chat_input("Ask a graph question...", key="graph_chat_input")
    user_query = prompt or demo_query

    if user_query:
        st.session_state.graph_messages.append({"role": "user", "content": user_query})
        st.chat_message("user").write(user_query)

        try:
            # Create assistant bubble immediately (prevents flicker)
            with st.chat_message("assistant"):
                answer_slot = st.empty()

                with st.spinner("Thinking..."):
                    standalone = _rewrite_followup_to_standalone(
                        settings,
                        st.session_state.graph_messages,
                        user_query,
                    )
                    out = graph_answer(settings, standalone)
                    answer = str(out.get("answer", ""))

                answer_slot.write(answer)

            # Append to history AFTER we have the final answer
            st.session_state.graph_messages.append({"role": "assistant", "content": answer})

            with st.expander("Debug (Cypher + results)"):
                st.write("Standalone question:", standalone)
                st.json(out.get("debug", {}))

        except Exception as e:
            st.error(str(e))

elif page == "Tools Agent":
    st.title("🧰 Tools Agent (Chat)")
    st.write("Tool-Assisted Research Chat (Web + Wikipedia + arXiv + Calculator).")

    if "tools_messages" not in st.session_state:
        st.session_state.tools_messages = [{"role": "assistant", "content": "Ask a question — I'll search web/Wikipedia/arXiv and use tools when needed."}]

    c1, _ = st.columns([1, 5])
    with c1:
        if st.button("Reset chat", key="reset_tools"):
            st.session_state.tools_messages = [{"role": "assistant", "content": "Chat reset. Ask a tools question!"}]
            st.rerun()

    for m in st.session_state.tools_messages:
        st.chat_message(m["role"]).write(m["content"])

    prompt = st.chat_input("Ask a tools question...", key="tools_chat_input")
    if prompt:
        st.session_state.tools_messages.append({"role": "user", "content": prompt})
        st.chat_message("user").write(prompt)

        try:
            with st.chat_message("assistant"):
                tools_slot = st.empty()
                answer_slot = st.empty()

                with st.spinner("Thinking..."):
                    tools_graph = _tools_graph_cached(settings.llm_model)
                    msgs = _dict_messages_to_lc(st.session_state.tools_messages)

                    out = tools_graph.invoke({"messages": msgs})
                    out_msgs = out.get("messages", []) or []

                    last_ai = next((mm for mm in reversed(out_msgs) if isinstance(mm, AIMessage)), None)
                    answer = last_ai.content if last_ai else "(no answer)"
                    tools_used = _extract_tool_names_from_messages(out_msgs)

                if show_tools_used and tools_used:
                    tools_slot.caption("🧰 Tools used: " + ", ".join([f"`{t}`" for t in tools_used]))
                answer_slot.write(answer)

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

            with st.expander("Debug (tool messages)"):
                st.write("Tools used:", tools_used)
                st.write("Messages produced:", len(out_msgs))

        except Exception as e:
            st.error(str(e))

    # Optional: keep your old "run once each" tester as a quick health check
    with st.expander("Quick tool health-check (run each tool once)"):
        q = st.text_input("Query for one-shot tools test", key="tools_q_once")
        if st.button("Run one-shot tools", type="secondary"):
            try:
                results = run_tools_once(
                    q,
                    wiki_chars=settings.wiki_doc_content_chars_max,
                )
                for r in results:
                    with st.expander(r.tool):
                        st.write(r.output)
            except Exception as e:
                st.error(str(e))

else:
    st.title("⚙️ Settings / Health Check")
    st.write("Use this page to confirm your keys and connections.")

    if not settings.groq_api_key:
        st.warning("GROQ_API_KEY is not set. Add it in your environment or .env.")
    else:
        st.success("GROQ_API_KEY is set.")

    st.write("**Current model:**", settings.llm_model)
    st.write("**SQLite path:**", settings.sqlite_path)

    if settings.neo4j_uri:
        st.write("**Neo4j URI:**", settings.neo4j_uri)
    else:
        st.info("Neo4j not configured yet (NEO4J_URI empty). Graph Agent will fail until set.")