import os
import asyncio
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from tools import search_knowledge_base, add_to_knowledge_base

load_dotenv()

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 = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

agent = create_deep_agent(
    model=llm,
    tools=[search_knowledge_base, add_to_knowledge_base],
    backend=backend,
    system_prompt="You are a helpful knowledge assistant. Use the tools to search and add documents.",
)

async def interactive_loop():
    thread_id = "interactive-session"
    print("Welcome to RAG Agent. Commands: /add <title> <content>, /search <query>, /quit")
    while True:
        user_input = input(">> ")
        if user_input.strip() == "/quit":
            print("Goodbye.")
            break
        if user_input.startswith("/add"):
            try:
                _, title, content = user_input.split(" ", 2)
            except ValueError:
                print("Usage: /add <title> <content>")
                continue
            message = HumanMessage(content=f"Add document titled '{title}' with content: {content}")
        elif user_input.startswith("/search"):
            query = user_input[len("/search"):].strip()
            message = HumanMessage(content=f"Search knowledge base for: {query}")
        else:
            message = HumanMessage(content=user_input)
        result = await agent.ainvoke(
            {"messages": [message]},
            {"configurable": {"thread_id": thread_id}},
        )
        print(result["messages"][-1].content)

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