# 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

"""Embeddings integration utilities for MySQL Connector/Python.

Provides MyEmbeddings class to generate embeddings via MySQL HeatWave
using ML_EMBED_TABLE and ML_EMBED_ROW.
"""

from typing import Dict, List, Optional

import pandas as pd

from langchain_core.embeddings import Embeddings
from pydantic import PrivateAttr

from mysql.ai.utils import (
    atomic_transaction,
    execute_sql,
    format_value_sql,
    source_schema,
    sql_table_from_df,
    sql_table_to_df,
    temporary_sql_tables,
)
from mysql.connector.abstracts import MySQLConnectionAbstract


class MyEmbeddings(Embeddings):
    """
    Embedding generator class that uses a MySQL database to compute embeddings for input text.

    This class batches input text into temporary SQL tables, invokes MySQL's ML_EMBED_TABLE
    to generate embeddings, and retrieves the results as lists of floats.

    Attributes:
        _db_connection (MySQLConnectionAbstract): MySQL connection used for all database operations.
        schema_name (str): Name of the database schema to use.
        options_placeholder (str): SQL-ready placeholder string for ML_EMBED_TABLE options.
        options_params (dict): Dictionary of concrete option values to be passed as SQL parameters.
    """

    _db_connection: MySQLConnectionAbstract = PrivateAttr()

    def __init__(
        self, db_connection: MySQLConnectionAbstract, options: Optional[Dict] = None
    ):
        """
        Initialize MyEmbeddings with a database connection and optional embedding parameters.

        References:
            https://dev.mysql.com/doc/heatwave/en/mys-hwgenai-ml-embed-row.html
                A full list of supported options can be found under "options"

        NOTE: The supported "options" are the intersection of the options provided in
            https://dev.mysql.com/doc/heatwave/en/mys-hwgenai-ml-embed-row.html
            https://dev.mysql.com/doc/heatwave/en/mys-hwgenai-ml-embed-table.html

        Args:
            db_connection: Active MySQL connector database connection.
            options: Optional dictionary of options for embedding operations.

        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._db_connection = db_connection
        self.schema_name = source_schema(db_connection)
        options = options or {}
        self.options_placeholder, self.options_params = format_value_sql(options)

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """
        Generate embeddings for a list of input texts using the MySQL ML embedding procedure.

        References:
            https://dev.mysql.com/doc/heatwave/en/mys-hwgenai-ml-embed-table.html

        Args:
            texts: List of input strings to embed.

        Returns:
            List of lists of floats, with each inner list containing the embedding for a text.

        Raises:
            DatabaseError:
                If provided options are invalid or unsupported.
                If a database connection issue occurs.
                If an operational error occurs during execution.
            ValueError:
                If one or more text entries were unable to be embedded.

        Implementation notes:
            - Creates a temporary table to pass input text to the MySQL embedding service.
            - Adds a primary key to ensure results preserve input order.
            - Calls ML_EMBED_TABLE and fetches the resulting embeddings.
            - Deletes the temporary table after use to avoid polluting the database.
            - Embedding vectors are extracted from the "embeddings" column of the result table.
        """
        if not texts:
            return []

        df = pd.DataFrame({"id": range(len(texts)), "text": texts})

        with (
            atomic_transaction(self._db_connection) as cursor,
            temporary_sql_tables(self._db_connection) as temporary_tables,
        ):
            qualified_table_name, table_name = sql_table_from_df(
                cursor, self.schema_name, df
            )
            temporary_tables.append((self.schema_name, table_name))

            # ML_EMBED_TABLE expects input/output columns and options as parameters
            embed_query = (
                "CALL sys.ML_EMBED_TABLE("
                f"'{qualified_table_name}.text', "
                f"'{qualified_table_name}.embeddings', "
                f"{self.options_placeholder}"
                ")"
            )
            execute_sql(cursor, embed_query, params=self.options_params)

            # Read back all columns, including "embeddings"
            df_embeddings = sql_table_to_df(cursor, self.schema_name, table_name)

            if df_embeddings["embeddings"].isnull().any() or any(
                e is None for e in df_embeddings["embeddings"]
            ):
                raise ValueError(
                    "Failure to generate embeddings for one or more text entry."
                )

            # Convert fetched embeddings to lists of floats
            embeddings = df_embeddings["embeddings"].tolist()
            embeddings = [list(e) for e in embeddings]

            return embeddings

    def embed_query(self, text: str) -> List[float]:
        """
        Generate an embedding for a single text string.

        References:
            https://dev.mysql.com/doc/heatwave/en/mys-hwgenai-ml-embed-row.html

        Args:
            text: The input string to embed.

        Returns:
            List of floats representing the embedding vector.

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

        Example:
            >>> MyEmbeddings(db_conn).embed_query("Hello world")
            [0.1, 0.2, ...]
        """
        with atomic_transaction(self._db_connection) as cursor:
            execute_sql(
                cursor,
                f'SELECT sys.ML_EMBED_ROW("%s", {self.options_placeholder})',
                params=(text, *self.options_params),
            )
            return list(cursor.fetchone()[0])