from floki.storage.vectorstores import VectorStoreBase
from floki.document.embedder import SentenceTransformerEmbedder
from floki.document.embedder.base import EmbedderBase
from typing import List, Dict, Optional, Iterable, Any, Literal, Union
from pydantic import Field, ConfigDict
import uuid
import logging

logger = logging.getLogger(__name__)


class PostgresVectorStore(VectorStoreBase):
    """
    A PostgreSQL-based vector store implementation leveraging pgvector for similarity search.
    Supports automatic embedding generation, metadata filtering, and flexible indexing.
    """

    connection_string: str = Field(..., description="PostgreSQL connection string.")
    table_name: str = Field("vector_store", description="The table name for storing vectors.")
    embedding_dim: Optional[int] = Field(None, description="Fixed dimensionality of embedding vectors (optional).")
    embedding_function: Optional[EmbedderBase] = Field(default_factory=SentenceTransformerEmbedder, description="Embedding function for embedding generation.")
    pool_config: Dict = Field(default_factory=lambda: {"min_size": 1, "max_size": 10, "timeout": 30}, description="Connection pool settings.")
    index_type: Literal["hnsw", "ivfflat", "flat"] = Field("ivfflat", description="The type of index to use for vector search.")
    index_params: Dict = Field(default_factory=lambda: {"lists": 100, "m": 16, "ef_construction": 64}, description="Parameters for the vector index.")

    pool: Optional[Any] = Field(default=None, init=False, description="Connection pool for PostgreSQL.")
    tracked_dimensions: set = Field(default_factory=set, init=False, description="Set of tracked vector dimensions for partial indexing.")
    
    model_config = ConfigDict(arbitrary_types_allowed=True)

    def model_post_init(self, __context: Any) -> None:
        """
        Initializes the PostgreSQL connection pool, ensures table and initial index setup.
        """
        super().model_post_init(__context)

        try:
            from psycopg_pool import ConnectionPool
            from pgvector.psycopg import register_vector
        except ImportError as e:
            raise ImportError(
                "The psycopg, psycopg-pool, and pgvector libraries are required to use this store. "
                "Install them using pip install 'psycopg[binary,pool]' pgvector"
            ) from e

        try:
            self.pool: ConnectionPool = ConnectionPool(self.connection_string, **self.pool_config)
            with self.pool.connection() as conn:
                # Enable the pgvector extension if not already enabled
                with conn.cursor() as cursor:
                    cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")
                    logger.info("pgvector extension ensured in the database.")

                # Register vector type with psycopg
                register_vector(conn)

                # Create the table
                with conn.cursor() as cursor:
                    embedding_column = f"VECTOR({self.embedding_dim})" if self.embedding_dim else "VECTOR"
                    cursor.execute(f"""
                        CREATE TABLE IF NOT EXISTS {self.table_name} (
                            id UUID PRIMARY KEY,
                            document TEXT,
                            metadata JSONB,
                            embedding {embedding_column}
                        );
                    """)
                    logger.info(f"Table '{self.table_name}' ensured.")

                # Create global index for fixed-dimension embeddings
                if self.embedding_dim:
                    with conn.cursor() as cursor:
                        cursor.execute(
                            f"""
                            CREATE INDEX IF NOT EXISTS {self.table_name}_embedding_idx
                            ON {self.table_name} USING {self.index_type} (embedding vector_cosine_ops)
                            WITH ({", ".join(f"{k}={v}" for k, v in self.index_params.items())});
                            """
                        )
                        logger.info(f"Global index created for fixed-dimension embeddings.")
                else:
                    logger.info(f"No fixed dimension specified; relying on dynamic partial indexing.")
        except Exception as e:
            logger.error(f"Failed to initialize PostgresVectorStore: {e}")
            raise

    def _ensure_partial_index(self, dimension: int):
        """
        Creates a partial index for embeddings with the specified dimension if it doesn't already exist.
        """
        if dimension in self.tracked_dimensions:
            return

        try:
            with self.pool.connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(
                        f"""
                        CREATE INDEX IF NOT EXISTS {self.table_name}_embedding_{dimension}_idx
                        ON {self.table_name} USING hnsw ((embedding::vector({dimension})) vector_cosine_ops)
                        WHERE vector_dims(embedding) = {dimension};
                        """
                    )
                    logger.info(f"Partial index created for embedding dimension {dimension}.")
            self.tracked_dimensions.add(dimension)
        except Exception as e:
            logger.error(f"Failed to create partial index for dimension {dimension}: {e}")

    def add(
        self,
        documents: Iterable[str],
        embeddings: Optional[List[List[float]]] = None,
        metadatas: Optional[List[Dict]] = None,
        ids: Optional[List[str]] = None,
        upsert: bool = False
    ) -> List[str]:
        """
        Adds or upserts documents into the vector store.

        Args:
            documents (Iterable[str]): The documents to add or upsert.
            embeddings (Optional[List[List[float]]]): Precomputed embeddings.
            metadatas (Optional[List[Dict]]): Metadata for each document.
            ids (Optional[List[str]]): Unique IDs for the documents.
            upsert (bool): Whether to update existing records on conflict. Defaults to False.

        Returns:
            List[str]: List of IDs for the added or upserted documents.
        """

        try:
            from psycopg.types.json import Jsonb
        except ImportError as e:
            raise ImportError("Required library 'psycopg' is missing.") from e
        
        try:
            if embeddings is None:
                embeddings = self.embedding_function(list(documents))
                logger.info("Generated embeddings using the provided embedding function.")

            if ids is None:
                ids = [str(uuid.uuid4()) for _ in documents]

            on_conflict_action = (
                """
                DO UPDATE SET
                    document = EXCLUDED.document,
                    metadata = EXCLUDED.metadata,
                    embedding = EXCLUDED.embedding
                """
                if upsert
                else "DO NOTHING"
            )

            with self.pool.connection() as conn:
                with conn.cursor() as cursor:
                    for i, doc in enumerate(documents):
                        dimension = len(embeddings[i])
                        if not self.embedding_dim:
                            self._ensure_partial_index(dimension)

                        metadata = Jsonb(metadatas[i]) if metadatas else Jsonb({})
                        query = f"""
                        INSERT INTO {self.table_name} (id, document, metadata, embedding)
                        VALUES (%s, %s, %s, %s)
                        ON CONFLICT (id) {on_conflict_action};
                        """
                        cursor.execute(
                            query,
                            (ids[i], doc, metadata, embeddings[i])
                        )
            logger.info(f"{'Upserted' if upsert else 'Added'} {len(documents)} documents.")
            return ids
        except Exception as e:
            logger.error(f"Failed to {'upsert' if upsert else 'add'} documents: {e}")
            raise
    
    def update(
        self,
        ids: List[str],
        embeddings: Optional[List[List[float]]] = None,
        metadatas: Optional[List[Dict]] = None,
        documents: Optional[List[str]] = None,
    ) -> None:
        """
        Updates existing documents in the vector store.

        Args:
            ids (List[str]): A list of document IDs to update.
            embeddings (Optional[List[List[float]]]): The new embedding vectors for the documents.
            metadatas (Optional[List[Dict]]): The new metadata for the documents.
            documents (Optional[List[str]]): The new text for the documents.

        Returns:
            None
        """
        try:
            from psycopg.types.json import Jsonb
        except ImportError as e:
            raise ImportError("Required library 'psycopg' is missing.") from e
        
        try:

            if not any([embeddings, metadatas, documents]):
                raise ValueError("At least one of embeddings, metadatas, or documents must be provided for update.")

            with self.pool.connection() as conn:
                with conn.cursor() as cursor:
                    for i, doc_id in enumerate(ids):
                        update_fields = []
                        params = []

                        if embeddings and embeddings[i]:
                            dimension = len(embeddings[i])
                            if not self.embedding_dim:
                                self._ensure_partial_index(dimension)
                            update_fields.append("embedding = %s")
                            params.append(embeddings[i])

                        if documents and documents[i]:
                            update_fields.append("document = %s")
                            params.append(documents[i])

                        if metadatas and metadatas[i]:
                            update_fields.append("metadata = %s")
                            params.append(Jsonb(metadatas[i]))

                        if not update_fields:
                            continue

                        params.append(doc_id)  # Add the ID for the WHERE clause

                        query = f"""
                        UPDATE {self.table_name}
                        SET {", ".join(update_fields)}
                        WHERE id = %s;
                        """
                        cursor.execute(query, params)

            logger.info(f"Updated {len(ids)} documents in {self.table_name}.")
        except Exception as e:
            logger.error(f"Failed to update documents: {e}")
            raise

    def delete(self, ids: List[str]) -> bool:
        """
        Deletes documents from the vector store by their IDs.

        Args:
            ids (List[str]): List of document IDs to delete.

        Returns:
            bool: True if the deletion was successful, False otherwise.
        """
        try:
            with self.pool.connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(
                        f"DELETE FROM {self.table_name} WHERE id = ANY(%s)", (ids,)
                    )
            logger.info(f"Deleted {len(ids)} documents.")
            return True
        except Exception as e:
            logger.error(f"Failed to delete documents: {e}")
            return False
    
    def get(self, ids: Optional[List[str]] = None, with_embedding: bool = False) -> List[Dict]:
        """
        Retrieves items from the PostgreSQL vector store by IDs. If no IDs are provided, retrieves all items.

        Args:
            ids (Optional[List[str]]): The IDs of the items to retrieve. If None, retrieves all items.
            with_embedding (bool): Whether to include embeddings in the retrieved data.

        Returns:
            List[Dict]: A list of dictionaries containing the metadata, documents, and optionally embeddings of the retrieved items.
        """
        try:
            with self.pool.connection() as conn:
                with conn.cursor() as cursor:
                    if ids:
                        query = "SELECT id, document, metadata"
                        if with_embedding:
                            query += ", embedding"
                        query += f" FROM {self.table_name} WHERE id = ANY(%s)"
                        cursor.execute(query, (ids,))
                    else:
                        query = "SELECT id, document, metadata"
                        if with_embedding:
                            query += ", embedding"
                        query += f" FROM {self.table_name}"
                        cursor.execute(query)

                    # Get column names for row-to-dict conversion
                    colnames = [desc[0] for desc in cursor.description]
                    rows = cursor.fetchall()

            return [
                dict(zip(colnames, row))
                for row in rows
            ]
        except Exception as e:
            logger.error(f"Failed to retrieve documents: {e}")
            raise
    
    def search_similar(
        self,
        query_texts: Optional[Union[str, List[str]]] = None,
        query_embeddings: Optional[Union[List[float], List[List[float]]]] = None,
        k: int = 4,
        distance_metric: str = "cosine",
        metadata_filter: Optional[Dict] = None,
    ) -> List[Dict]:
        """
        Perform a similarity search in the vector store with optional metadata filtering.

        Args:
            query_texts (Optional[Union[str, List[str]]]): Text queries to embed and search. 
                If provided, embeddings are generated using the configured embedding function.
            query_embeddings (Optional[Union[List[float], List[List[float]]]]): Precomputed embeddings 
                for similarity search. Provide either `query_texts` or `query_embeddings`, not both.
            k (int): Number of top results to return. Defaults to 4.
            distance_metric (str): Distance metric to use for similarity computation. 
                Supported values are:
                    - "cosine": Cosine similarity.
                    - "l2": Euclidean distance.
                    - "inner_product": Inner product similarity.
                Defaults to "cosine".
            metadata_filter (Optional[Dict]): Metadata conditions for filtering results. 
                Keys are metadata fields, and values are the required values for those fields.

        Returns:
            List[Dict]: A list of dictionaries, each representing a search result. Each dictionary contains:
                - "id" (str): The ID of the matched item.
                - "document" (str): The document content of the matched item.
                - "metadata" (Dict): Metadata associated with the matched item.
                - "similarity" (float): The similarity score of the matched item.

        Raises:
            ValueError: If neither `query_texts` nor `query_embeddings` is provided, or if both are provided.
            Exception: For any issues during the database query execution.
        """
        try:
            # Validate inputs
            if not query_texts and not query_embeddings:
                raise ValueError("Either `query_texts` or `query_embeddings` must be provided.")
            if query_texts and query_embeddings:
                raise ValueError("Provide either `query_texts` or `query_embeddings`, not both.")

            # Generate embeddings if query_texts is provided
            if query_texts:
                if isinstance(query_texts, str):
                    query_texts = [query_texts]
                query_embeddings = self.embedding_function(query_texts)

            # Handle single embedding case
            if isinstance(query_embeddings[0], (int, float)):
                query_embeddings = [query_embeddings]

            # Map distance metrics to PostgreSQL operators
            operator_map = {
                "cosine": "<=>",  # Cosine similarity
                "l2": "<->",      # Euclidean distance
                "inner_product": "<#>",  # Inner product
            }
            operator = operator_map.get(distance_metric, "<=>")

            results = []
            with self.pool.connection() as conn:
                # Use a compatible row factory or default cursor
                with conn.cursor() as cursor:
                    for embedding in query_embeddings:
                        # Convert the embedding to a PostgreSQL-compatible vector format
                        embedding_vector = f"ARRAY{embedding}::vector"

                        # Base query
                        query = f"""
                        SELECT id, document, metadata, 1 - (embedding {operator} {embedding_vector}) AS similarity
                        FROM {self.table_name}
                        """

                        # Add metadata filtering conditions if provided
                        params = []
                        if metadata_filter:
                            filter_conditions = " AND ".join(
                                [f"metadata ->> %s = %s" for _ in metadata_filter]
                            )
                            query += f" WHERE {filter_conditions}"
                            params.extend(sum(([k, v] for k, v in metadata_filter.items()), []))

                        # Add ordering and limit
                        query += " ORDER BY similarity DESC LIMIT %s"
                        params.append(k)

                        # Execute the query
                        cursor.execute(query, tuple(params))
                        for row in cursor.fetchall():
                            results.append(row)

            # Format results for consistency
            formatted_results = [
                {
                    "id": row[0],
                    "document": row[1],
                    "metadata": row[2],
                    "similarity": row[3],
                }
                for row in results
            ]

            return formatted_results
        except Exception as e:
            logger.error(f"Failed to perform similarity search: {e}")
            raise

    def reset(self):
        """
        Resets the vector store by truncating the table, deleting all stored data.

        Returns:
            None
        """
        try:
            with self.pool.connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(f"TRUNCATE TABLE {self.table_name}")
            logger.info(f"Table '{self.table_name}' reset.")
        except Exception as e:
            logger.error(f"Failed to reset table: {e}")
            raise
    
    def count(self) -> int:
        """
        Counts the number of documents in the vector store.

        Returns:
            int: The total number of documents in the store.
        """
        try:
            with self.pool.connection() as conn:
                with conn.cursor() as cursor:
                    cursor.execute(f"SELECT COUNT(*) FROM {self.table_name}")
                    count = cursor.fetchone()[0]
            return count
        except Exception as e:
            logger.error(f"Failed to count documents in the store: {e}")
            raise