# app.py
import os
import uuid
import gradio as gr

from langchain_core.chat_history import (
    InMemoryChatMessageHistory,
    BaseChatMessageHistory,
)
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.runnables import RunnableLambda, RunnableParallel
import os


from core import answer_as_table

# Ensure GOOGLE_API_KEY is set in environment before running.
# Example:
# os.environ["GOOGLE_API_KEY"] = "your-google-api-key"

# Prompt scaffolding (used only to satisfy history insertion)
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a helpful academic assistant that generates literature reviews.",
        ),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{question}"),
    ]
)


def identity(inputs: dict) -> dict:
    return {
        "question": (inputs.get("question") or "").strip(),
        "use_web": bool(inputs.get("use_web", False)),
        "region": (inputs.get("region") or "us-en"),
        "safesearch": (inputs.get("safesearch") or "moderate"),
        "timelimit": (inputs.get("timelimit") or None),
        "backend": (inputs.get("backend") or None),
        "max_results": int(inputs.get("max_results") or 20),
    }


id_runnable = RunnableLambda(identity)


def _orchestrate(inputs: dict) -> str:
    text = (inputs.get("question") or "").strip()
    use_web = bool(inputs.get("use_web", False))
    region = inputs.get("region") or "us-en"
    safesearch = inputs.get("safesearch") or "moderate"
    timelimit = inputs.get("timelimit") or None
    backend = inputs.get("backend") or None
    max_results = int(inputs.get("max_results") or 20)

    if not text:
        # Return a small info table only when user provided nothing
        return (
            "| Intent | Reply |\n"
            "|--------|-------|\n"
            "| Help | Please enter a research topic or a message. |\n"
        )

    # Route to web TABLE or plain chat text depending on use_web
    return answer_as_table(
        text,
        region=region,
        max_results=max_results,
        safesearch=safesearch,
        timelimit=timelimit,
        backend=backend,
        force_web=use_web,
    )


core_runnable = RunnableLambda(_orchestrate)

# Run prompt and identity in parallel, then pick the identity output to feed core.
# Prompt runs solely to let RunnableWithMessageHistory insert 'history'.
combined = (
    RunnableParallel(prompt=prompt, data=id_runnable).pick("data")
) | core_runnable

# Session-scoped history
_store: dict[str, BaseChatMessageHistory] = {}


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in _store:
        _store[session_id] = InMemoryChatMessageHistory()
    return _store[session_id]


with_history = RunnableWithMessageHistory(
    combined,
    get_session_history,
    input_messages_key="question",
    history_messages_key="history",
)  # requires config={"configurable": {"session_id": "<id>"}} on invoke


def respond(
    message,
    history,
    use_web,
    session_state,
    region,
    safesearch,
    timelimit,
    backend,
    max_results,
):
    """
    - message: dict or str (ChatInterface type='messages' passes a dict with 'text')
    - history: UI history (Gradio-managed; LangChain history is separate)
    - use_web: checkbox
    - session_state: gr.State carrying a stable session_id to isolate histories across users
    - region, safesearch, timelimit, backend, max_results: web search controls
    """
    text = (message.get("text") if isinstance(message, dict) else message) or ""
    text = text.strip()

    if not text:
        return (
            "| Intent | Reply |\n"
            "|--------|-------|\n"
            "| Help | Please enter a research topic or a message. |\n"
        ), session_state

    # Ensure a per-user session id for RunnableWithMessageHistory
    session_id = session_state.get("session_id")
    if not session_id:
        session_id = f"conv-{uuid.uuid4().hex}"
        session_state["session_id"] = session_id

    try:
        output = with_history.invoke(
            {
                "question": text,
                "use_web": bool(use_web),
                "region": (region or "us-en"),
                "safesearch": (safesearch or "moderate"),
                "timelimit": (timelimit or None),
                "backend": (backend or None),
                "max_results": int(max_results or 20),
            },
            config={"configurable": {"session_id": session_id}},
        )
        # output is either a Markdown TABLE (web) or plain chat text (no web)
        return output, session_state
    except Exception as e:
        return (
            f"| Intent | Reply |\n|--------|-------|\n| Error | {str(e)} |\n"
        ), session_state


with gr.Blocks(title="Literature Review Chat") as demo:
    gr.Markdown(
        "Enter a research topic to generate a Markdown literature review table (enable web), or chat for quick help (plain text)."
    )
    session_state = gr.State(
        {"session_id": None}
    )  # session-persistent state in the browser tab

    with gr.Row():
        use_web = gr.Checkbox(label="Use web search (academic sources)", value=True)
        region = gr.Dropdown(
            choices=["us-en", "wt-wt", "uk-en", "ca-en", "in-en", "de-de", "fr-fr"],
            value="us-en",
            label="Region",
        )
        safesearch = gr.Dropdown(
            choices=["on", "moderate", "off"], value="moderate", label="SafeSearch"
        )
        timelimit = gr.Dropdown(
            choices=[None, "d", "w", "m", "y"], value=None, label="Time limit"
        )
        backend = gr.Dropdown(
            choices=[None, "api", "html", "lite"], value=None, label="DDG backend"
        )
        max_results = gr.Slider(
            minimum=5, maximum=50, value=20, step=1, label="Max results"
        )

    chat = gr.ChatInterface(
        fn=respond,
        additional_inputs=[
            use_web,
            session_state,
            region,
            safesearch,
            timelimit,
            backend,
            max_results,
        ],
        additional_outputs=[session_state],
        type="messages",
        title="Literature Review Chat",
        description="Toggle the checkbox to search the web and produce a literature review table; otherwise, get a concise plain-text chat reply.",
        save_history=True,
    )

if __name__ == "__main__":
    demo.launch()