"""
Qdrant client wrapper for Library Desk.

Provides async vector operations with:
- Collection-per-user multi-tenancy
- Document chunk storage with embeddings
- Semantic search
- Similarity queries
"""

from qdrant_client import QdrantClient
from qdrant_client.models import (
    Distance, VectorParams, PointStruct,
    Filter, FieldCondition, MatchValue, Range
)
from typing import List, Dict, Any, Optional
import uuid
import logging

from src.core.multi_tenancy import get_qdrant_collection_name

logger = logging.getLogger(__name__)


class QdrantClientWrapper:
    """
    Qdrant client wrapper with multi-tenancy support.

    Pattern: Collection per user (from qdrant_memory.py)
    Each user has isolated vector collection for their documents.
    """

    def __init__(self, url: str, embedding_dim: int = 768):
        """
        Initialize Qdrant client.

        Args:
            url: Qdrant server URL (e.g., "http://qdrant:6333")
            embedding_dim: Vector embedding dimension (default 768 for nomic-embed-text)
        """
        self.client = QdrantClient(url=url)
        self.embedding_dim = embedding_dim
        logger.info(f"Initialized Qdrant client: {url}")

    def get_collection_name(self, user: str) -> str:
        """
        Get collection name for user.

        Args:
            user: User identifier

        Returns:
            Collection name (e.g., "library_desk_jpmschweitzer")
        """
        return get_qdrant_collection_name(user)

    async def ensure_collection(self, collection_name: str):
        """
        Create collection if it doesn't exist.

        Args:
            collection_name: Collection name
        """
        try:
            collections = self.client.get_collections()
            existing = [c.name for c in collections.collections]

            if collection_name not in existing:
                logger.info(f"Creating Qdrant collection: {collection_name}")
                self.client.create_collection(
                    collection_name=collection_name,
                    vectors_config=VectorParams(
                        size=self.embedding_dim,
                        distance=Distance.COSINE
                    )
                )
                logger.info(f"Created collection: {collection_name}")
        except Exception as e:
            logger.error(f"Error ensuring collection: {e}", exc_info=True)
            raise

    async def collection_exists(self, collection_name: str) -> bool:
        """
        Check if collection exists.

        Args:
            collection_name: Collection name

        Returns:
            True if collection exists
        """
        try:
            collections = self.client.get_collections()
            existing = [c.name for c in collections.collections]
            return collection_name in existing
        except Exception as e:
            logger.error(f"Error checking collection: {e}", exc_info=True)
            return False

    async def upsert_document_chunks(
        self,
        user: str,
        doc_id: str,
        chunks: List[Dict[str, Any]],
        embeddings: List[List[float]]
    ) -> int:
        """
        Upsert document chunks with embeddings.

        Point structure:
        {
            id: "doc_123_chunk_0",
            vector: [...],
            payload: {
                doc_id: "doc_123",
                chunk_index: 0,
                content: "text content",
                metadata: {...}
            }
        }

        Args:
            user: User identifier
            doc_id: Document ID
            chunks: List of chunk dictionaries with content
            embeddings: List of embedding vectors

        Returns:
            Number of chunks upserted

        Raises:
            ValueError: If chunks and embeddings length mismatch
        """
        if len(chunks) != len(embeddings):
            raise ValueError(
                f"Chunks ({len(chunks)}) and embeddings ({len(embeddings)}) length mismatch"
            )

        collection_name = self.get_collection_name(user)
        await self.ensure_collection(user)

        points = []
        for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
            # Generate deterministic point ID
            point_id = str(uuid.uuid5(
                uuid.NAMESPACE_DNS,
                f"{doc_id}_chunk_{i}"
            ))

            points.append(PointStruct(
                id=point_id,
                vector=embedding,
                payload={
                    "doc_id": doc_id,
                    "chunk_index": i,
                    "content": chunk.get("content", ""),
                    "metadata": chunk.get("metadata", {})
                }
            ))

        try:
            self.client.upsert(
                collection_name=collection_name,
                points=points
            )
            logger.info(f"Upserted {len(points)} chunks for {doc_id}")
            return len(points)
        except Exception as e:
            logger.error(f"Failed to upsert chunks: {e}", exc_info=True)
            raise

    async def search(
        self,
        user: str,
        query_vector: List[float],
        limit: int = 10,
        score_threshold: float = 0.7,
        filter_dict: Optional[Dict[str, Any]] = None
    ) -> List[Dict[str, Any]]:
        """
        Semantic search in user's collection.

        Args:
            user: User identifier
            query_vector: Query embedding vector
            limit: Maximum results to return
            score_threshold: Minimum similarity score (0.0-1.0)
            filter_dict: Optional payload filters

        Returns:
            List of matches with scores and payloads

        Example:
            results = await client.search(
                user="jpmschweitzer",
                query_vector=[0.1, 0.2, ...],
                limit=5,
                filter_dict={"doc_id": "doc_123"}
            )
        """
        collection_name = self.get_collection_name(user)

        # Build filter if provided
        query_filter = None
        if filter_dict:
            conditions = []
            for key, value in filter_dict.items():
                conditions.append(
                    FieldCondition(key=key, match=MatchValue(value=value))
                )
            query_filter = Filter(must=conditions)

        try:
            response = self.client.query_points(
                collection_name=collection_name,
                query=query_vector,
                limit=limit,
                score_threshold=score_threshold,
                query_filter=query_filter,
                with_payload=True
            )

            return [
                {
                    "id": str(point.id),
                    "score": point.score,
                    "doc_id": point.payload["doc_id"],
                    "chunk_index": point.payload["chunk_index"],
                    "content": point.payload["content"],
                    "metadata": point.payload.get("metadata", {})
                }
                for point in response.points
            ]
        except Exception as e:
            logger.error(f"Search failed: {e}", exc_info=True)
            return []

    async def delete_document(self, user: str, doc_id: str) -> bool:
        """
        Delete all chunks for a document.

        Args:
            user: User identifier
            doc_id: Document ID

        Returns:
            True if deleted successfully
        """
        collection_name = self.get_collection_name(user)

        try:
            self.client.delete(
                collection_name=collection_name,
                points_selector=Filter(
                    must=[
                        FieldCondition(
                            key="doc_id",
                            match=MatchValue(value=doc_id)
                        )
                    ]
                )
            )
            logger.info(f"Deleted chunks for {doc_id}")
            return True
        except Exception as e:
            logger.error(f"Failed to delete document: {e}", exc_info=True)
            return False

    async def get_document_chunks(
        self,
        user: str,
        doc_id: str
    ) -> List[Dict[str, Any]]:
        """
        Get all chunks for a document.

        Args:
            user: User identifier
            doc_id: Document ID

        Returns:
            List of chunks with content and metadata
        """
        collection_name = self.get_collection_name(user)

        try:
            # Scroll through points with doc_id filter
            points, _ = self.client.scroll(
                collection_name=collection_name,
                scroll_filter=Filter(
                    must=[
                        FieldCondition(
                            key="doc_id",
                            match=MatchValue(value=doc_id)
                        )
                    ]
                ),
                limit=1000,
                with_payload=True,
                with_vectors=False
            )

            return [
                {
                    "id": str(point.id),
                    "chunk_index": point.payload["chunk_index"],
                    "content": point.payload["content"],
                    "metadata": point.payload.get("metadata", {})
                }
                for point in points
            ]
        except Exception as e:
            logger.error(f"Failed to get document chunks: {e}", exc_info=True)
            return []

    async def count_documents(self, user: str) -> int:
        """
        Count total number of unique documents in user's collection.

        Args:
            user: User identifier

        Returns:
            Number of unique documents
        """
        collection_name = self.get_collection_name(user)

        try:
            # Get collection info
            collection_info = self.client.get_collection(collection_name)
            # This gives total points, not unique docs
            # For unique docs, would need to aggregate by doc_id
            return collection_info.points_count
        except Exception as e:
            logger.error(f"Failed to count documents: {e}", exc_info=True)
            return 0

    async def delete_collection(self, user: str) -> bool:
        """
        Delete user's entire collection.

        Warning: This removes all data for the user!

        Args:
            user: User identifier

        Returns:
            True if deleted successfully
        """
        collection_name = self.get_collection_name(user)

        try:
            self.client.delete_collection(collection_name)
            logger.warning(f"Deleted collection: {collection_name}")
            return True
        except Exception as e:
            logger.error(f"Failed to delete collection: {e}", exc_info=True)
            return False

    async def find_similar_chunks(
        self,
        user: str,
        doc_id: str,
        limit: int = 10
    ) -> List[Dict[str, Any]]:
        """
        Find chunks similar to those in a given document.

        Strategy: Get all chunks from doc, use their vectors to find similar chunks.

        Args:
            user: User identifier
            doc_id: Source document ID
            limit: Maximum results per chunk

        Returns:
            List of similar chunks from other documents
        """
        collection_name = self.get_collection_name(user)

        try:
            # Get source document chunks with vectors
            source_points, _ = self.client.scroll(
                collection_name=collection_name,
                scroll_filter=Filter(
                    must=[
                        FieldCondition(
                            key="doc_id",
                            match=MatchValue(value=doc_id)
                        )
                    ]
                ),
                limit=10,  # Sample first 10 chunks
                with_payload=True,
                with_vectors=True
            )

            if not source_points:
                return []

            # Search using first chunk's vector
            # (could aggregate multiple chunks for better results)
            first_vector = source_points[0].vector

            response = self.client.query_points(
                collection_name=collection_name,
                query=first_vector,
                limit=limit * 2,  # Get more to filter out same doc
                with_payload=True
            )

            # Filter out chunks from same document
            similar = [
                {
                    "id": str(point.id),
                    "score": point.score,
                    "doc_id": point.payload["doc_id"],
                    "content": point.payload["content"]
                }
                for point in response.points
                if point.payload["doc_id"] != doc_id
            ]

            return similar[:limit]

        except Exception as e:
            logger.error(f"Failed to find similar chunks: {e}", exc_info=True)
            return []

    async def upsert_vector(
        self,
        collection_name: str,
        vector_id: str,
        vector: List[float],
        payload: Dict[str, Any]
    ) -> bool:
        """
        Upsert a single vector point.

        Args:
            collection_name: Collection name
            vector_id: Point ID
            vector: Embedding vector
            payload: Point payload/metadata

        Returns:
            True if successful
        """
        try:
            self.client.upsert(
                collection_name=collection_name,
                points=[PointStruct(
                    id=vector_id,
                    vector=vector,
                    payload=payload
                )]
            )
            return True
        except Exception as e:
            logger.error(f"Failed to upsert vector: {e}", exc_info=True)
            return False

    async def delete_by_filter(
        self,
        collection_name: str,
        filter_conditions: Dict[str, Any]
    ) -> int:
        """
        Delete points matching filter conditions.

        Args:
            collection_name: Collection name
            filter_conditions: Filter conditions (e.g., {"page_id": 5})

        Returns:
            Number of points deleted (approximation)
        """
        try:
            # Build filter
            conditions = []
            for key, value in filter_conditions.items():
                conditions.append(
                    FieldCondition(key=key, match=MatchValue(value=value))
                )
            query_filter = Filter(must=conditions)

            # Delete points
            result = self.client.delete(
                collection_name=collection_name,
                points_selector=query_filter
            )

            # Return operation status (Qdrant doesn't return count directly)
            logger.info(f"Deleted points with filter {filter_conditions}")
            return 1  # Placeholder, actual count not available from API

        except Exception as e:
            logger.error(f"Failed to delete by filter: {e}", exc_info=True)
            return 0

    async def search_vectors(
        self,
        collection_name: str,
        query_vector: List[float],
        limit: int = 10,
        score_threshold: float = 0.5,
        filter_conditions: Optional[Dict[str, Any]] = None
    ) -> List[Dict[str, Any]]:
        """
        Search vectors in collection.

        Args:
            collection_name: Collection name
            query_vector: Query embedding vector
            limit: Maximum results
            score_threshold: Minimum similarity score
            filter_conditions: Optional filter conditions

        Returns:
            List of search results with scores and payloads
        """
        # Build filter if provided
        query_filter = None
        if filter_conditions:
            conditions = []
            for key, value in filter_conditions.items():
                conditions.append(
                    FieldCondition(key=key, match=MatchValue(value=value))
                )
            query_filter = Filter(must=conditions)

        try:
            response = self.client.query_points(
                collection_name=collection_name,
                query=query_vector,
                limit=limit,
                score_threshold=score_threshold,
                query_filter=query_filter,
                with_payload=True
            )

            return [
                {
                    "id": str(point.id),
                    "score": point.score,
                    "payload": dict(point.payload)
                }
                for point in response.points
            ]
        except Exception as e:
            logger.error(f"Search failed: {e}", exc_info=True)
            return []

    async def scroll_all_points(
        self,
        collection_name: str,
        batch_size: int = 100,
        with_payload: bool = True,
        with_vectors: bool = False,
        filter_conditions: Optional[Dict[str, Any]] = None
    ) -> List[Dict[str, Any]]:
        """
        Scroll through all points in a collection.

        Args:
            collection_name: Collection name
            batch_size: Number of points per batch
            with_payload: Include payload in results
            with_vectors: Include vectors in results
            filter_conditions: Optional filter conditions

        Returns:
            List of all points with id and payload
        """
        all_points = []
        offset = None

        # Build filter if provided
        scroll_filter = None
        if filter_conditions:
            conditions = []
            for key, value in filter_conditions.items():
                conditions.append(
                    FieldCondition(key=key, match=MatchValue(value=value))
                )
            scroll_filter = Filter(must=conditions)

        try:
            while True:
                points, next_offset = self.client.scroll(
                    collection_name=collection_name,
                    scroll_filter=scroll_filter,
                    limit=batch_size,
                    offset=offset,
                    with_payload=with_payload,
                    with_vectors=with_vectors
                )

                for point in points:
                    all_points.append({
                        "id": str(point.id),
                        "payload": dict(point.payload) if point.payload else {}
                    })

                if next_offset is None:
                    break
                offset = next_offset

            return all_points

        except Exception as e:
            logger.error(f"Failed to scroll collection {collection_name}: {e}", exc_info=True)
            return []

    async def delete_by_ids(
        self,
        collection_name: str,
        point_ids: List[str]
    ) -> int:
        """
        Delete points by their IDs.

        Args:
            collection_name: Collection name
            point_ids: List of point IDs to delete

        Returns:
            Number of points deleted
        """
        if not point_ids:
            return 0

        try:
            self.client.delete(
                collection_name=collection_name,
                points_selector=point_ids
            )
            logger.info(f"Deleted {len(point_ids)} points from {collection_name}")
            return len(point_ids)

        except Exception as e:
            logger.error(f"Failed to delete points by IDs: {e}", exc_info=True)
            return 0

    async def list_collections(self) -> List[Dict[str, Any]]:
        """
        List all collections with stats.

        Returns:
            List of collection info dictionaries
        """
        try:
            collections = self.client.get_collections()
            result = []

            for coll in collections.collections:
                # Get detailed collection info
                try:
                    info = self.client.get_collection(coll.name)
                    result.append({
                        "name": coll.name,
                        "vectors_count": info.vectors_count or 0,
                        "points_count": info.points_count or 0,
                        "segments_count": info.segments_count or 0
                    })
                except Exception as e:
                    logger.warning(f"Failed to get info for collection {coll.name}: {e}")
                    result.append({
                        "name": coll.name,
                        "vectors_count": 0,
                        "points_count": 0,
                        "segments_count": 0
                    })

            return result

        except Exception as e:
            logger.error(f"Failed to list collections: {e}", exc_info=True)
            return []

    # ========== Volatile Data Methods ==========

    async def search_with_expiry_filter(
        self,
        collection_name: str,
        query_vector: List[float],
        current_timestamp: int,
        limit: int = 10,
        score_threshold: float = 0.7
    ) -> List[Dict[str, Any]]:
        """
        Search vectors filtering out expired records.

        Args:
            collection_name: Collection name
            query_vector: Query embedding vector
            current_timestamp: Current time in milliseconds
            limit: Maximum results
            score_threshold: Minimum similarity score

        Returns:
            List of non-expired search results
        """
        # Filter: ttl_expiry > current_timestamp (not expired)
        expiry_filter = Filter(
            must=[
                FieldCondition(
                    key="ttl_expiry",
                    range=Range(gt=current_timestamp)
                )
            ]
        )

        try:
            response = self.client.query_points(
                collection_name=collection_name,
                query=query_vector,
                limit=limit,
                score_threshold=score_threshold,
                query_filter=expiry_filter,
                with_payload=True
            )

            return [
                {
                    "id": str(point.id),
                    "score": point.score,
                    "payload": dict(point.payload)
                }
                for point in response.points
            ]
        except Exception as e:
            logger.error(f"Volatile search failed: {e}", exc_info=True)
            return []

    async def delete_expired_vectors(
        self,
        collection_name: str,
        current_timestamp: int
    ) -> int:
        """
        Delete all vectors where ttl_expiry < current_timestamp.

        Args:
            collection_name: Collection name
            current_timestamp: Current time in milliseconds

        Returns:
            Number of points deleted (approximate)
        """
        # Filter: ttl_expiry < current_timestamp (expired)
        expiry_filter = Filter(
            must=[
                FieldCondition(
                    key="ttl_expiry",
                    range=Range(lt=current_timestamp)
                )
            ]
        )

        try:
            # First count how many will be deleted (scroll to count)
            count = 0
            offset = None
            while True:
                points, next_offset = self.client.scroll(
                    collection_name=collection_name,
                    scroll_filter=expiry_filter,
                    limit=100,
                    offset=offset,
                    with_payload=False
                )
                count += len(points)
                if next_offset is None:
                    break
                offset = next_offset

            if count == 0:
                return 0

            # Delete expired points
            self.client.delete(
                collection_name=collection_name,
                points_selector=expiry_filter
            )

            logger.info(f"Deleted {count} expired vectors from {collection_name}")
            return count

        except Exception as e:
            logger.error(f"Failed to delete expired vectors: {e}", exc_info=True)
            return 0

    async def get_volatile_collections(self) -> List[str]:
        """
        Get all volatile collections (prefixed with 'volatile_').

        Returns:
            List of volatile collection names
        """
        try:
            collections = self.client.get_collections()
            return [
                c.name for c in collections.collections
                if c.name.startswith("volatile_")
            ]
        except Exception as e:
            logger.error(f"Failed to list volatile collections: {e}", exc_info=True)
            return []