import numpy as np
from docarray import BaseDoc
from docarray.index.backends.weaviate import WeaviateDocumentIndex
from docarray.typing import NdArray
from pydantic import Field


class Document(BaseDoc):
    text: str
    embedding: NdArray[2] = Field(dims=2, is_embedding=True)
    file: NdArray[100] = Field(dims=100)


batch_config = dict(batch_size=20, dynamic=False, timeout_retries=3, num_workers=1)
runtime_config = WeaviateDocumentIndex.RuntimeConfig(batch_config=batch_config)

db_config = WeaviateDocumentIndex.DBConfig(host="http://localhost:8080")
store = WeaviateDocumentIndex[Document](db_config=db_config)
store.configure(runtime_config)

docs = [
    Document(text="Hello world", embedding=np.array([1, 2]), file=np.random.rand(100), id="1"),
    Document(text="Hello world, how are you?", embedding=np.array([3, 4]), file=np.random.rand(100), id="2"),
    Document(text="Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut",
             embedding=np.array([5, 6]), file=np.random.rand(100), id="3", ),
]
store.index(docs)

hybrid_search = store.build_query().text_search("world").find([1, 2]).limit(2).build()
hybrid_search_result = store.execute_query(hybrid_search)
print(hybrid_search_result.text)