from langchain.vectorstores import Weaviate
from langchain_weaviate.vectorstores import WeaviateVectorStore

import weaviate
from weaviate.classes.init import Auth
from weaviate.config import AdditionalConfig, Timeout
from app.embedding import WeaviateCustomEmbeddings
import os

internal_url = "vector-store-01jse60hep6416b15n1nsj6ch8.cheetah.svc.cluster.local"
external_url = "183.111.96.67"
# http_host = os.getenv("WEAVIATE_HTTP_HOST", "183.111.96.67")
http_host = os.getenv("WEAVIATE_HTTP_HOST", external_url)
# grpc_host = os.getenv("WEAVIATE_GRPC_HOST", "183.111.96.67")
grpc_host = os.getenv("WEAVIATE_GRPC_HOST", external_url)
weaviate_api_key = os.getenv("WEAVIATE_API_KEY", "01jse60hwsf37za5kmnkzzcwa9")

def get_retriever():
    client = weaviate.connect_to_custom(
        http_host=http_host,        # Hostname for the HTTP API connection
        http_port=31008,            # Default is 80, WCD uses 443
        http_secure=False,          # Whether to use https (secure) for the HTTP API connection
        grpc_host=grpc_host,        # Hostname for the gRPC API connection
        grpc_port=30311,            # Default is 50051, WCD uses 443
        grpc_secure=False,          # Whether to use a secure channel for the gRPC API connection
        auth_credentials=Auth.api_key(weaviate_api_key),    # API key for authentication
    )

    embedding = WeaviateCustomEmbeddings()

    vectorstore = WeaviateVectorStore(
        client=client,
        index_name="Loca",
        text_key="page_content",
        embedding=embedding,
        attributes=["source", "page_content"]
    )

    retriever = vectorstore.as_retriever(
        search_kwargs={
            "k": 5
        }
    )

    return retriever