import hashlib
import json, uuid, time, os
from typing import Dict, List, Tuple

from qdrant_client import models, AsyncQdrantClient

from animuz_core.utils import init_logger

LOGGER = init_logger(__name__)

def merge_result(result: List[List[models.ScoredPoint]]) -> List[dict]:
    """
        combine dense and sparse search results by summing the scores (following the bge-m3 paper)
        the results must be [[dense_scores,...], [sparse_scores,...]]
    """
    assert len(result) == 2, "the results must be [[dense_scores,...], [sparse_scores,...]]"

    dense = {p.id:p for p in result[0]}
    sparse = {p.id:p for p in result[1]}

    new_result = []
    for k,v in sparse.items():
        if k in dense:
            v = v.model_copy()
            v.score += dense[k].score
            dense.pop(k)

        new_result.append(v)

    for k,v in dense.items(): new_result.append(v)
        
    sorted_res = sorted(new_result, key=lambda x: x.score, reverse=True)
    return [res.dict() for res in sorted_res]

class MyQdrantClient():

    def __init__(self, host:str, port:int, collection_name:str) -> None:
        self.HOST = host
        self.PORT = port
        self.COLLECTION_NAME = collection_name

        self.client = AsyncQdrantClient(host=self.HOST, port= self.PORT, api_key=os.environ.get("QDRANT_CLOUD_API_KEY"))
        
    
    # async def init_collection(self) -> str:
    #     """
    #     Initializes a collection in the database if it does not already exist.
    #     If the collection does not exist, it creates a new collection with the specified parameters.
    #     Returns a string indicating whether the collection was created or if it already exists.
    #     """
    #     if not await self.client.collection_exists(self.COLLECTION_NAME):
    #         await self.client.create_collection(
    #             collection_name=self.COLLECTION_NAME,
    #             vectors_config={
    #                 "dense": models.VectorParams(
    #                     size=1024,
    #                     distance=models.Distance.DOT,
    #                 )
    #             },
    #             sparse_vectors_config={
    #                 "sparse": models.SparseVectorParams(index=models.SparseIndexParams(on_disk=True,)),
    #             }
    #         )
    #         return "collection created"
        
    #     return "collection exists"
    
    def set_collection(self, collection_name: str) -> None:
        self.COLLECTION_NAME = collection_name

    async def upload(self, dense_embedding: List[float], indices: List[str], values: List[float], properties: Dict[str, str]) -> None:
        """
            Upload the given dense_embedding and sparse embedding (in the form of `indices` and `values`) and a dict of properties to the vector db

        """
        start = time.time()
        self.client.upload_points(
            collection_name=self.COLLECTION_NAME,
            points=[
                models.PointStruct(
                    id=str(uuid.uuid4()),
                    vector={
                            "dense" : dense_embedding,
                            "sparse" : models.SparseVector(indices=indices, values=values),
                            },
                    payload=properties
                )
            ],
        )
        LOGGER.debug(json.dumps({
            "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
            "message": "Qdrant upload latency",
            "latency": round(time.time() - start, 4),
            "exception": False
        }))

    async def upload_batch(self, dense_embeddings: List[List[float]], list_of_indices: List[List[str]], list_of_values: List[List[float]], list_of_properties: List[Dict[str, str]]) -> None:
        """
            Batch upload the given dense_embeddings and sparse embeddings (in the form of `indices` and `values`) and dicts of properties to the vector db
            
        """
        start = time.time()
        points = [
            models.PointStruct(
                    id=str(uuid.uuid4()), 
                    vector={
                            "dense" : dense_embedding,
                            "sparse" : models.SparseVector(indices=indices, values=values),
                            }, 
                    payload=properties
            ) for dense_embedding, indices, values, properties in zip(dense_embeddings, list_of_indices, list_of_values, list_of_properties)
        ]

        self.client.upload_points(
            collection_name=self.COLLECTION_NAME,
            points=points,
        )
        LOGGER.debug(json.dumps({
            "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
            "message": "Qdrant batch upload latency",
            "latency": round(time.time() - start, 4),
            "exception": False
        }))

    async def upload_batch_v2(
        self,
        dense_embeddings: List[List[float]],
        list_of_sparse_vectors: List[Dict[str, Dict[str, List]]],
        list_of_properties: List[Dict[str, str]]
    ) -> None:
        """
        Batch upload the given dense_embeddings and multiple sparse embeddings (e.g., 'sparse', 'bm25') and dicts of properties to the vector db.
        list_of_sparse_vectors: List of dicts, each mapping sparse vector name to {"indices": [...], "values": [...]}

        This is better for cases where we have multiple sparse embeddings per point.
        We also generate point IDs here based on the properties.
        """
        start = time.time()
        points = []
        for dense_embedding, sparse_vectors, properties in zip(dense_embeddings, list_of_sparse_vectors, list_of_properties):
            vector_dict = {"dense": dense_embedding}
            for name, vec in sparse_vectors.items():
                vector_dict[name] = models.SparseVector(indices=vec["indices"], values=vec["values"])
            
            # Point ID Logic
            if "point_id" in properties:
                point_id = properties["point_id"]
                # Remove point_id from properties so it is not included in the payload
                properties.pop("point_id", None)
            else:
                # Use md5 as uuid (128 bits / 16 bytes / 32 hex chars)
                prop_str = json.dumps(properties, sort_keys=True)
                point_id = hashlib.md5(prop_str.encode("utf-8")).hexdigest()

            points.append(
                models.PointStruct(
                    id=point_id, 
                    vector=vector_dict,
                    payload=properties
                )
            )

        self.client.upload_points(
            collection_name=self.COLLECTION_NAME,
            points=points,
        )
        LOGGER.debug(json.dumps({
            "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
            "message": "Qdrant batch upload latency",
            "latency": round(time.time() - start, 4),
            "exception": False
        }))

    async def hybrid_search(self, dense_embedding: List[float], indices: List[str], values: List[float], user_chat_id: str, top_k: int = 3):
        """Hybrid search using server-side RRF fusion (single round trip)."""
        start = time.time()
        limit = 10 if top_k == 3 else 20

        result = await self.client.query_points(
            collection_name=self.COLLECTION_NAME,
            prefetch=[
                models.Prefetch(
                    query=models.SparseVector(indices=indices, values=values),
                    using="sparse",
                    limit=limit,
                ),
                models.Prefetch(
                    query=dense_embedding,
                    using="dense",
                    limit=limit,
                ),
            ],
            query_filter=models.Filter(
                must=[
                    models.FieldCondition(
                        key="userChatID",
                        match=models.MatchValue(value=user_chat_id),
                    )
                ]
            ),
            query=models.FusionQuery(fusion=models.Fusion.RRF),
            limit=top_k,
        )

        LOGGER.debug(json.dumps({
            "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
            "message": "Hybrid search latency",
            "latency": round(time.time() - start, 4),
            "exception": False
        }))
        return result

    async def delete_doc(self, conditions: List[Tuple[str, str]]) -> None:
        start = time.time()
        await self.client.delete(
            collection_name=self.COLLECTION_NAME,
            points_selector=models.FilterSelector(
                filter=models.Filter(
                    must=[
                        models.FieldCondition(
                            key=cond[0],
                            match=models.MatchValue(value=cond[1]),
                        ) for cond in conditions
                    ],
                )
            ),
        )
        LOGGER.debug(json.dumps({
            "timestamp": time.strftime('[%Y-%b-%d %H:%M:%S %z]'),
            "message": "Qdrant delete doc latency",
            "latency": round(time.time() - start, 4),
            "exception": False
        }))