"""
Deep agent that can search the web, create virtual files and finally dump them to disk.

The agent is built on top of LangChain 0.2+ and uses the "deep agents from scratch"
approach described in the course.  It is intentionally minimal but fully functional.
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import Any, Dict

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import Runnable
from langchain_ollama import ChatOllama
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma

# Local modules
from virtual_fs import virtual_fs

# ---------------------------------------------------------------------------
# 1. Tools
# ---------------------------------------------------------------------------

# 1.1 Web search tool – simple HTTP GET + title extraction
import requests
from bs4 import BeautifulSoup


def web_search(query: str) -> str:
    """Return the title and first paragraph of the first search result.

    This is a very small wrapper around a Google search.  For a production
    system you would use a real search API.
    """
    # Simple Bing search URL – works without API key for a few requests
    url = f"https://www.bing.com/search?q={requests.utils.quote(query)}"
    resp = requests.get(url, timeout=10)
    resp.raise_for_status()
    soup = BeautifulSoup(resp.text, "html.parser")
    results = soup.select("li.b_algo")
    if not results:
        return "No results found."
    first = results[0]
    title = first.select_one("h2").get_text(strip=True)
    snippet = first.select_one("p").get_text(strip=True)
    return f"Title: {title}\nSnippet: {snippet}"

# 1.2 Write file tool

def write_file_tool(path: str, content: str) -> str:
    virtual_fs.write(path, content)
    return f"File written to {path}."

# 1.3 Read file tool

def read_file_tool(path: str) -> str:
    try:
        return virtual_fs.read(path)
    except KeyError:
        return f"File {path} does not exist in virtual FS."

# 1.4 Dump virtual FS to disk

def dump_virtual_fs_tool(output_dir: str = "output") -> str:
    root = Path(output_dir)
    virtual_fs.dump_to_disk(root)
    return f"Virtual FS dumped to {root.resolve()}"

# ---------------------------------------------------------------------------
# 2. Agent definition – deep agent style
# ---------------------------------------------------------------------------

# 2.1 LLM
llm = ChatOllama(model="llama3.1", temperature=0.7)

# 2.2 Prompt template – instruct the agent how to use tools
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant that can search the web, write files, read files and dump virtual files to disk.")
])

# 2.3 Tool mapping
from langchain.tools import tool

# Wrap tools with langchain Tool objects
from langchain.tools import Tool

search_tool = Tool(
    name="WebSearch",
    func=web_search,
    description="Use this to search the web for information. Input should be a natural language query.",
)
write_tool = Tool(
    name="WriteFile",
    func=write_file_tool,
    description="Write content to a file in the virtual file system. Input: path and content.",
)
read_tool = Tool(
    name="ReadFile",
    func=read_file_tool,
    description="Read a file from the virtual file system. Input: path.",
)
dump_tool = Tool(
    name="DumpVirtualFS",
    func=dump_virtual_fs_tool,
    description="Dump all virtual files to the real file system. Input: output directory (optional).",
)

# 2.4 Agent chain – simple chain that lets the LLM decide which tool to call
from langchain.agents import AgentExecutor, ZeroShotAgent

# Define the tool names and descriptions for the prompt
tool_names = [search_tool.name, write_tool.name, read_tool.name, dump_tool.name]
tool_descriptions = [t.description for t in [search_tool, write_tool, read_tool, dump_tool]]

# Build the agent
agent = ZeroShotAgent.from_llm_and_tools(
    llm=llm,
    tools=[search_tool, write_tool, read_tool, dump_tool],
    prefix="You are a helpful assistant. Use the following tools when needed.",
    suffix="When you are finished, output the final answer.",
    tool_prompt="You can use the following tools: {tool_names}. {tool_descriptions}"
)

# Executor
executor = AgentExecutor.from_agent_and_tools(
    agent=agent,
    tools=[search_tool, write_tool, read_tool, dump_tool],
    verbose=True,
)

# ---------------------------------------------------------------------------
# 3. Demo / entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    print("Deep Agent Demo – type your question. Type 'exit' to quit.")
    while True:
        user_input = input("> ")
        if user_input.lower() in {"exit", "quit"}:
            print("Exiting…")
            break
        try:
            result = executor.invoke({"input": user_input})
            print("\nResult:\n", result)
        except Exception as e:
            print("Error:", e)