import os
import sqlite3
import hyperon

class DBMeTTa:
    def __init__(self, path, initial=None):
        self.metta = hyperon.MeTTa()
        self.path = path
        if not os.path.exists(self.path):
            if initial:
                self.metta.run(initial)
            self.save()
        else:
            self.load()
    
    def save(self):
        """Save space atoms to an SQLite database by stringifying them."""
        # Connect to SQLite DB
        conn = sqlite3.connect(self.path)
        cursor = conn.cursor()
    
        cursor.execute('BEGIN;')
        
        # Create table to store atoms if it doesn't exist
        cursor.execute('''CREATE TABLE IF NOT EXISTS atoms (
                            atom_str TEXT PRIMARY KEY)''')
    
        cursor.execute('DELETE FROM atoms;')
    
        # Iterate over atoms in the space
        for atom in self.metta.space().get_atoms():
            atom_str = repr(atom)  # Use the string representation of the atom
            cursor.execute('INSERT INTO atoms (atom_str) VALUES (?)', [atom_str,])
        
        # Commit and close
        conn.commit()
        conn.close()

    def load(self):
        """Load space atoms from an SQLite database and add them back to the space."""
        # Connect to SQLite DB
        conn = sqlite3.connect(self.path)
        cursor = conn.cursor()
        
        # Fetch all atom strings
        cursor.execute('SELECT atom_str FROM atoms')
        rows = cursor.fetchall()
        
        # Recreate atoms from the stored strings
        space = self.metta.space()
        existing_atom_strs = set([hash(str(a)) for a in space.get_atoms()])
        for atom_str, in rows:
            if hash(atom_str) in existing_atom_strs:
                continue
    
            # Convert the atom string back to an Atom object
            atom = self.metta.parse_single(atom_str)
            space.add_atom(atom)
        
        # Commit and close
        conn.commit()
        conn.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.save()

if __name__ == "__main__":
    # Initialize a space and a MeTTa instance
    db = DBMeTTa('space_data.db', '''
            (isa red color)
            (isa green color)
            (isa blue color)
            ;(isa comment color)
            !(match &self (isa $color color) $color)
    
            (= (f) (+ 2 3))
            !(f)
    ''')
    atoms = list(db.metta.space().get_atoms())

    # Save the space to the database
    db.save()

    # Load the space from the database
    db.load()

    assert atoms == db.metta.space().get_atoms()

    # Load the space from the database
    db = DBMeTTa(db.path)

    assert [str(atom) for atom in atoms] == [str(atom) for atom in db.metta.space().get_atoms()]
    
    print("Space saved and loaded successfully.")