"""vector_databases_demo.py — Chroma vector store"""

import shutil
from pathlib import Path

from dotenv import load_dotenv
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings

load_dotenv()

TEXTS = [
    "The <a> tag creates a hyperlink.",
    "The <title> tag sets the browser tab title.",
    "The <h1> tag marks the main heading.",
]

QUERY = "How do I make a link on a page?"
DB_DIR = Path(__file__).parent / "chroma_demo_db"


def main() -> None:
    if DB_DIR.exists():
        shutil.rmtree(DB_DIR)

    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

    print("Storing chunks in Chroma...")
    vectorstore = Chroma.from_texts(
        texts=TEXTS,
        embedding=embeddings,
        persist_directory=str(DB_DIR),
    )
    print(f"Saved to {DB_DIR}")

    print(f"\nQuery: {QUERY}\n")
    results = vectorstore.similarity_search_with_score(QUERY, k=3)

    for i, (doc, score) in enumerate(results):
        print(f"--- result {i} (score {score:.4f}) ---")
        print(doc.page_content)


if __name__ == "__main__":
    main()