import sqlite3
from contextlib import contextmanager

DATABASE_URL = "E:/geminicli/face_recognition_service_v2/face_recognition.db"

@contextmanager
def get_db_connection():
    """获取数据库连接，并使用上下文管理器确保连接关闭"""
    conn = sqlite3.connect(DATABASE_URL)
    conn.row_factory = sqlite3.Row
    try:
        yield conn
    finally:
        conn.close()

def initialize_database():
    """初始化数据库，创建所需的表"""
    with get_db_connection() as conn:
        cursor = conn.cursor()
        
        # 创建用户表
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL
        );
        """)
        
        # 创建人脸特征表
        cursor.execute("""
        CREATE TABLE IF NOT EXISTS face_features (
            feature_id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            embedding BLOB NOT NULL,
            FOREIGN KEY (user_id) REFERENCES users (id) ON DELETE CASCADE
        );
        """)
        
        # 创建索引以加速查询
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_id ON face_features (user_id);")
        
        conn.commit()
    print("数据库初始化完成。")

if __name__ == '__main__':
    initialize_database()