import argparse

from dotenv import load_dotenv

from agent import agent
from vectorstore import create_vectorstore, load_documents


def ingest_documents(directory: str) -> None:
    vectorstore = create_vectorstore()
    chunk_count = load_documents(directory, vectorstore)
    print(f"Loaded {chunk_count} chunks from {directory} into ChromaDB.")


def run_chat() -> None:
    print("RAG agent ready. Type 'exit' to quit.")
    while True:
        user_input = input("You: ").strip()
        if user_input.lower() == "exit":
            print("Goodbye!")
            return
        if not user_input:
            continue

        result = agent.invoke(
            {"messages": [{"role": "user", "content": user_input}]}
        )
        final_message = result["messages"][-1]
        print(f"Assistant: {final_message.content}")


def main() -> None:
    load_dotenv()

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--ingest",
        metavar="DIRECTORY",
        help="Load .txt and .md files from DIRECTORY into ChromaDB before the demo.",
    )
    args = parser.parse_args()

    if args.ingest:
        ingest_documents(args.ingest)
        return

    run_chat()


if __name__ == "__main__":
    main()