import json
import boto3
import psycopg
import requests
from botocore.exceptions import ClientError
from datetime import datetime, timezone

# Secrets Manager client
secrets_client = boto3.client("secretsmanager")
s3 = boto3.client("s3")
# Target bucket and file details
bucket_name = "suno-database-monior-log"
now = datetime.now(timezone.utc)
webhook_url = (
    "https://hooks.slack.com/services/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"
)
# Format as "YYYY-MM-DD HH"
date = now.strftime("%Y-%m-%d")
file_name = (
    str(date) + "/" + str(datetime.now(timezone.utc).hour) + "/" + now.strftime("%M:%S")
)


def get_db_credentials(secret_name):
    """
    Retrieve database credentials from AWS Secrets Manager.
    """
    try:
        response = secrets_client.get_secret_value(SecretId=secret_name)
        secret = json.loads(response["SecretString"])
        return secret
    except ClientError as e:
        print(f"Error retrieving secret: {e}")
        raise e


def get_long_running_jobs(connection, database_name):
    try:
    # alert all the pid which are running for more than 8 minutes, except autovacuum task
        cursor = connection.cursor()
        cursor.execute("""
            SELECT 
                pid,
                usename AS user_name,
                application_name,
                client_addr,
                state,
                state_change,
                wait_event_type,
                wait_event,
                now() - state_change AS execution_time,
                query
            FROM pg_stat_activity
            WHERE now() - query_start > interval '8 minutes' and query not ILIKE 'autovacuum:%';
        """)
        results = cursor.fetchall()
        m = []
        for row in results:
            pid = row[0]
            state = row[4]
            cur = {
                "pid": str(row[0]),
                "user_name": str(row[1]),
                "application_name": str(row[2]),
                "client_addr": str(row[3]),
                "state": str(row[4]),
                "state_change": str(row[5]),
                "wait_event_type": str(row[6]),
                "wait_event": str(row[7]),
                "execution_time": str(row[8]),
                "query": str(row[9])
            }
            m.append(cur)
        
        if len(m) > 0:
            payload = {
                "text": str(len(m))
                + database_name + " database jobs are running more than 8 minutes, please check " + (file_name + "-long-running") + " in s3 suno-database-monior-log. (link: https://us-east-2.console.aws.amazon.com/s3/buckets/suno-database-monior-log?region=us-east-2&bucketType=general&tab=objects)"
            }
            response = requests.post(webhook_url, json=payload)
            result = json.dumps(m)
            s3.put_object(
                Bucket=bucket_name,
                Key=file_name + "-long-running",
                Body=result,
                ContentType="application/json",
            )
        cursor.close()
        # Close the connection
        connection.close()
    except Exception as e:
        print(f"Error connecting to database: {e}")
        payload = {
            "text": "database moniotr jobs failed, reason: " + str(e)
        }
        # Send the POST request
        response = requests.post(webhook_url, json=payload)
    finally:
        if cursor:
            cursor.close()

        if connection:
            connection.close()


def handler(event, context):
    # 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)

    # Database connection parameters
    db_host = "suno-main-postgres-prod-cluster.cluster-cnfvffydbwvc.us-east-2.rds.amazonaws.com"
    db_user = credentials["username"]
    db_password = credentials["password"]
    db_name = "suno_main"
    db_port = credentials.get("port", 5432)  # Default port for PostgreSQL

    read_db_host = "suno-main-postgres-prod-cluster.cluster-ro-cnfvffydbwvc.us-east-2.rds.amazonaws.com"
    read_db_user = credentials["username"]
    read_db_password = credentials["password"]
    read_db_name = "suno_main"
    read_db_port = credentials.get("port", 5432)  # Default port for PostgreSQL

    # Connect to the RDS instance
    try:
        connection = psycopg.connect(
            host=db_host,
            user=db_user,
            password=db_password,
            dbname=db_name,
            port=db_port,
        )
        get_long_running_jobs(connection, "write")
    except Exception as e:
        print(f"Error connecting to database: {e}")
        payload = {
            "text": "database moniotr jobs failed, reason: " + str(e)
        }
        # Send the POST request
        requests.post(webhook_url, json=payload)
    finally:
        if connection:
            connection.close()
        
    
    try:
        read_connection = psycopg.connect(
            host=read_db_host,
            user=read_db_user,
            password=read_db_password,
            dbname=read_db_name,
            port=read_db_port,
        )
        get_long_running_jobs(read_connection, "read")
    except Exception as e:
        print(f"Error connecting to database: {e}")
        payload = {
            "text": "database moniotr jobs failed, reason: " + str(e)
        }
        # Send the POST request
        requests.post(webhook_url, json=payload)
    finally:
        if read_connection:
            read_connection.close()