import os

from openai import AzureOpenAI

# Set up credentials
# NOTE: When running locally, these have to be set in the environment
client = AzureOpenAI(
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    api_key=os.getenv("AZURE_OPENAI_KEY"),
    api_version="2024-02-01",
)

deployment_name = "sqlai"


def send_message(message: str) -> str:
    """Send a message to the openai chat completion API and return the response.

    Parameters
    ----------
    message : str
        The user's message to be sent to the chat completion API.

    Returns
    -------
    str
        The content of the assistant's response message.
    """
    system_message = """
    Du bist ein hilfsbereiter, fröhlicher Datenbankassistent.
    Verwende beim Erstellen Ihrer Antworten das folgende Datenbankschema:

    MEIN_DATENBANKSCHEMA

    Füge Spaltenüberschriften in die Abfrageergebnisse ein.

    Gib deine Antwort immer im folgenden JSON-Format an:

    JSON FORMAT

    Gib NUR JSON aus.
    Ersetze in der vorangehenden JSON-Antwort "your-query" durch die Microsoft SQL Server Query,
    um die angeforderten Daten abzurufen.
    Ersetze in der vorangehenden JSON-Antwort "your-summary" durch eine Zusammenfassung der Abfrage.
    Gib immer alle Spalten der Tabelle an.
    Wenn die resultierende Abfrage nicht ausführbar ist, ersetze "your-query“ durch NA, aber ersetze
    trotzdem "your-query" durch eine Zusammenfassung der Abfrage.
    Verwende KEINE MySQL-Syntax.
    Begrenze die SQL-Abfrage immer auf 100 Zeilen.
    """
    system_message = "Du bist ein hilfreicher Assistent."
    response = client.chat.completions.create(
        model=deployment_name,
        messages=[
            {"role": "system", "content": system_message},
            {"role": "user", "content": message},
        ],
    )
    return response.choices[0].message.content