from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import httpx
import pinecone
from pinecone import Index
from elasticsearch import AsyncElasticsearch
import weaviate
import logging
from typing import Dict, Any
from sentence_transformers import SentenceTransformer
from pymilvus import Milvus, DataType

from cassandra.cluster import Cluster
from cassandra.auth import PlainTextAuthProvider

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class VectorQueryRequest(BaseModel):
    api_key: str
    vector_db_type: str
    user_query: str
    top_k: int = 5

class VectorDBQueryHandler:
    def __init__(self, vector_details: Dict[str, Any], user_query: str):
        self.api_key = vector_details.get("auth")
        self.vector_db_type = vector_details.get("type")
        self.endpoint = vector_details.get("url")
        self.user_query = user_query
        self.index_name = vector_details.get("index")
        self.top_k = vector_details.get("top_k", 5)
        self.model = vector_details.get("model", "all-MiniLM-L6-v2")

        # Initialize your model here
        # Initializing the model
        self.model = self.initialize_model()

    def initialize_model(self):
        # Loading a pre-trained model using the model name
        # The try-except block is to handle cases where the model name might be incorrect or not supported
        try:
            model = SentenceTransformer(self.model_name)
            return model
        except Exception as e:
            # Log the error or handle it as per your application's requirements
            print(f"Error loading model {self.model_name}: {e}")
            return None

    async def vector_query(self):
        if self.vector_db_type == "pinecone":
            return await self.query_pinecone()
        elif self.vector_db_type == "elasticsearch":
            return await self.query_elasticsearch()
        elif self.vector_db_type == "weaviate":
            return await self.query_weaviate()
        elif self.vector_db_type == "milvus":
            return await self.query_milvus()
        elif  self.vector_db_type == "chroma":
            return await self.query_chroma()
        elif  self.vector_db_type == "vectara":
            return await self.query_vectara()
        elif  self.vector_db_type == "astra":
            return await self.query_astra()
        else:
            return {"error": "Unsupported vector database type"}


    async def query_pinecone(self):
        #endpoint, api_key, user_query, index_name, top_k are the parameters
        # Ensure Pinecone is initialized and the index is set up
        pinecone.init(api_key=self.api_key)
        index = pinecone.Index(self.index_name)

        # Convert the text query to a vector using the loaded model
        query_vector = self.model.encode([self.user_query]).tolist()

        # Perform the vector search in Pinecone
        response = index.query(queries=query_vector, top_k=self.top_k)

        # Process Pinecone's response
        # Adjust the response handling according to Pinecone's actual response format and your requirements
        results = []
        if "matches" in response and len(response["matches"]) > 0:
            for match in response["matches"][0]:
                # Assuming you have a way to fetch document details (e.g., from another database or metadata stored in Pinecone)
                doc_details = await self.fetch_document_details(match["id"])
                results.append({
                    "id": match["id"],
                    "score": match["score"],
                    "link": doc_details.get("link", "No link available"),  # Handle according to your document details retrieval method
                })
        return results

    async def fetch_document_details(self, doc_id):
        # Placeholder for an asynchronous method to fetch document details based on document ID
        # Implement this method based on how your application stores and retrieves document metadata
        # For illustration, returning a dummy link; replace with your actual data retrieval logic
        return {"link": f"http://example.com/document/{doc_id}"}

    async def query_elasticsearch(self):
        #endpoint, api_key, user_query, index_name, top_k are the parameters
        es_endpoint = self.endpoint
        es = AsyncElasticsearch([es_endpoint], headers={"Authorization": f"ApiKey {self.api_key}"})

        query_body = {
            "query": {
                "simple_query_string": {
                    "query": self.user_query,
                    "fields": ["title", "content"],  # Adjust based on your document structure
                }
            },
            "size": self.top_k
        }

        try:
            response = await es.search(index=self.index_name, body=query_body)
            hits = response['hits']['hits']
            logger.info(f"Found {len(hits)} hits for query: {self.user_query}")

            results = [{
                "id": hit["_id"],
                "score": hit["_score"],
                "source": hit["_source"]  # Contains the document content and metadata
            } for hit in hits]
            return results
        except Exception as e:
            logger.error(f"Error performing search in Elasticsearch: {e}")
            # Proper use of HTTPException for FastAPI applications, ensuring correct HTTP status codes are used
            raise HTTPException(status_code=500, detail="An error occurred while performing the search.")
        finally:
            await es.close()

    async def query_weaviate(self):
        # Construct your query. Adjust based on whether you're using text or vector search.
        # The following is a basic text search example, assuming Weaviate's automatic text-to-vector conversion is enabled.
        
        #endpoint, index_name, user_query, top_k are the parameters
        weaviate_client = weaviate.Client(self.endpoint)
        try:
            # Perform the search query using near_text for text-based search
            response = weaviate_client.query.get(
                self.index_name
            ).with_limit(
                self.top_k
            ).with_near_text({
                "concepts": [self.user_query]
            }).do()

            hits = response["data"]["Get"][self.index_name]
            
            # Process and return the response
            results = [{
                "id": hit["id"],
                "name": hit.get("name", "No name available"),
                "content": hit.get("content", "No content available")
            } for hit in hits]
            
            return results
        except Exception as e:
            logger.error(f"Error performing search in Weaviate: {e}")
            # Decide how to handle the exception, e.g., re-raise or return a specific error message
            raise

    async def query_milvus(self):
        #endpoint, customer_id, api_key, user_query, index_name, top_k are the parameters
        milvus_client = Milvus(uri=self.endpoint)
        try:
            # Convert the user query to a vector
            query_vector = self.model.encode([self.user_query]).tolist()

            # Prepare the search parameters
            search_params = {
                "collection_name": self.index_name,
                "query_records": query_vector,
                "top_k": self.top_k,
                "params": {"metric_type": MetricType.L2, "params": {"nprobe": 10}},
            }

            # Perform the search
            status, results = milvus_client.search(**search_params)

            if not status.OK():
                logger.error(f"Milvus search error: {status}")
                return []

            # Process and return results
            processed_results = []
            for result in results[0]:  # Assuming the structure of results
                processed_results.append({
                    "id": result.id,
                    "distance": result.distance  # Or use 'score': 1 / (1 + result.distance) for similarity
                })

            return processed_results
        except Exception as e:
            logger.error(f"Error performing search in Milvus: {e}")
            raise

    async def query_vectara(self):
        #endpoint, customer_id, api_key, user_query, index_name, top_k are the parameters

        endpoint = f"{self.endpoint}"  # Adjust with actual Vectara query endpoint
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        payload = {
            "query": self.user_query,
            "customer_id": "your_customer_id",  # Replace with your actual customer ID
            "corpus_key": self.index_name,  # Assuming this corresponds to Vectara's corpus concept
            "top_k": self.top_k
        }

        try:
            async with httpx.AsyncClient() as client:
                response = await client.post(endpoint, headers=headers, json=payload)
                response.raise_for_status()  # Raises exception for 4XX/5XX responses
                
                data = response.json()
                hits = data.get("results", [])  # Adjust based on actual response structure
                
                # Process and return the response
                results = [{
                    "id": hit["documentId"],  # Adjust according to Vectara's response structure
                    "score": hit.get("score"),
                    # Extract other relevant fields as needed
                } for hit in hits]

                return results
        except httpx.HTTPStatusError as e:
            logger.error(f"HTTP error occurred during Vectara query: {e}")
        except Exception as e:
            logger.error(f"Error performing search in Vectara: {e}")
            # Handle other types of errors appropriately
            raise

    async def query_astra(self):
        # Setup authentication and connect to AstraDB

        #client_id, client_secret, endpoint, keyspace, table_name are the parameters
        
        auth_provider = PlainTextAuthProvider(username=self.client_id, password=self.client_secret)
        cluster = Cluster(cloud={'secure_connect_bundle': self.endpoint}, auth_provider=auth_provider)
        session = cluster.connect()

        try:
            # Adjust the CQL query based on your specific use case and schema
            # This is a simple SELECT statement as an example. You'll need to customize it.
            cql_query = f"SELECT * FROM {self.keyspace}.{self.table_name} WHERE your_search_field = %s"
            prepared = session.prepare(cql_query)
            rows = session.execute(prepared, [self.user_query])

            results = [{"id": row.id, "content": row.content} for row in rows]  # Adjust based on your schema
            return results
        except Exception as e:
            logger.error(f"Error querying AstraDB: {e}")
            raise
        finally:
            session.shutdown()

    async def perform_vector_search(query_request: VectorQueryRequest):
        handler = VectorDBQueryHandler(
            api_key=query_request.api_key,
            vector_db_type=query_request.vector_db_type,
            user_query=query_request.user_query,
            index_name="your_index_name",  # This needs to be passed or set correctly
            top_k=query_request.top_k
        )
        try:
            results = await handler.vector_query()
            return {"results": results}
        except HTTPException as e:
            raise e
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))
