import os
import asyncio
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
import requests
from bs4 import BeautifulSoup

# LLM
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

# Backend: workspace + output
workspace_dir = "./workspace"
output_dir = "./output"

backend = CompositeBackend([
    LocalShellBackend(workspace_dir=workspace_dir),
    FilesystemBackend(root_dir=output_dir),
])

# Tools

@tool
def search_web(query: str) -> str:
    """Search the web using DuckDuckGo and return the first result title."""
    url = f"https://duckduckgo.com/html/?q={query}"
    headers = {"User-Agent": "Mozilla/5.0"}
    resp = requests.get(url, headers=headers, timeout=10)
    resp.raise_for_status()
    soup = BeautifulSoup(resp.text, "html.parser")
    result = soup.select_one(".result__a")
    return result.get_text() if result else "No results found"

@tool
def fetch_page(url: str) -> str:
    """Fetch the content of a web page."""
    headers = {"User-Agent": "Mozilla/5.0"}
    resp = requests.get(url, headers=headers, timeout=10)
    resp.raise_for_status()
    return resp.text[:2000]  # return first 2000 chars

@tool
def write_file(file_path: str, content: str) -> str:
    """Write content to a file in the output directory."""
    backend.write_file(file_path, content)
    return f"File written: {file_path}"

@tool
def read_file(file_path: str) -> str:
    """Read content from a file in the output directory."""
    return backend.read_file(file_path)

@tool
def list_files() -> str:
    """List files in the output directory."""
    files = []
    for root, dirs, filenames in os.walk(output_dir):
        for f in filenames:
            files.append(os.path.relpath(os.path.join(root, f), output_dir))
    return "\\n".join(files) if files else "No files found"

# Agent
agent = create_deep_agent(
    model=llm,
    tools=[search_web, fetch_page, write_file, read_file, list_files],
    backend=backend,
    system_prompt=(
        "You are a research assistant. You can search the web, fetch pages, "
        "create and read files in the output directory. "
        "When you finish your task, list all files you created."
    ),
)

async def main():
    user_prompt = (
        "Please research 'Python async programming', write a brief report "
        "in report.txt explaining key concepts, and create a README.md "
        "summarizing the steps taken. Then list the files you created."
    )
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=user_prompt)]},
        {"configurable": {"thread_id": "research-session"}},
    )
    # Print final assistant message
    print("\\n=== Final Assistant Message ===")
    print(result["messages"][-1].content)

if __name__ == "__main__":
    asyncio.run(main())