# 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

"""GenAI LLM integration utilities for MySQL Connector/Python.

Provides MyLLM wrapper that issues ML_GENERATE calls via SQL.
"""

import json

from typing import Any, List, Optional

try:
    from langchain_core.language_models.llms import LLM
except ImportError:
    from langchain.llms.base import LLM
from pydantic import PrivateAttr

from mysql.ai.utils import atomic_transaction, execute_sql, format_value_sql
from mysql.connector.abstracts import MySQLConnectionAbstract


class MyLLM(LLM):
    """
    Custom Large Language Model (LLM) interface for MySQL HeatWave.

    This class wraps the generation functionality provided by HeatWave LLMs,
    exposing an interface compatible with common LLM APIs for text generation.
    It provides full support for generative queries and limited support for
    agentic queries.

    Attributes:
        _db_connection (MySQLConnectionAbstract):
            Underlying MySQL connector database connection.
    """

    _db_connection: MySQLConnectionAbstract = PrivateAttr()

    class Config:
        """
        Pydantic config for the model.

        By default, LangChain (through Pydantic BaseModel) does not allow
        setting or storing undeclared attributes such as _db_connection.
        Setting extra = "allow" makes it possible to store extra attributes
        on the class instance, which is required for MyLLM.
        """

        extra = "allow"

    def __init__(self, db_connection: MySQLConnectionAbstract):
        """
        Initialize the MyLLM instance with an active MySQL database connection.

        Args:
            db_connection: A MySQL connection object used to run LLM queries.

        Notes:
            The db_connection is stored as a private attribute via object.__setattr__,
            which is compatible with Pydantic models.
        """
        super().__init__()

        self._db_connection = db_connection

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> str:
        """
        Generate a text completion from the LLM for a given input prompt.

        References:
            https://dev.mysql.com/doc/heatwave/en/mys-hwgenai-ml-generate.html
                A full list of supported options (specified by kwargs) can be found under "options"

        Args:
            prompt: The input prompt string for the language model.
            stop: Optional list of stop strings to support agentic and chain-of-thought
                  reasoning workflows.
            **kwargs: Additional keyword arguments providing generation options to
                      the LLM (these are serialized to JSON and passed to the HeatWave syscall).

        Returns:
            The generated model output as a string.
            (The actual completion does NOT include the input prompt.)

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

        Implementation Notes:
            - Serializes kwargs into a SQL-compatible JSON string.
            - Calls the LLM stored procedure using a database cursor context.
            - Uses `sys.ML_GENERATE` on the server to produce the model output.
            - Expects the server response to be a JSON object with a 'text' key.
        """
        options = kwargs.copy()
        if stop is not None:
            options["stop_sequences"] = stop

        options_placeholder, options_params = format_value_sql(options)
        with atomic_transaction(self._db_connection) as cursor:
            # The prompt is passed as a parameterized argument (avoids SQL injection).
            generate_query = f"""SELECT sys.ML_GENERATE("%s", {options_placeholder});"""
            execute_sql(cursor, generate_query, params=(prompt, *options_params))
            # Expect a JSON-encoded result from MySQL; parse to extract the output.
            llm_response = json.loads(cursor.fetchone()[0])["text"]

        return llm_response

    @property
    def _identifying_params(self) -> dict:
        """
        Return a dictionary of params that uniquely identify this LLM instance.

        Returns:
            dict: Dictionary of identifier parameters (should include
                model_name for tracing/caching).
        """
        return {
            "model_name": "mysql_heatwave_llm",
        }

    @property
    def _llm_type(self) -> str:
        """
        Get the type name of this LLM implementation.

        Returns:
            A string identifying the LLM provider (used for logging or metrics).
        """
        return "mysql_heatwave_llm"