import psycopg2
import boto3
import json

secrets_client = boto3.client("secretsmanager")


def get_db_credentials(secret_name):
    """
    Retrieve database credentials from AWS Secrets Manager.
    """
    response = secrets_client.get_secret_value(SecretId=secret_name)
    secret = json.loads(response["SecretString"])
    return secret


# Name of the secret where RDS credentials are stored
secret_name = "app-user-main-db-secret"

# Retrieve credentials from Secrets Manager
credentials = get_db_credentials(secret_name)

db_host = credentials["host"]
db_read_host = credentials["read_host"]
db_user = credentials["username"]
db_password = credentials["password"]
db_name = credentials["database_name"]
db_port = credentials.get("port", 5432)  # Default port for PostgreSQL


def _truncate_application_name(application_name):
    return f"glue_{application_name}"[:64]


def get_postgres_connection(application_name):
    connection = psycopg2.connect(
        host=db_host,
        user=db_user,
        password=db_password,
        dbname=db_name,
        port=db_port,
        application_name=_truncate_application_name(application_name),
    )
    return connection


# Retrieve credentials from Secrets Manager
def get_db_read_connection(application_name):
    connection = psycopg2.connect(
        host=db_read_host,
        user=db_user,
        password=db_password,
        dbname=db_name,
        port=db_port,
        application_name=_truncate_application_name(application_name),
    )
    return connection
