# Copyright (c) 2025 Oracle and/or its affiliates.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License, version 2.0, as
# published by the Free Software Foundation.
#
# This program is designed to work with certain software (including
# but not limited to OpenSSL) that is licensed under separate terms,
# as designated in a particular file or component or in included license
# documentation. The authors of MySQL hereby grant you an
# additional permission to link the program and your derivative works
# with the separately licensed software that they have either included with
# the program or referenced in the documentation.
#
# Without limiting anything contained in the foregoing, this file,
# which is part of MySQL Connector/Python, is also subject to the
# Universal FOSS Exception, version 1.0, a copy of which can be found at
# http://oss.oracle.com/licenses/universal-foss-exception.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU General Public License, version 2.0, for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin St, Fifth Floor, Boston, MA 02110-1301  USA

"""MySQL-backed vector store for embeddings and semantic document retrieval.

Provides a VectorStore implementation persisting documents, metadata, and
embeddings in MySQL, plus similarity search utilities.
"""

import json

from typing import Any, Iterable, List, Optional, Sequence, Union

import pandas as pd

from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import VectorStore
from pydantic import PrivateAttr

from mysql.ai.genai.embedding import MyEmbeddings
from mysql.ai.utils import (
    VAR_NAME_SPACE,
    atomic_transaction,
    delete_sql_table,
    execute_sql,
    extend_sql_table,
    format_value_sql,
    get_random_name,
    is_table_empty,
    source_schema,
    table_exists,
)
from mysql.connector.abstracts import MySQLConnectionAbstract

BASIC_EMBEDDING_QUERY = "Hello world!"
EMBEDDING_SOURCE = "external_source"

VAR_EMBEDDING = f"{VAR_NAME_SPACE}.embedding"
VAR_CONTEXT = f"{VAR_NAME_SPACE}.context"
VAR_CONTEXT_MAP = f"{VAR_NAME_SPACE}.context_map"
VAR_RETRIEVAL_INFO = f"{VAR_NAME_SPACE}.retrieval_info"
VAR_OPTIONS = f"{VAR_NAME_SPACE}.options"

ID_SPACE = "internal_ai_id_"


