backend/database.py

# backend/database.py (ROBUST VERSION)
import sqlite3
from datetime import datetime

DB_NAME = "email_sorter.db"

def init_db():
    """Initializes the database using a context manager for safety."""
    with sqlite3.connect(DB_NAME) as conn:
        cursor = conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS emails (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                timestamp TEXT NOT NULL,
                email_from TEXT,
                subject TEXT,
                body TEXT,
                predicted_folder TEXT,
                confidence REAL,
                user_feedback TEXT 
            )
        ''')
    print("Database initialized successfully.")

def add_email_records(records_df):
    """Adds multiple email records from a pandas DataFrame using a context manager."""
    with sqlite3.connect(DB_NAME) as conn:
        records_df['timestamp'] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        records_df['user_feedback'] = None
        
        df_to_save = records_df.rename(columns={
            'From': 'email_from',
            'Subject': 'subject',
            'Body': 'body'
        })
        
        db_columns = ['timestamp', 'email_from', 'subject', 'body', 'predicted_folder', 'confidence', 'user_feedback']
        df_to_save = df_to_save[[col for col in db_columns if col in df_to_save.columns]]

        df_to_save.to_sql('emails', conn, if_exists='append', index=False)

def get_all_emails():
    """Fetches all email records using a context manager."""
    with sqlite3.connect(DB_NAME) as conn:
        cursor = conn.cursor()
        cursor.execute("SELECT id, timestamp, email_from, subject, predicted_folder, confidence, user_feedback FROM emails ORDER BY timestamp DESC")
        records = cursor.fetchall()
        columns = [description[0] for description in cursor.description]
    return records, columns

def update_feedback(record_id, corrected_label):
    """Updates the user_feedback for a specific record using a context manager."""
    with sqlite3.connect(DB_NAME) as conn:
        cursor = conn.cursor()
        cursor.execute("UPDATE emails SET user_feedback = ? WHERE id = ?", (corrected_label, record_id))
    print(f"Updated record ID {record_id} with feedback: {corrected_label}")