#!/usr/bin/env python3

"""
Agent for interacting with Zotero with human-in-the-loop features
"""

import logging
from typing import Any, Dict
import hydra

from langchain_core.language_models.chat_models import BaseChatModel
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import create_react_agent, ToolNode
from langgraph.checkpoint.memory import MemorySaver
from ..state.state_talk2scholars import Talk2Scholars
from ..tools.zotero.zotero_read import zotero_read
from ..tools.zotero.zotero_review import zotero_review
from ..tools.zotero.zotero_write import zotero_write
from ..tools.s2.display_results import display_results as s2_display
from ..tools.s2.query_results import query_results as s2_query_results
from ..tools.s2.retrieve_semantic_scholar_paper_id import (
    retrieve_semantic_scholar_paper_id,
)

# Initialize logger
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def get_app(uniq_id, llm_model: BaseChatModel):
    """
    Initializes and returns the LangGraph application for the Zotero agent.

    This function sets up the Zotero agent, which integrates various tools to search,
    retrieve, and display research papers from Zotero. The agent follows the ReAct
    pattern for structured interaction and includes human-in-the-loop features.

    Args:
        uniq_id (str): Unique identifier for the current conversation session.
        llm_model (BaseChatModel, optional): The language model to be used by the agent.
            Defaults to `ChatOpenAI(model="gpt-4o-mini", temperature=0)`.

    Returns:
        StateGraph: A compiled LangGraph application that enables the Zotero agent to
            process user queries and retrieve research papers.

    Example:
        >>> app = get_app("thread_123")
        >>> result = app.invoke(initial_state)
    """

    def agent_zotero_node(state: Talk2Scholars) -> Dict[str, Any]:
        """
        Processes the user query and retrieves relevant research papers from Zotero.

        This function calls the language model using the configured `ReAct` agent to
        analyze the state and generate an appropriate response. The function then
        returns control to the main supervisor.

        Args:
            state (Talk2Scholars): The current conversation state, including messages exchanged
                and any previously retrieved research papers.

        Returns:
            Dict[str, Any]: A dictionary containing the updated conversation state.

        Example:
            >>> result = agent_zotero_node(current_state)
            >>> papers = result.get("papers", [])
        """
        logger.log(
            logging.INFO, "Creating Agent_Zotero node with thread_id %s", uniq_id
        )
        result = model.invoke(state, {"configurable": {"thread_id": uniq_id}})

        return result

    # Load hydra configuration
    logger.log(logging.INFO, "Load Hydra configuration for Talk2Scholars Zotero agent.")
    with hydra.initialize(version_base=None, config_path="../configs"):
        cfg = hydra.compose(
            config_name="config",
            overrides=["agents/talk2scholars/zotero_agent=default"],
        )
        cfg = cfg.agents.talk2scholars.zotero_agent
        logger.log(logging.INFO, "Loaded configuration for Zotero agent")

    # Define the tools
    tools = ToolNode(
        [
            zotero_read,
            s2_display,
            s2_query_results,
            retrieve_semantic_scholar_paper_id,
            zotero_review,  # First review
            zotero_write,  # Then save with user confirmation
        ]
    )

    # Define the model
    logger.log(logging.INFO, "Using model %s", llm_model)

    # Create the agent
    model = create_react_agent(
        llm_model,
        tools=tools,
        state_schema=Talk2Scholars,
        prompt=cfg.zotero_agent,
        checkpointer=MemorySaver(),  # Required for interrupts to work
    )

    workflow = StateGraph(Talk2Scholars)
    workflow.add_node("agent_zotero", agent_zotero_node)
    workflow.add_edge(START, "agent_zotero")

    # Initialize memory to persist state between graph runs
    checkpointer = MemorySaver()

    # Compile the graph
    app = workflow.compile(checkpointer=checkpointer, name="agent_zotero")
    logger.log(
        logging.INFO,
        "Compiled the graph with thread_id %s and llm_model %s",
        uniq_id,
        llm_model,
    )

    return app