"""
Vector database abstraction layer.

Provides a unified interface for different vector store backends:
- Chroma
- Qdrant  
- Weaviate
"""

import uuid
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

from .logging import get_logger

logger = get_logger(__name__)


@dataclass
class Document:
    """Represents a document chunk with metadata."""
    id: str
    content: str
    embedding: Optional[List[float]] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    @classmethod
    def create(
        cls,
        content: str,
        metadata: Optional[Dict[str, Any]] = None,
        embedding: Optional[List[float]] = None
    ) -> "Document":
        """Create a new document with auto-generated ID."""
        return cls(
            id=str(uuid.uuid4()),
            content=content,
            embedding=embedding,
            metadata=metadata or {}
        )


@dataclass
class SearchResult:
    """Represents a search result from the vector store."""
    document: Document
    score: float
    
    @property
    def content(self) -> str:
        return self.document.content
    
    @property
    def metadata(self) -> Dict[str, Any]:
        return self.document.metadata


class VectorStore(ABC):
    """Abstract base class for vector store implementations."""
    
    @abstractmethod
    def add_documents(self, documents: List[Document]) -> List[str]:
        """
        Add documents to the vector store.
        
        Args:
            documents: List of documents to add
            
        Returns:
            List of document IDs
        """
        pass
    
    @abstractmethod
    def search(
        self,
        query_embedding: List[float],
        top_k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """
        Search for similar documents.
        
        Args:
            query_embedding: Query vector
            top_k: Number of results to return
            filter_metadata: Optional metadata filters
            
        Returns:
            List of search results
        """
        pass
    
    @abstractmethod
    def delete(self, document_ids: List[str]) -> None:
        """Delete documents by ID."""
        pass
    
    @abstractmethod
    def clear(self) -> None:
        """Clear all documents from the store."""
        pass
    
    @abstractmethod
    def count(self) -> int:
        """Get the number of documents in the store."""
        pass
    
    def search_by_text(
        self,
        query: str,
        embed_fn: callable,
        top_k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """
        Search using text query (requires embedding function).
        
        Args:
            query: Text query
            embed_fn: Function to embed text
            top_k: Number of results
            filter_metadata: Optional filters
            
        Returns:
            List of search results
        """
        query_embedding = embed_fn(query)
        return self.search(query_embedding, top_k, filter_metadata)


class ChromaStore(VectorStore):
    """Chroma vector store implementation."""
    
    def __init__(
        self,
        collection_name: str = "readanybook",
        persist_directory: Optional[str] = None,
        host: Optional[str] = None,
        port: Optional[int] = None,
        distance_metric: str = "cosine"
    ):
        """
        Initialize Chroma store.
        
        Args:
            collection_name: Name of the collection
            persist_directory: Directory for persistent storage (local mode)
            host: Chroma server host (server mode)
            port: Chroma server port (server mode)
            distance_metric: Distance metric (cosine, l2, ip)
        """
        try:
            import chromadb
            from chromadb.config import Settings as ChromaSettings
        except ImportError:
            raise ImportError("chromadb is required. Install with: pip install chromadb")
        
        self.collection_name = collection_name
        self.distance_metric = distance_metric
        
        # Initialize client
        if host and port:
            # Server mode
            self.client = chromadb.HttpClient(host=host, port=port)
            logger.info(f"Connected to Chroma server at {host}:{port}")
        elif persist_directory:
            # Persistent local mode
            Path(persist_directory).mkdir(parents=True, exist_ok=True)
            self.client = chromadb.PersistentClient(path=persist_directory)
            logger.info(f"Initialized persistent Chroma at {persist_directory}")
        else:
            # In-memory mode
            self.client = chromadb.Client()
            logger.info("Initialized in-memory Chroma")
        
        # Get or create collection
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": distance_metric}
        )
    
    def add_documents(self, documents: List[Document]) -> List[str]:
        """Add documents to Chroma."""
        if not documents:
            return []
        
        ids = [doc.id for doc in documents]
        contents = [doc.content for doc in documents]
        metadatas = [doc.metadata for doc in documents]
        embeddings = [doc.embedding for doc in documents if doc.embedding]
        
        add_kwargs = {
            "ids": ids,
            "documents": contents,
            "metadatas": metadatas
        }
        
        if embeddings and len(embeddings) == len(documents):
            add_kwargs["embeddings"] = embeddings
        
        self.collection.add(**add_kwargs)
        logger.info(f"Added {len(documents)} documents to Chroma")
        
        return ids
    
    def search(
        self,
        query_embedding: List[float],
        top_k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """Search Chroma for similar documents."""
        query_kwargs = {
            "query_embeddings": [query_embedding],
            "n_results": top_k,
            "include": ["documents", "metadatas", "distances", "embeddings"]
        }
        
        if filter_metadata:
            query_kwargs["where"] = filter_metadata
        
        results = self.collection.query(**query_kwargs)
        
        search_results = []
        if results["ids"] and results["ids"][0]:
            for i, doc_id in enumerate(results["ids"][0]):
                doc = Document(
                    id=doc_id,
                    content=results["documents"][0][i] if results["documents"] else "",
                    metadata=results["metadatas"][0][i] if results["metadatas"] else {},
                    embedding=results["embeddings"][0][i] if results.get("embeddings") else None
                )
                # Chroma returns distances, convert to similarity score
                distance = results["distances"][0][i] if results["distances"] else 0
                score = 1 - distance if self.distance_metric == "cosine" else 1 / (1 + distance)
                search_results.append(SearchResult(document=doc, score=score))
        
        return search_results
    
    def delete(self, document_ids: List[str]) -> None:
        """Delete documents from Chroma."""
        if document_ids:
            self.collection.delete(ids=document_ids)
            logger.info(f"Deleted {len(document_ids)} documents from Chroma")
    
    def clear(self) -> None:
        """Clear all documents from collection."""
        self.client.delete_collection(self.collection_name)
        self.collection = self.client.create_collection(
            name=self.collection_name,
            metadata={"hnsw:space": self.distance_metric}
        )
        logger.info("Cleared Chroma collection")
    
    def count(self) -> int:
        """Get document count."""
        return self.collection.count()


class QdrantStore(VectorStore):
    """Qdrant vector store implementation."""
    
    def __init__(
        self,
        collection_name: str = "readanybook",
        host: str = "localhost",
        port: int = 6333,
        grpc_port: int = 6334,
        vector_size: int = 384,
        distance_metric: str = "cosine",
        persist_directory: Optional[str] = None
    ):
        """
        Initialize Qdrant store.
        
        Args:
            collection_name: Name of the collection
            host: Qdrant server host
            port: Qdrant REST API port
            grpc_port: Qdrant gRPC port
            vector_size: Dimension of vectors
            distance_metric: Distance metric
            persist_directory: Path for local storage (if no server)
        """
        try:
            from qdrant_client import QdrantClient
            from qdrant_client.models import Distance, VectorParams
        except ImportError:
            raise ImportError("qdrant-client is required. Install with: pip install qdrant-client")
        
        self.collection_name = collection_name
        self.vector_size = vector_size
        
        # Map distance metric
        distance_map = {
            "cosine": Distance.COSINE,
            "l2": Distance.EUCLID,
            "ip": Distance.DOT
        }
        self.distance = distance_map.get(distance_metric, Distance.COSINE)
        
        # Initialize client
        if persist_directory:
            Path(persist_directory).mkdir(parents=True, exist_ok=True)
            self.client = QdrantClient(path=persist_directory)
            logger.info(f"Initialized local Qdrant at {persist_directory}")
        else:
            self.client = QdrantClient(host=host, port=port, grpc_port=grpc_port)
            logger.info(f"Connected to Qdrant server at {host}:{port}")
        
        # Create collection if not exists
        collections = self.client.get_collections().collections
        if not any(c.name == collection_name for c in collections):
            self.client.create_collection(
                collection_name=collection_name,
                vectors_config=VectorParams(size=vector_size, distance=self.distance)
            )
            logger.info(f"Created Qdrant collection: {collection_name}")
    
    def add_documents(self, documents: List[Document]) -> List[str]:
        """Add documents to Qdrant."""
        if not documents:
            return []
        
        try:
            from qdrant_client.models import PointStruct
        except ImportError:
            raise ImportError("qdrant-client is required")
        
        points = []
        for doc in documents:
            if doc.embedding is None:
                continue
            
            payload = {
                "content": doc.content,
                **doc.metadata
            }
            
            points.append(PointStruct(
                id=doc.id,
                vector=doc.embedding,
                payload=payload
            ))
        
        if points:
            self.client.upsert(collection_name=self.collection_name, points=points)
            logger.info(f"Added {len(points)} documents to Qdrant")
        
        return [doc.id for doc in documents]
    
    def search(
        self,
        query_embedding: List[float],
        top_k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """Search Qdrant for similar documents."""
        search_params = {
            "collection_name": self.collection_name,
            "query_vector": query_embedding,
            "limit": top_k,
            "with_payload": True,
            "with_vectors": True
        }
        
        if filter_metadata:
            from qdrant_client.models import Filter, FieldCondition, MatchValue
            conditions = [
                FieldCondition(key=k, match=MatchValue(value=v))
                for k, v in filter_metadata.items()
            ]
            search_params["query_filter"] = Filter(must=conditions)
        
        results = self.client.search(**search_params)
        
        search_results = []
        for hit in results:
            content = hit.payload.pop("content", "")
            doc = Document(
                id=str(hit.id),
                content=content,
                metadata=hit.payload,
                embedding=hit.vector
            )
            search_results.append(SearchResult(document=doc, score=hit.score))
        
        return search_results
    
    def delete(self, document_ids: List[str]) -> None:
        """Delete documents from Qdrant."""
        if document_ids:
            from qdrant_client.models import PointIdsList
            self.client.delete(
                collection_name=self.collection_name,
                points_selector=PointIdsList(points=document_ids)
            )
            logger.info(f"Deleted {len(document_ids)} documents from Qdrant")
    
    def clear(self) -> None:
        """Clear all documents from collection."""
        from qdrant_client.models import VectorParams
        
        self.client.delete_collection(self.collection_name)
        self.client.create_collection(
            collection_name=self.collection_name,
            vectors_config=VectorParams(size=self.vector_size, distance=self.distance)
        )
        logger.info("Cleared Qdrant collection")
    
    def count(self) -> int:
        """Get document count."""
        info = self.client.get_collection(self.collection_name)
        return info.points_count


class WeaviateStore(VectorStore):
    """Weaviate vector store implementation."""
    
    def __init__(
        self,
        collection_name: str = "ReadAnyBook",
        host: str = "localhost",
        port: int = 8080,
        distance_metric: str = "cosine"
    ):
        """
        Initialize Weaviate store.
        
        Args:
            collection_name: Name of the class/collection
            host: Weaviate server host
            port: Weaviate server port
            distance_metric: Distance metric
        """
        try:
            import weaviate
        except ImportError:
            raise ImportError("weaviate-client is required. Install with: pip install weaviate-client")
        
        self.collection_name = collection_name
        
        self.client = weaviate.Client(f"http://{host}:{port}")
        logger.info(f"Connected to Weaviate at {host}:{port}")
        
        # Create schema if not exists
        if not self.client.schema.exists(collection_name):
            class_schema = {
                "class": collection_name,
                "vectorizer": "none",
                "properties": [
                    {"name": "content", "dataType": ["text"]},
                    {"name": "doc_id", "dataType": ["string"]},
                    {"name": "metadata", "dataType": ["text"]}
                ]
            }
            self.client.schema.create_class(class_schema)
            logger.info(f"Created Weaviate class: {collection_name}")
    
    def add_documents(self, documents: List[Document]) -> List[str]:
        """Add documents to Weaviate."""
        import json
        
        if not documents:
            return []
        
        with self.client.batch as batch:
            for doc in documents:
                if doc.embedding is None:
                    continue
                
                properties = {
                    "content": doc.content,
                    "doc_id": doc.id,
                    "metadata": json.dumps(doc.metadata)
                }
                
                batch.add_data_object(
                    properties,
                    self.collection_name,
                    vector=doc.embedding,
                    uuid=doc.id
                )
        
        logger.info(f"Added {len(documents)} documents to Weaviate")
        return [doc.id for doc in documents]
    
    def search(
        self,
        query_embedding: List[float],
        top_k: int = 10,
        filter_metadata: Optional[Dict[str, Any]] = None
    ) -> List[SearchResult]:
        """Search Weaviate for similar documents."""
        import json
        
        query = (
            self.client.query
            .get(self.collection_name, ["content", "doc_id", "metadata"])
            .with_near_vector({"vector": query_embedding})
            .with_limit(top_k)
            .with_additional(["certainty", "vector"])
        )
        
        results = query.do()
        
        search_results = []
        if results.get("data", {}).get("Get", {}).get(self.collection_name):
            for item in results["data"]["Get"][self.collection_name]:
                metadata = json.loads(item.get("metadata", "{}"))
                doc = Document(
                    id=item.get("doc_id", ""),
                    content=item.get("content", ""),
                    metadata=metadata,
                    embedding=item.get("_additional", {}).get("vector")
                )
                score = item.get("_additional", {}).get("certainty", 0)
                search_results.append(SearchResult(document=doc, score=score))
        
        return search_results
    
    def delete(self, document_ids: List[str]) -> None:
        """Delete documents from Weaviate."""
        for doc_id in document_ids:
            self.client.data_object.delete(doc_id, self.collection_name)
        logger.info(f"Deleted {len(document_ids)} documents from Weaviate")
    
    def clear(self) -> None:
        """Clear all documents from collection."""
        self.client.schema.delete_class(self.collection_name)
        # Recreate class
        class_schema = {
            "class": self.collection_name,
            "vectorizer": "none",
            "properties": [
                {"name": "content", "dataType": ["text"]},
                {"name": "doc_id", "dataType": ["string"]},
                {"name": "metadata", "dataType": ["text"]}
            ]
        }
        self.client.schema.create_class(class_schema)
        logger.info("Cleared Weaviate collection")
    
    def count(self) -> int:
        """Get document count."""
        result = (
            self.client.query
            .aggregate(self.collection_name)
            .with_meta_count()
            .do()
        )
        return result.get("data", {}).get("Aggregate", {}).get(self.collection_name, [{}])[0].get("meta", {}).get("count", 0)


def create_vector_store(settings) -> VectorStore:
    """
    Factory function to create a vector store from settings.
    
    Args:
        settings: VectorDBSettings or Settings object
        
    Returns:
        VectorStore instance
    """
    if hasattr(settings, "vectordb"):
        vdb_settings = settings.vectordb
    else:
        vdb_settings = settings
    
    backend = vdb_settings.backend
    
    if backend == "chroma":
        return ChromaStore(
            collection_name=vdb_settings.collection_name,
            persist_directory=vdb_settings.persist_directory,
            host=vdb_settings.chroma.host,
            port=vdb_settings.chroma.port,
            distance_metric=vdb_settings.distance_metric
        )
    elif backend == "qdrant":
        return QdrantStore(
            collection_name=vdb_settings.collection_name,
            host=vdb_settings.qdrant.host,
            port=vdb_settings.qdrant.port,
            grpc_port=vdb_settings.qdrant.grpc_port,
            distance_metric=vdb_settings.distance_metric,
            persist_directory=vdb_settings.persist_directory
        )
    elif backend == "weaviate":
        return WeaviateStore(
            collection_name=vdb_settings.collection_name,
            host=vdb_settings.weaviate.host,
            port=vdb_settings.weaviate.port,
            distance_metric=vdb_settings.distance_metric
        )
    else:
        raise ValueError(f"Unknown vector store backend: {backend}")
