import os
from typing import Union

import pandas as pd
import pyodbc


def test_db_connection() -> bool:
    """Test the connection to Azure SQL Database.

    This function attempts to establish a connection to an Azure SQL Database
    using the connection string stored in the environment variable
    'AZURE_SQL_CONNECTION_STRING'. It makes up to 5 attempts to connect,
    with a timeout of 480 seconds for each attempt.

    Returns
    -------
    bool
        True if the connection was successful, False otherwise.
    """
    connection_string = os.environ.get("AZURE_SQL_CONNECTION_STRING")
    for i in range(5):
        print(f"Trying to connect to Azure SQL Database... Attempt {i + 1}")
        try:
            pyodbc.connect(connection_string, timeout=480)
            print("Connected to Azure SQL Database successfully!")
            connected = True
            break

        except Exception as e:
            print(f"Error connecting to Azure SQL Database: {e}")
            connected = False
            continue

    return connected


def execute_query(query: str) -> Union[pd.DataFrame, str]:
    """Execute a SQL query on an Azure SQL Database and return the results.

    This function connects to an Azure SQL Database using the connection string
    stored in the environment variable 'AZURE_SQL_CONNECTION_STRING', executes
    the provided SQL query, and returns the results as a pandas DataFrame.

    Parameters
    ----------
    query : str
        The SQL query to execute.

    Returns
    -------
    Union[pd.DataFrame, str]
        A pandas DataFrame containing the query results if successful,
        or a string containing the error message if an exception occurs.
    """
    try:
        connection_string = os.environ.get("AZURE_SQL_CONNECTION_STRING")
        conn = pyodbc.connect(connection_string, timeout=240)
        df = pd.read_sql(query, conn)
        conn.close()
        return df
    except Exception as e:
        return str(e)
    finally:
        if conn in locals():
            conn.close()