import psycopg2
import requests
import socket
import os

class SystemAwakeningPulse:
    """
    Verifies Kernel/RLM/AIVA connectivity.
    """

    def __init__(self, pg_host, pg_db, pg_user, pg_password, kernel_host, qdrant_host):
        """
        Initializes the SystemAwakeningPulse object.

        Args:
            pg_host (str): PostgreSQL host.
            pg_db (str): PostgreSQL database name.
            pg_user (str): PostgreSQL username.
            pg_password (str): PostgreSQL password.
            kernel_host (str): Genesis Kernel host.
            qdrant_host (str): Qdrant Vector Database host.
        """
        self.pg_host = pg_host
        self.pg_db = pg_db
        self.pg_user = pg_user
        self.pg_password = pg_password
        self.kernel_host = kernel_host
        self.qdrant_host = qdrant_host

    def check_postgres_connection(self):
        """
        Checks connection to PostgreSQL (AIVA memory).

        Returns:
            bool: True if connection is successful, False otherwise.
        """
        try:
            conn = psycopg2.connect(
                host=self.pg_host,
                database=self.pg_db,
                user=self.pg_user,
                password=self.pg_password,
                connect_timeout=5  # Timeout after 5 seconds
            )
            conn.close()
            return True
        except psycopg2.Error as e:
            print(f"Error connecting to PostgreSQL: {e}")
            return False

    def ping_genesis_kernel(self):
        """
        Pings Genesis Kernel.

        Returns:
            bool: True if ping is successful, False otherwise.
        """
        try:
            # Use socket to check if the Kernel is listening on the port
            host, port = self.kernel_host.split(':')
            port = int(port)
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(5)  # Timeout after 5 seconds
                s.connect((host, port))
                return True
        except (socket.error, ValueError) as e:
            print(f"Error pinging Genesis Kernel: {e}")
            return False

    def verify_qdrant_status(self):
        """
        Verifies Vector Database (Qdrant) status.

        Returns:
            bool: True if Qdrant is reachable, False otherwise.
        """
        try:
            response = requests.get(f"http://{self.qdrant_host}/", timeout=5)
            response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
            return True
        except requests.exceptions.RequestException as e:
            print(f"Error verifying Qdrant status: {e}")
            return False

    def run_pulse(self):
        """
        Runs the system awakening pulse check.

        Returns:
            str: "SYSTEM AWAKE" or specific error message.
        """
        if not self.check_postgres_connection():
            return "ERROR: Could not connect to PostgreSQL (AIVA memory)."
        if not self.ping_genesis_kernel():
            return "ERROR: Could not ping Genesis Kernel."
        if not self.verify_qdrant_status():
            return "ERROR: Could not verify Qdrant status."

        return "SYSTEM AWAKE"

if __name__ == '__main__':
    # Example usage (replace with your actual configuration)
    pg_host = os.environ.get("POSTGRES_HOST", "localhost")  # Default to localhost
    pg_db = os.environ.get("POSTGRES_DB", "aiva_memory")  # Default to aiva_memory
    pg_user = os.environ.get("POSTGRES_USER", "aiva")  # Default to aiva
    pg_password = os.environ.get("POSTGRES_PASSWORD", "aiva_password")  # Default to aiva_password
    kernel_host = os.environ.get("KERNEL_HOST", "localhost:8000")  # Default to localhost:8000
    qdrant_host = os.environ.get("QDRANT_HOST", "localhost:6333")  # Default to localhost:6333

    pulse = SystemAwakeningPulse(pg_host, pg_db, pg_user, pg_password, kernel_host, qdrant_host)
    status = pulse.run_pulse()
    print(status)