from prometheus_client import start_http_server, Gauge, REGISTRY
import psycopg2
import time
import os

# Configuration from environment variables
DB_HOST = os.environ.get("POSTGRES_HOST", "admin")
DB_NAME = os.environ.get("POSTGRES_DB", "admin")
DB_USER = os.environ.get("POSTGRES_USER", "admin")
DB_PASSWORD = os.environ.get("POSTGRES_PASSWORD", "admin")

conn_params = {
    "host": DB_HOST,
    "database": DB_NAME,
    "user": DB_USER,
    "password": DB_PASSWORD
}

QUERY_CALLS = Gauge('postgresql_query_calls', 'Number of PostgreSQL query calls', ['query'])
QUERY_TOTAL_TIME = Gauge('postgresql_query_total_time_ms', 'Total time of PostgreSQL queries (ms)', ['query'])


def generate_metrics():
    try:
        conn = psycopg2.connect(**conn_params)
        cur = conn.cursor()
        cur.execute("""
            SELECT query, calls, total_time
            FROM pg_stat_statements
            ORDER BY total_time DESC;
        """)

        for row in cur:
            query = row[0].replace("\\", "\\\\").replace('"', '\\"')
            QUERY_CALLS.labels(query=query).set(row[1])
            QUERY_TOTAL_TIME.labels(query=query).set(row[2] * 1000)

        cur.close()
        conn.close()
    except psycopg2.Error as e:
        print(f"Error connecting to the database: {e}", flush=True)


if __name__ == '__main__':
    start_http_server(8000)
    while True:
        generate_metrics()
        time.sleep(60)