database.py

import sqlite3
import numpy as np
import json
from typing import List, Tuple, Optional

class FaceDatabase:
    """Database manager for storing face embeddings and names."""
    
    def __init__(self, db_path: str = "faces.db"):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        """Initialize the database with the faces table."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS faces (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                embedding TEXT NOT NULL,
                list_type TEXT DEFAULT 'none',
                photo_path TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # Add list_type column if it doesn't exist (for existing databases)
        try:
            cursor.execute('ALTER TABLE faces ADD COLUMN list_type TEXT DEFAULT "none"')
        except sqlite3.OperationalError:
            pass  # Column already exists
        
        # Add photo_path column if it doesn't exist
        try:
            cursor.execute('ALTER TABLE faces ADD COLUMN photo_path TEXT')
        except sqlite3.OperationalError:
            pass  # Column already exists
        
        # Create notifications table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS notifications (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                screenshot_path TEXT NOT NULL,
                face_encoding TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                status TEXT DEFAULT 'pending'
            )
        ''')
        
        # Create settings table for email config
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS settings (
                key TEXT PRIMARY KEY,
                value TEXT
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def add_face(self, name: str, embedding: np.ndarray, list_type: str = 'none', photo_path: str = None) -> bool:
        """
        Add a face embedding with a name to the database.
        
        Args:
            name: Name associated with the face
            embedding: 128-dimensional face embedding vector
            list_type: 'whitelist', 'blacklist', or 'none' (default)
            photo_path: Path to the photo file (optional)
            
        Returns:
            True if successful, False otherwise
        """
        try:
            # Convert numpy array to JSON string
            embedding_json = json.dumps(embedding.tolist())
            
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT INTO faces (name, embedding, list_type, photo_path)
                VALUES (?, ?, ?, ?)
            ''', (name, embedding_json, list_type, photo_path))
            
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"Error adding face: {e}")
            return False
    
    def add_face_azure(self, name: str, person_id: str, list_type: str = 'none', photo_path: str = None) -> bool:
        """
        Add a face using Azure person ID (instead of local embedding).
        
        Args:
            name: Name associated with the face
            person_id: Azure Face API person ID
            list_type: 'whitelist', 'blacklist', or 'none' (default)
            photo_path: Path to the photo file (optional)
            
        Returns:
            True if successful, False otherwise
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Store person_id in embedding field (as JSON string for compatibility)
            cursor.execute('''
                INSERT INTO faces (name, embedding, list_type, photo_path)
                VALUES (?, ?, ?, ?)
            ''', (name, json.dumps({'azure_person_id': person_id}), list_type, photo_path))
            
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"Error adding Azure face: {e}")
            return False
    
    def get_all_faces(self) -> List[Tuple[int, str, np.ndarray, str, str]]:
        """
        Retrieve all face embeddings from the database.
        
        Returns:
            List of tuples (id, name, embedding, list_type, photo_path)
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('SELECT id, name, embedding, COALESCE(list_type, "none"), photo_path FROM faces')
        rows = cursor.fetchall()
        
        faces = []
        for row in rows:
            face_id, name, embedding_json, list_type, photo_path = row
            embedding = np.array(json.loads(embedding_json))
            faces.append((face_id, name, embedding, list_type, photo_path))
        
        conn.close()
        return faces
    
    def get_faces_by_list_type(self, list_type: str) -> List[Tuple[int, str, np.ndarray, str]]:
        """
        Get faces by list type (whitelist, blacklist, or none).
        
        Args:
            list_type: 'whitelist', 'blacklist', or 'none'
            
        Returns:
            List of tuples (id, name, embedding, photo_path)
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('SELECT id, name, embedding, photo_path FROM faces WHERE COALESCE(list_type, "none") = ?', (list_type,))
        rows = cursor.fetchall()
        
        faces = []
        for row in rows:
            face_id, name, embedding_json, photo_path = row
            embedding = np.array(json.loads(embedding_json))
            faces.append((face_id, name, embedding, photo_path))
        
        conn.close()
        return faces
    
    def set_list_type(self, name: str, list_type: str) -> bool:
        """
        Set the list type for a face (whitelist, blacklist, or none).
        
        Args:
            name: Name of the face
            list_type: 'whitelist', 'blacklist', or 'none'
            
        Returns:
            True if successful, False otherwise
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('UPDATE faces SET list_type = ? WHERE name = ?', (list_type, name))
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"Error setting list type: {e}")
            return False
    
    def find_match(self, query_embedding: np.ndarray, threshold: float = 0.6, 
                   use_whitelist: bool = False, use_blacklist: bool = False) -> Optional[Tuple[str, float]]:
        """
        Find the best matching face in the database.
        
        Args:
            query_embedding: Face embedding to match
            threshold: Similarity threshold (lower is more similar, 0.6 is default)
            use_whitelist: If True, only search in whitelist
            use_blacklist: If True, exclude blacklist from search
            
        Returns:
            Tuple of (name, distance) if match found, None otherwise
        """
        if use_whitelist:
            all_faces = [(f[0], f[1], f[2]) for f in self.get_all_faces() if f[3] == 'whitelist']
        else:
            all_faces = self.get_all_faces()
        
        if not all_faces:
            return None
        
        # Filter out blacklist if needed
        if use_blacklist:
            blacklist_names = {f[1] for f in self.get_faces_by_list_type('blacklist')}
            all_faces = [(f[0], f[1], f[2]) for f in all_faces if f[1] not in blacklist_names]
        
        best_match = None
        best_distance = float('inf')
        
        for face_id, name, stored_embedding in all_faces:
            # Calculate Euclidean distance (face_recognition uses this)
            distance = np.linalg.norm(query_embedding - stored_embedding)
            
            if distance < best_distance:
                best_distance = distance
                best_match = (name, distance)
        
        # Check if best match is within threshold
        if best_match and best_match[1] <= threshold:
            return best_match
        
        return None
    
    def delete_face(self, name: str) -> bool:
        """
        Delete all faces with a given name.
        
        Args:
            name: Name of the face to delete
            
        Returns:
            True if successful, False otherwise
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('DELETE FROM faces WHERE name = ?', (name,))
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"Error deleting face: {e}")
            return False
    
    def clear_all_faces(self) -> bool:
        """
        Delete all faces from the database.
        
        Returns:
            True if successful, False otherwise
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('DELETE FROM faces')
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"Error clearing database: {e}")
            return False
    
    def add_notification(self, screenshot_path: str, face_encoding: np.ndarray = None) -> int:
        """
        Add a notification for a detected unknown face.
        
        Args:
            screenshot_path: Path to the screenshot
            face_encoding: Face encoding (optional)
            
        Returns:
            Notification ID
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            encoding_json = None
            if face_encoding is not None:
                encoding_json = json.dumps(face_encoding.tolist())
            
            cursor.execute('''
                INSERT INTO notifications (screenshot_path, face_encoding)
                VALUES (?, ?)
            ''', (screenshot_path, encoding_json))
            
            notification_id = cursor.lastrowid
            conn.commit()
            conn.close()
            return notification_id
        except Exception as e:
            print(f"Error adding notification: {e}")
            return None
    
    def get_notifications(self) -> List[Tuple[int, str, str, str, str]]:
        """
        Get all notifications.
        
        Returns:
            List of tuples (id, screenshot_path, face_encoding, created_at, status)
        """
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('SELECT id, screenshot_path, face_encoding, created_at, status FROM notifications ORDER BY created_at DESC')
        rows = cursor.fetchall()
        
        conn.close()
        return rows
    
    def update_notification_status(self, notification_id: int, status: str) -> bool:
        """Update notification status."""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('UPDATE notifications SET status = ? WHERE id = ?', (status, notification_id))
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"Error updating notification: {e}")
            return False
    
    def clear_all_notifications(self) -> int:
        """Clear all notifications from the database. Returns number of rows deleted."""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute('DELETE FROM notifications')
            rows_deleted = cursor.rowcount
            conn.commit()
            conn.close()
            return rows_deleted
        except Exception as e:
            print(f"Error clearing notifications: {e}")
            return 0
    
    def get_notification_screenshot_paths(self) -> List[str]:
        """Get all notification screenshot paths."""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute('SELECT screenshot_path FROM notifications')
            paths = [row[0] for row in cursor.fetchall()]
            conn.close()
            return paths
        except Exception as e:
            print(f"Error getting notification paths: {e}")
            return []
    
    def delete_notification(self, notification_id: int) -> bool:
        """Delete a single notification by ID."""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute('DELETE FROM notifications WHERE id = ?', (notification_id,))
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"Error deleting notification: {e}")
            return False
    
    def get_notification_by_id(self, notification_id: int) -> Optional[Tuple]:
        """Get a notification by ID."""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            cursor.execute('SELECT id, screenshot_path, face_encoding, created_at, status FROM notifications WHERE id = ?', (notification_id,))
            row = cursor.fetchone()
            conn.close()
            return row
        except Exception as e:
            print(f"Error getting notification: {e}")
            return None
    
    def get_setting(self, key: str) -> Optional[str]:
        """Get a setting value."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('SELECT value FROM settings WHERE key = ?', (key,))
        row = cursor.fetchone()
        
        conn.close()
        return row[0] if row else None
    
    def set_setting(self, key: str, value: str) -> bool:
        """Set a setting value."""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT OR REPLACE INTO settings (key, value)
                VALUES (?, ?)
            ''', (key, value))
            
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"Error setting setting: {e}")
            return False