"""
CLI for interacting with the RAG agent.

Commands:
  /add <file_path>   – add document to knowledge base
  /search <query>    – search knowledge base
  /quit              – exit

Optional argument: --load-dir DIR – load all .txt files from directory at start.
"""

import argparse
import sys
from pathlib import Path

from rag_agent import agent, add_to_knowledge_base, search_knowledge_base
from langchain_core.messages import HumanMessage


def load_directory(dir_path: str):
    """Load all .txt files from a directory into the knowledge base."""
    p = Path(dir_path)
    if not p.is_dir():
        print(f"{dir_path} is not a directory")
        return
    for txt_file in p.rglob("*.txt"):
        try:
            content = txt_file.read_text(encoding="utf-8")
            title = txt_file.stem
            add_to_knowledge_base(content=content, title=title)
            print(f"Loaded {txt_file}")
        except Exception as e:
            print(f"Failed to load {txt_file}: {e}")


def main():
    parser = argparse.ArgumentParser(description="RAG agent CLI")
    parser.add_argument("--load-dir", type=str, help="Directory with .txt files to preload")
    args = parser.parse_args()

    if args.load_dir:
        load_directory(args.load_dir)

    print("Enter commands (/add <file>, /search <query>, /quit). Press Ctrl+C to exit.")
    while True:
        try:
            line = input(">>> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nExiting.")
            break
        if not line:
            continue
        if line.lower() == "/quit":
            print("Bye!")
            break
        if line.startswith("/add "):
            path = line[5:].strip()
            try:
                content = Path(path).read_text(encoding="utf-8")
                title = Path(path).stem
                res = add_to_knowledge_base(content=content, title=title)
                print(res)
            except Exception as e:
                print(f"Error adding file: {e}")
        elif line.startswith("/search "):
            query = line[8:].strip()
            try:
                res = search_knowledge_base(query=query, max_results=5)
                print(res)
            except Exception as e:
                print(f"Error searching: {e}")
        else:
            # Treat any other input as a message to the agent
            try:
                result = agent.ainvoke({"messages": [HumanMessage(content=line)]})
                print(result["messages"][-1].content)
            except Exception as e:
                print(f"Agent error: {e}")

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