import psycopg2
from prometheus_client import start_http_server, Gauge
import time

# Configuration for database connection
DB_PARAMS = {
    "host": "admin",
    "database": "admin",
    "user": "admin",
    "password": "admin"
}

# Prometheus metrics
QUERY_CALLS = Gauge('postgresql_query_calls', 'Number of calls for each query', ['query'])
QUERY_TOTAL_TIME = Gauge('postgresql_query_total_time_ms', 'Total execution time for each query in ms', ['query'])

def fetch_metrics():
    try:
        # Log connection attempt
        print("Connecting to PostgreSQL database...")

        conn = psycopg2.connect(**DB_PARAMS)
        cur = conn.cursor()

        # Execute query to get data
        cur.execute("""
            SELECT query, calls, total_exec_time
            FROM pg_stat_statements
            ORDER BY total_exec_time DESC;
        """)

        # Iterate through results and set Prometheus metrics
        for row in cur:
            query = row[0].replace("\\", "\\\\").replace('"', '\\"')  # Escape special characters
            calls = row[1]
            total_time = row[2] * 1000  # Convert seconds to milliseconds

            QUERY_CALLS.labels(query=query).set(calls)
            QUERY_TOTAL_TIME.labels(query=query).set(total_time)

            # Log the metrics being set
            print(f"Metrics set for query: {query} | Calls: {calls} | Total execution time: {total_time} ms")

        cur.close()
        conn.close()
    except psycopg2.Error as e:
        print(f"Error fetching data: {e}")
    except Exception as e:
        print(f"Unexpected error: {e}")

if __name__ == '__main__':
    # Start Prometheus HTTP server on port 8000
    start_http_server(8000)
    print("Exporter running on http://localhost:8000/metrics")

    # Main loop to fetch metrics at regular intervals
    while True:
        fetch_metrics()
        time.sleep(60)  # Scrape every 60 seconds