import os

from haystack import Document
from haystack.components.embedders import OpenAITextEmbedder
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.utils import Secret


model = os.getenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")
dimensions_env = os.getenv("OPENAI_EMBEDDING_DIMENSIONS", "512")
dimensions = int(dimensions_env) if dimensions_env else None
api_base_url = os.getenv("OPENAI_API_BASE_URL")

embedder = OpenAITextEmbedder(
    api_key=Secret.from_env_var("OPENAI_API_KEY"),
    model=model,
    dimensions=dimensions,
    api_base_url=api_base_url,
    timeout=30.0,
    max_retries=1,
)

result = embedder.run(text="password reset approval policy")
embedding = result["embedding"]
meta = result["meta"]

document_store = InMemoryDocumentStore(embedding_similarity_function="cosine")
document_store.write_documents(
    [
        Document(
            content="Password reset requests require help desk approval.",
            meta={"name": "password-reset-policy"},
            embedding=embedding,
        )
    ]
)

retriever = InMemoryEmbeddingRetriever(document_store=document_store, top_k=1)
match = retriever.run(query_embedding=embedding)["documents"][0]
usage = meta.get("usage", {})

print(f"model={meta.get('model', model)}")
print(f"embedding_dimensions={len(embedding)}")
print(f"first_values={embedding[0]:.4f},{embedding[1]:.4f},{embedding[2]:.4f}")
print(f"usage_total_tokens={usage.get('total_tokens', 'unknown')}")
print(f"retriever_top_document={match.meta['name']}")