import argparse
from tools import add_to_knowledge_base, search_knowledge_base


def main():
    parser = argparse.ArgumentParser(description="RAG Agent CLI")
    subparsers = parser.add_subparsers(dest='command')

    add_parser = subparsers.add_parser('add', help='Add document to knowledge base')
    add_parser.add_argument('--title', required=True, help='Title of the document')
    add_parser.add_argument('--content', required=True, help='Content of the document')

    search_parser = subparsers.add_parser('search', help='Search knowledge base')
    search_parser.add_argument('--query', required=True, help='Query string')
    search_parser.add_argument('--max', type=int, default=5, help='Max results')

    args = parser.parse_args()

    if args.command == 'add':
        result = add_to_knowledge_base(args.content, args.title)
        print(result)
    elif args.command == 'search':
        results = search_knowledge_base(args.query, args.max)
        for i, r in enumerate(results, 1):
            print(f"{i}. {r}")
    else:
        parser.print_help()

if __name__ == "__main__":
    main()