import os
import re
import logging
import sys
from azure.core.credentials import AzureKeyCredential
from azure.search.documents import SearchClient
from azure.search.documents.indexes import SearchIndexClient
from llama_index.core import (
    SimpleDirectoryReader,
    StorageContext,
    VectorStoreIndex,
)
from llama_index.core.settings import Settings
from llama_index.llms.azure_openai import AzureOpenAI
from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
from llama_index.vector_stores.azureaisearch import AzureAISearchVectorStore
from llama_index.vector_stores.azureaisearch import (
    IndexManagement,
    MetadataIndexFieldType,
)

from dotenv import load_dotenv

load_dotenv()

def azure_search(query):
    aoai_api_key = os.getenv("AZURE_OPENAI_API_KEY")
    aoai_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
    aoai_api_version = "2024-02-01"


    llm = AzureOpenAI(
        model="gpt-4",
        deployment_name="gpt-4",
        api_key=aoai_api_key,
        azure_endpoint=aoai_endpoint,
        api_version=aoai_api_version,
    )

    # You need to deploy your own embedding model as well as your own chat completion model
    embed_model = AzureOpenAIEmbedding(
        model="text-embedding-ada-002",
        deployment_name="text-embedding-ada-002",
        api_key=aoai_api_key,
        azure_endpoint=aoai_endpoint,
        api_version=aoai_api_version,
    )

    search_service_api_key = os.getenv("SEARCH_SERVICE_API_KEY")
    search_service_endpoint = os.getenv("SEARCH_SERVICE_ENDPOINT")
    search_service_api_version = "2023-11-01"
    credential = AzureKeyCredential(os.getenv("SEARCH_SERVICE_API_KEY"))

    # Index name to use
    index_name = "llamaindex-vector-demo"

    # Use index client to demonstrate creating an index
    index_client = SearchIndexClient(
        endpoint=search_service_endpoint,
        credential=credential,
    )

    # Use search client to demonstration using existing index
    search_client = SearchClient(
        endpoint=search_service_endpoint,
        index_name=index_name,
        credential=credential,
    )

    index_name = "llamaindex-vector-demo"

    metadata_fields = {
        "document_name": "document_name",
    }
    vector_store = AzureAISearchVectorStore(
        search_or_index_client=search_client,
        filterable_metadata_field_keys=metadata_fields,
        index_management=IndexManagement.VALIDATE_INDEX,
        id_field_key="id",
        chunk_field_key="chunk",
        embedding_field_key="embedding",
        embedding_dimensionality=1536,
        metadata_string_field_key="metadata",
        doc_id_field_key="doc_id",
    )

    Settings.llm = llm
    Settings.embed_model = embed_model
    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    index = VectorStoreIndex.from_documents(
        [],
        storage_context=storage_context,
    )

    # Query Data
    query_engine = index.as_query_engine(similarity_top_k=2)
    response = query_engine.query(query)
    return response.response