"""Deep Agent implementation based on LangGraph and LangChain.

The module exposes two public functions:

* ``create_agent`` – returns a :class:`langgraph.graph.StateGraph` that can be
  executed.
* ``create_agent_executor`` – returns an :class:`langchain.agents.AgentExecutor`
  that can be used directly.

Both functions lazily import the heavy LangChain/LangGraph dependencies so
that the module can be imported even if those packages are not installed.
"""

from __future__ import annotations

from typing import Any

# Lightweight imports – can be imported eagerly
from virtual_fs import VirtualFileSystem
from tools import WebSearch, CreateVirtualFile, ExportVirtualFiles

# Shared virtual file system instance used by all tools
vfs = VirtualFileSystem()

# Define the tools
web_search_tool = WebSearch()
create_file_tool = CreateVirtualFile(vfs)
export_files_tool = ExportVirtualFiles(vfs)

# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def create_agent() -> Any:
    """Return a LangGraph graph.

    The function performs a lazy import of :mod:`langgraph` and related
    dependencies.  If the imports fail, a clear ``ImportError`` is raised.
    """
    try:
        from langgraph.graph import StateGraph
        from langgraph.prebuilt import create_react_agent as create_agent
        from langchain_ollama import ChatOllama
        from langchain_core.prompts import ChatPromptTemplate
        from langchain_core.output_parsers import StrOutputParser
    except Exception as exc:  # pragma: no cover
        raise ImportError(
            "Failed to import LangGraph/LangChain dependencies. Ensure that the\n"
            "required packages are installed and the environment is correctly\n"
            "configured."
        ) from exc

    # LLM and prompt
    llm = ChatOllama(model="llama3.1")
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "You are a helpful assistant that can search the web, create virtual files, and export them to disk."),
            ("human", "{input}"),
        ]
    )
    parser = StrOutputParser()

    # Build a simple agent using LangGraph prebuilt
    react_agent = create_agent(
        llm=llm,
        tools=[web_search_tool, create_file_tool, export_files_tool],
        prompt=prompt,
        output_parser=parser,
    )

    # Create the graph
    graph = StateGraph()
    graph.add_node("react_agent", react_agent)
    graph.set_entry_point("react_agent")
    graph.set_finish_point("react_agent")
    return graph


def create_agent_executor() -> Any:
    """Return an :class:`langchain.agents.AgentExecutor`.

    The function lazily imports the required classes.  It is useful for
    quick experimentation and for environments where the full graph is
    unnecessary.
    """
    try:
        from langchain.agents import AgentExecutor
        from langgraph.prebuilt import create_react_agent as create_agent
        from langchain_ollama import ChatOllama
        from langchain_core.prompts import ChatPromptTemplate
        from langchain_core.output_parsers import StrOutputParser
    except Exception as exc:  # pragma: no cover
        raise ImportError(
            "Failed to import LangChain dependencies for the AgentExecutor."
        ) from exc

    llm = ChatOllama(model="llama3.1")
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "You are a helpful assistant that can search the web, create virtual files, and export them to disk."),
            ("human", "{input}"),
        ]
    )
    parser = StrOutputParser()

    # Build the agent
    react_agent = create_agent(
        llm=llm,
        tools=[web_search_tool, create_file_tool, export_files_tool],
        prompt=prompt,
        output_parser=parser,
    )

    return AgentExecutor.from_agent_and_tools(
        agent=react_agent,
        tools=[web_search_tool, create_file_tool, export_files_tool],
        verbose=True,
    )

# Expose public names for tests
__all__ = ["create_agent", "create_agent_executor", "vfs"]