import os
import asyncio
from pathlib import Path

from dotenv import load_dotenv
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults

from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from langchain_core.messages import HumanMessage

load_dotenv()

# ---------- 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,
)

# ---------- Vector Store ----------
embeddings = OpenAIEmbeddings(
    model="text-embedding-3-small",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
)

def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma:
    return Chroma(
        collection_name="knowledge",
        embedding_function=embeddings,
        persist_directory=persist_directory,
    )

vectorstore = create_vectorstore()

def load_documents(directory: str, vectorstore: Chroma) -> None:
    path = Path(directory)
    if not path.is_dir():
        return
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    docs = []
    for file_path in path.rglob("*"):
        if file_path.suffix.lower() not in {".txt", ".md"}:
            continue
        text = file_path.read_text(encoding="utf-8")
        chunks = splitter.split_text(text)
        for i, chunk in enumerate(chunks):
            metadata = {"source": str(file_path), "chunk_index": i}
            docs.append(Document(page_content=chunk, metadata=metadata))
    if docs:
        vectorstore.add_documents(docs)

# Load initial documents (run once or each start)
load_documents("./documents", vectorstore)

# ---------- Tools ----------
@tool
def search_local_kb(query: str, top_k: int = 3) -> str:
    """
    Perform a semantic search in the local Chroma knowledge base.
    Returns the concatenated contents of the most relevant documents.
    """
    docs = vectorstore.similarity_search(query, k=top_k)
    if not docs:
        return "No relevant information found in the local knowledge base."
    return "\n---\n".join(doc.page_content for doc in docs)

tavily_tool = TavilySearchResults(max_results=5)

@tool
def web_search(query: str) -> str:
    """
    Search the web using Tavily and return a short summary of the top results.
    """
    results = tavily_tool.run(query)
    if not results:
        return "No web results found."
    # results is a list of dicts; extract title and url
    lines = []
    for r in results:
        title = r.get("title", "No title")
        url = r.get("url", "")
        lines.append(f"{title}: {url}")
    return "\n".join(lines)

# ---------- Backend ----------
backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

# ---------- Agent ----------
system_prompt = (
    "You are an AI assistant that can answer questions using two sources: "
    "a local knowledge base (accessed via the tool `search_local_kb`) and the web (accessed via the tool `web_search`). "
    "Decide which tool to use based on the user query. "
    "When you use a tool, include the source name in your final answer: "
    "`Source: chromadb` for local knowledge, `Source: tavily` for web results. "
    "If both sources are needed, combine them and list both sources."
)

agent = create_deep_agent(
    model=llm,
    tools=[search_local_kb, web_search],
    backend=backend,
    system_prompt=system_prompt,
)

# ---------- CLI ----------
async def chat_loop() -> None:
    print("AI Assistant (type 'exit' to quit)")
    thread_id = "session-1"
    while True:
        user_input = input("\nYou: ").strip()
        if user_input.lower() in {"exit", "quit"}:
            print("Goodbye!")
            break

        result = await agent.ainvoke(
            {"messages": [HumanMessage(content=user_input)]},
            {"configurable": {"thread_id": thread_id}},
        )
        answer = result["messages"][-1].content
        print(f"\nAssistant: {answer}")

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