import sqlite3  
from datetime import datetime  
from cryptography.fernet import Fernet  
import os  

# Generiraj ili učitaj enkripcijski ključ  
def load_key():  
    key_path = "secrets/encryption_key.bin"  
    if not os.path.exists(key_path):  
        os.makedirs("secrets", exist_ok=True)  
        key = Fernet.generate_key()  
        with open(key_path, "wb") as key_file:  
            key_file.write(key)  
    return open(key_path, "rb").read()  

key = load_key()  
cipher = Fernet(key)  

def encrypt_data(data):  
    return cipher.encrypt(data.encode()).decode()  

def decrypt_data(encrypted_data):  
    return cipher.decrypt(encrypted_data.encode()).decode()  

# Modificirane DB funkcije  
def save_to_db(user_input, bot_response, mode):  
    conn = sqlite3.connect('memory.db')  
    cursor = conn.cursor()  
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")  
    cursor.execute(  
        "INSERT INTO chat_history VALUES (?, ?, ?, ?)",  
        (timestamp, encrypt_data(user_input), encrypt_data(bot_response), mode)  
    conn.commit()  
    conn.close()