import os
import asyncio
from pathlib import Path
from typing import List

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

vector_store = Chroma(
    collection_name="knowledge",
    embedding_function=embeddings,
)

# ---------- Text Splitter ----------
splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
    separators=["\n\n", "\n", " "],
)

# ---------- RAG Tools ----------
@tool
def search_knowledge_base(query: str, max_results: int = 3) -> str:
    """
    Perform a semantic search in the knowledge base.
    Returns the concatenated contents of the most relevant documents.
    """
    docs: List[Document] = vector_store.similarity_search(query, k=max_results)
    if not docs:
        return "No relevant documents found."
    return "\n---\n".join(doc.page_content for doc in docs)


@tool
def add_to_knowledge_base(content: str, title: str = "document") -> str:
    """
    Add a new document to the knowledge base.
    The content will be split into chunks before indexing.
    """
    chunks = splitter.split_text(content)
    docs = [
        Document(page_content=chunk, metadata={"title": title, "chunk_index": i})
        for i, chunk in enumerate(chunks)
    ]
    vector_store.add_documents(docs)
    return f"Added {len(docs)} chunks from '{title}' to the knowledge base."


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

# ---------- Agent ----------
agent = create_deep_agent(
    model=llm,
    tools=[search_knowledge_base, add_to_knowledge_base],
    backend=backend,
    system_prompt=(
        "You are an AI assistant with access to a local knowledge base. "
        "When you need factual information, use the provided tools: "
        "`search_knowledge_base` to retrieve data and `add_to_knowledge_base` to store new documents. "
        "Always cite sources from the knowledge base in your answers."
    ),
)

# ---------- Helper Functions ----------
def load_documents_from_directory(directory: Path) -> None:
    """
    Recursively read .txt files from the given directory and add them to the knowledge base.
    """
    for file_path in directory.rglob("*.txt"):
        try:
            content = file_path.read_text(encoding="utf-8")
            title = file_path.stem
            add_to_knowledge_base(content, title)
            print(f"Loaded {file_path}")
        except Exception as e:
            print(f"Failed to load {file_path}: {e}")


async def chat_loop() -> None:
    """
    Simple CLI loop.
    Commands:
      /add <path>   - add a text file or all txt files in a directory
      /search <q>   - search the knowledge base
      /quit         - exit
    Anything else is sent to the agent as a user message.
    """
    thread_id = "cli-session"
    print("AI assistant ready. Type /quit to exit.")
    while True:
        user_input = input(">>> ").strip()
        if not user_input:
            continue
        if user_input.lower() == "/quit":
            print("Goodbye!")
            break
        if user_input.startswith("/add"):
            parts = user_input.split(maxsplit=1)
            if len(parts) != 2:
                print("Usage: /add <path>")
                continue
            path = Path(parts[1]).expanduser().resolve()
            if path.is_dir():
                load_documents_from_directory(path)
            elif path.is_file() and path.suffix.lower() == ".txt":
                content = path.read_text(encoding="utf-8")
                add_to_knowledge_base(content, path.stem)
                print(f"Added file {path}")
            else:
                print("Provide a .txt file or a directory containing .txt files.")
            continue
        if user_input.startswith("/search"):
            parts = user_input.split(maxsplit=1)
            if len(parts) != 2:
                print("Usage: /search <query>")
                continue
            query = parts[1]
            result = search_knowledge_base(query)
            print(f"Search results:\n{result}")
            continue

        # Normal conversation with the agent
        try:
            response = await agent.ainvoke(
                {"messages": [HumanMessage(content=user_input)]},
                {"configurable": {"thread_id": thread_id}},
            )
            answer = response["messages"][-1].content
            print(answer)
        except Exception as e:
            print(f"Agent error: {e}")


if __name__ == "__main__":
    # Optional: preload a default docs folder
    default_dir = Path("./docs")
    if default_dir.is_dir():
        load_documents_from_directory(default_dir)
    asyncio.run(chat_loop())