import os
import psycopg
import json
import boto3
from botocore.exceptions import ClientError
import logging

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

def get_secret():
    secret_name = os.environ['SECRET_NAME']
    region_name = os.environ['REGION_NAME']

    # Create a Secrets Manager client
    session = boto3.session.Session()
    client = session.client(
        service_name='secretsmanager',
        region_name=region_name
    )

    try:
        get_secret_value_response = client.get_secret_value(
            SecretId=secret_name
        )
    except ClientError as e:
        raise e

    secret = get_secret_value_response['SecretString']
    return json.loads(secret)

def lambda_handler(event, context):
    conn = None
    try:
        # Fetch database credentials from Secrets Manager
        secret = get_secret()
        db_user = secret['username']
        db_password = secret['password']

        # Fetch other database connection parameters from environment variables
        db_host = os.environ['DB_HOST']
        db_name = os.environ['DB_NAME']
        sns_topic_arn = os.environ['SNS_TOPIC_ARN']
        query = os.environ['DB_QUERY']

        # Connect to the PostgreSQL database using psycopg3
        conn = psycopg.connect(
            host=db_host,
            dbname=db_name,
            user=db_user,
            password=db_password
        )

        # Call the stored procedure
        with conn.cursor() as cur:
            cur.execute(query)
        
        # Commit the transaction
        conn.commit()

        # Send success message to SNS
        sns_client = boto3.client('sns')
        sns_message = "SQL Stored procedure executed successfully."

        sns_client.publish(
            TopicArn=sns_topic_arn,
            Message=sns_message
        )

        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Stored procedure executed successfully'
            }),
            'headers': {
                'Content-Type': 'application/json'
            }
        }
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        return {
            'statusCode': 500,
            'body': str(e),
            'headers': {
                'Content-Type': 'application/json'
            }
        }
    finally:
        if conn is not None:
            conn.close()
            logger.info("Database connection closed")