class MyVectorStore(VectorStore):
    """
    MySQL-backed vector store for handling embeddings and semantic document retrieval.

    Supports adding, deleting, and searching high-dimensional vector representations
    of documents using efficient storage and HeatWave ML similarity search procedures.

    Supports use as a context manager: when used in a `with` statement, all backing
    tables/data are deleted automatically when the block exits (even on exception).

    Attributes:
        db_connection (MySQLConnectionAbstract): Active MySQL database connection.
        embedder (Embeddings): Embeddings generator for computing vector representations.
        schema_name (str): SQL schema for table storage.
        table_name (Optional[str]): Name of the active table backing the store
            (or None until created).
        embedding_dimension (int): Size of embedding vectors stored.
        next_id (int): Internal counter for unique document ID generation.
    """

    _db_connection: MySQLConnectionAbstract = PrivateAttr()
    _embedder: Embeddings = PrivateAttr()
    _schema_name: str = PrivateAttr()
    _table_name: Optional[str] = PrivateAttr()
    _embedding_dimension: int = PrivateAttr()
    _next_id: int = PrivateAttr()

    def __init__(
        self,
        db_connection: MySQLConnectionAbstract,
        embedder: Optional[Embeddings] = None,
    ) -> None:
        """
        Initialize a MyVectorStore with a database connection and embedding generator.

        Args:
            db_connection: MySQL database connection for all vector operations.
            embedder: Embeddings generator used for creating and querying embeddings.

        Raises:
            ValueError: If the schema name is not valid
            DatabaseError:
                If a database connection issue occurs.
                If an operational error occurs during execution.
        """
        super().__init__()
        self._next_id = 0

        self._schema_name = source_schema(db_connection)
        self._embedder = embedder or MyEmbeddings(db_connection)
        self._db_connection = db_connection
        self._table_name: Optional[str] = None

        # Embedding dimension determined using an example call.
        # Assumes embeddings have fixed length.
        self._embedding_dimension = len(
            self._embedder.embed_query(BASIC_EMBEDDING_QUERY)
        )

    def _get_ids(self, num_ids: int) -> list[str]:
        """
        Generate a batch of unique internal document IDs for vector storage.

        Args:
            num_ids: Number of IDs to create.

        Returns:
            List of sequentially numbered internal string IDs.
        """
        ids = [
            f"internal_ai_id_{i}" for i in range(self._next_id, self._next_id + num_ids)
        ]
        self._next_id += num_ids
        return ids

    def _make_vector_store(self) -> None:
        """
        Create a backing SQL table for storing vectors if not already created.

        Returns:
            None

        Raises:
            DatabaseError:
                If a database connection issue occurs.
                If an operational error occurs during execution.

        Notes:
            The table name is randomized to avoid collisions.
            Schema includes content, metadata, and embedding vector.
        """
        if self._table_name is None:
            with atomic_transaction(self._db_connection) as cursor:
                table_name = get_random_name(
                    lambda table_name: not table_exists(
                        cursor, self._schema_name, table_name
                    )
                )

                create_table_stmt = f"""
                CREATE TABLE {self._schema_name}.{table_name} (
                    `id` VARCHAR(128) NOT NULL,
                    `content` TEXT,
                    `metadata` JSON DEFAULT NULL,
                    `embed` vector(%s),
                    PRIMARY KEY (`id`)
                ) ENGINE=InnoDB;
                """
                execute_sql(
                    cursor, create_table_stmt, params=(self._embedding_dimension,)
                )

                self._table_name = table_name

    def delete(self, ids: Optional[Sequence[str]] = None, **_: Any) -> None:
        """
        Delete documents by ID. Optionally deletes the vector table if empty after deletions.

        Args:
            ids: Optional sequence of document IDs to delete. If None, no action is taken.

        Returns:
            None

        Raises:
            DatabaseError:
                If a database connection issue occurs.
                If an operational error occurs during execution.

        Notes:
            If the backing table is empty after deletions, the table is dropped and
            table_name is set to None.
        """
        with atomic_transaction(self._db_connection) as cursor:
            if ids:
                for _id in ids:
                    execute_sql(
                        cursor,
                        f"DELETE FROM {self._schema_name}.{self._table_name} WHERE id = %s",
                        params=(_id,),
                    )

            if is_table_empty(cursor, self._schema_name, self._table_name):
                self.delete_all()

    def delete_all(self) -> None:
        """
        Delete and drop the entire vector store table.

        Returns:
            None
        """
        if self._table_name is not None:
            with atomic_transaction(self._db_connection) as cursor:
                delete_sql_table(cursor, self._schema_name, self._table_name)
                self._table_name = None

    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: Optional[list[dict]] = None,
        ids: Optional[List[str]] = None,
        **_: dict,
    ) -> List[str]:
        """
        Add a batch of text strings and corresponding metadata to the vector store.

        Args:
            texts: List of strings to embed and store.
            metadatas: Optional list of metadata dicts (one per text).
            ids: Optional custom document IDs.

        Returns:
            List of document IDs corresponding to the added texts.

        Raises:
            DatabaseError:
                If a database connection issue occurs.
                If an operational error occurs during execution.

        Notes:
            If metadatas is None, an empty dict is assigned to each document.
        """
        texts = list(texts)

        documents = [
            Document(page_content=text, metadata=meta)
            for text, meta in zip(texts, metadatas or [{}] * len(texts))
        ]
        return self.add_documents(documents, ids=ids)

    @classmethod
    def from_texts(
        cls,
        texts: Iterable[str],
        embedder: Embeddings,
        metadatas: Optional[list[dict]] = None,
        db_connection: MySQLConnectionAbstract = None,
    ) -> VectorStore:
        """
        Construct and populate a MyVectorStore instance from raw texts and metadata.

        Args:
            texts: List of strings to vectorize and store.
            embedder: Embeddings generator to use.
            metadatas: Optional list of metadata dicts per text.
            db_connection: Active MySQL connection.

        Returns:
            Instance of MyVectorStore containing the added texts.

        Raises:
            ValueError: If db_connection is not provided.
            DatabaseError:
                If a database connection issue occurs.
                If an operational error occurs during execution.
        """
        if db_connection is None:
            raise ValueError(
                "db_connection must be specified to create a MyVectorStore object"
            )

        texts = list(texts)

        instance = cls(db_connection=db_connection, embedder=embedder)
        instance.add_texts(texts, metadatas=metadatas)

        return instance

    def add_documents(
        self, documents: list[Document], ids: list[str] = None
    ) -> list[str]:
        """
        Embed and store Document objects as high-dimensional vectors with metadata.

        Args:
            documents: List of Document objects (each with 'page_content' and 'metadata').
            ids: Optional list of explicit document IDs. Must match the length of documents.

        Returns:
            List of document IDs stored.

        Raises:
            ValueError: If provided IDs do not match the number of documents.
            DatabaseError:
                If a database connection issue occurs.
                If an operational error occurs during execution.

        Notes:
            Automatically creates the backing table if it does not exist.
        """
        if ids and len(ids) != len(documents):
            msg = (
                "ids must be the same length as documents. "
                f"Got {len(ids)} ids and {len(documents)} documents."
            )
            raise ValueError(msg)

        if len(documents) > 0:
            self._make_vector_store()
        else:
            return []

        if ids is None:
            ids = self._get_ids(len(documents))

        content = [doc.page_content for doc in documents]
        vectors = self._embedder.embed_documents(content)

        df = pd.DataFrame()
        df["id"] = ids
        df["content"] = content
        df["embed"] = vectors
        df["metadata"] = [doc.metadata for doc in documents]

        with atomic_transaction(self._db_connection) as cursor:
            extend_sql_table(
                cursor,
                self._schema_name,
                self._table_name,
                df,
                col_name_to_placeholder_string={"embed": "string_to_vector(%s)"},
            )

        return ids

    def similarity_search(
        self,
        query: str,
        k: int = 3,
        **kwargs: Any,
    ) -> list[Document]:
        """
        Search for and return the most similar documents in the store to the given query.

        Args:
            query: String query to embed and use for similarity search.
            k: Number of top documents to return.
            kwargs: options to pass to ML_SIMILARITY_SEARCH. Currently supports
                distance_metric, max_distance, percentage_distance, and segment_overlap

        Returns:
            List of Document objects, ordered from most to least similar.

        Raises:
            DatabaseError:
                If provided kwargs are invalid or unsupported.
                If a database connection issue occurs.
                If an operational error occurs during execution.

        Implementation Notes:
            - Calls ML similarity search within MySQL using stored procedures.
            - Retrieves IDs, content, and metadata for search matches.
            - Parsing and retrieval for context results are handled via intermediate JSONs.
        """
        if self._table_name is None:
            return []

        embedding = self._embedder.embed_query(query)

        with atomic_transaction(self._db_connection) as cursor:
            # Set the embedding variable for the similarity search SP
            execute_sql(
                cursor,
                f"SET @{VAR_EMBEDDING} = string_to_vector(%s)",
                params=[str(embedding)],
            )

            distance_metric = kwargs.get("distance_metric", "COSINE")
            retrieval_options = {
                "max_distance": kwargs.get("max_distance", 0.6),
                "percentage_distance": kwargs.get("percentage_distance", 20.0),
                "segment_overlap": kwargs.get("segment_overlap", 0),
            }

            retrieval_options_placeholder, retrieval_options_params = format_value_sql(
                retrieval_options
            )
            similarity_search_query = f"""
            CALL sys.ML_SIMILARITY_SEARCH(
                @{VAR_EMBEDDING},
                JSON_ARRAY(
                    '{self._schema_name}.{self._table_name}'
                ),
                JSON_OBJECT(
                    "segment", "content",
                    "segment_embedding", "embed",
                    "document_name", "id"
                ),
                {k},
                %s,
                NULL,
                NULL,
                {retrieval_options_placeholder},
                @{VAR_CONTEXT},
                @{VAR_CONTEXT_MAP},
                @{VAR_RETRIEVAL_INFO}
            )
            """

            execute_sql(
                cursor,
                similarity_search_query,
                params=[distance_metric, *retrieval_options_params],
            )
            execute_sql(cursor, f"SELECT @{VAR_CONTEXT_MAP}")

            results = []

            context_maps = json.loads(cursor.fetchone()[0])
            for context in context_maps:
                execute_sql(
                    cursor,
                    (
                        "SELECT id, content, metadata "
                        f"FROM {self._schema_name}.{self._table_name} "
                        "WHERE id = %s"
                    ),
                    params=(context["document_name"],),
                )
                doc_id, content, metadata = cursor.fetchone()

                doc_args = {
                    "id": doc_id,
                    "page_content": content,
                }
                if metadata is not None:
                    doc_args["metadata"] = json.loads(metadata)

                doc = Document(**doc_args)
                results.append(doc)

            return results

    def __enter__(self) -> "VectorStore":
        """
        Enter the runtime context related to this vector store instance.

        Returns:
            The current MyVectorStore object, allowing use within a `with` statement block.

        Usage Notes:
            - Intended for use in a `with` statement to ensure automatic
              cleanup of resources.
            - No special initialization occurs during context entry, but enables
              proper context-managed lifecycle.

        Example:
            with MyVectorStore(db_connection, embedder) as vectorstore:
                vectorstore.add_texts([...])
                # Vector store is active within this block.
            # All storage and resources are now cleaned up.
        """
        return self

    def __exit__(
        self,
        exc_type: Union[type, None],
        exc_val: Union[BaseException, None],
        exc_tb: Union[object, None],
    ) -> None:
        """
        Exit the runtime context for the vector store, ensuring all storage
        resources are cleaned up.

        Args:
            exc_type: The exception type, if any exception occurred in the context block.
            exc_val: The exception value, if any exception occurred in the context block.
            exc_tb:  The traceback object, if any exception occurred in the context block.

        Returns:
            None: Indicates that exceptions are never suppressed; they will propagate as normal.

        Implementation Notes:
            - Automatically deletes all vector store data and backing tables via `delete_all()`
            upon exiting the context.
            - This cleanup occurs whether the block exits normally or due to an exception.
            - Does not suppress exceptions; errors in the context block will continue to propagate.
            - Use when the vector store lifecycle is intended to be temporary or scoped.

        Example:
            with MyVectorStore(db_connection, embedder) as vectorstore:
                vectorstore.add_texts([...])
                # Vector store is active within this block.
            # All storage and resources are now cleaned up.
        """
        self.delete_all()
        # No return, so exceptions are never suppressed