database.py

import sqlite3
import hashlib
import re

DATABASE_FILE = "cbt.db"

def get_db_connection():
    """Establishes connection to the sqlite3 database file."""
    conn = sqlite3.connect(DATABASE_FILE)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    """Creates the database tables if they do not already exist."""
    conn = get_db_connection()
    cursor = conn.cursor()
    
    # Enable Write-Ahead Logging for better speed/performance
    cursor.execute("PRAGMA journal_mode=WAL;")
    
    # Create the users table
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS users (
        id TEXT PRIMARY KEY,
        full_name TEXT NOT NULL,
        role TEXT NOT NULL CHECK(role IN ('student', 'teacher')),
        class_or_subject TEXT NOT NULL,
        pin_hash TEXT NOT NULL,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    );
    """)
    
    conn.commit()
    conn.close()
    print("Database successfully initialized or verified!")

def hash_pin(pin: str) -> str:
    """Hashes a numeric PIN securely using SHA-256."""
    return hashlib.sha256(pin.strip().encode()).hexdigest()

def clean_input(text: str) -> str:
    """Cleans names and subjects of non-alphanumeric characters."""
    return re.sub(r'[^a-zA-Z0-9]', '', text).strip()

def get_next_sequence(prefix: str) -> str:
    """Queries SQLite to find the next sequence number for a prefix and returns pad-3 string (e.g., '002')."""
    conn = get_db_connection()
    cursor = conn.cursor()
    # Find all current user IDs starting with the given prefix
    cursor.execute("SELECT id FROM users WHERE id LIKE ? ORDER BY id DESC", (f"{prefix}%",))
    rows = cursor.fetchall()
    conn.close()
    
    max_num = 0
    for row in rows:
        user_id = row['id']
        parts = user_id.split('-')
        if len(parts) >= 4:
            try:
                num = int(parts[3])
                if num > max_num:
                    max_num = num
            except ValueError:
                continue
                
    next_num = max_num + 1
    return f"{next_num:03d}"

def calculate_student_id(full_name: str, class_name: str) -> str:
    """
    Generates ID in format: STU-SS1-ADM-001 (Role-Class-name-Number)
    """
    # 1. Clean Class name 
    clean_class = class_name.upper().replace(" ", "")
    if not clean_class:
        clean_class = "SS1"
        
    # 2. Extract 3 letter name identifier (e.g., ADM)
    parts = full_name.strip().split()
    first_part = clean_input(parts[0]).upper() if parts else "STU"
    if len(first_part) >= 3:
        name_code = first_part[:3]
    else:
        name_code = (first_part + "XXX")[:3]
        
    # 3. Form base prefix and get sequence
    prefix_base = f"STU-{clean_class}-{name_code}"
    next_num = get_next_sequence(prefix_base)
    
    return f"{prefix_base}-{next_num}"

def calculate_teacher_id(full_name: str, subject_name: str) -> str:
    """
    Generates ID in format: TCH-MTH-BI-001 (Role-Subject-name-Number)
    """
    SUBJECT_CODES = {
        "MATHEMATICS": "MTH",
        "ENGLISH": "ENG",
        "PHYSICS": "PHY",
        "CHEMISTRY": "CHM",
        "BIOLOGY": "BIO",
        "ICT": "ICT"
    }
    
    # 1. Clean Subject Code
    clean_sub = clean_input(subject_name).upper()
    sub_code = SUBJECT_CODES.get(clean_sub, clean_sub[:3])
    if len(sub_code) < 3:
         sub_code = (sub_code + "XXX")[:3]
         
    # 2. Extract 2 letter name initials (e.g. BI)
    parts = [clean_input(p) for p in full_name.strip().split() if clean_input(p)]
    if len(parts) >= 2:
        name_code = (parts[0][0] + parts[1][0]).upper()
    elif len(parts) == 1:
        name_code = parts[0][:2].upper().ljust(2, "X")
    else:
        name_code = "XX"
        
    # 3. Form base prefix and get sequence
    prefix_base = f"TCH-{sub_code}-{name_code}"
    next_num = get_next_sequence(prefix_base)
    
    return f"{prefix_base}-{next_num}"

def register_user(full_name: str, role: str, class_or_subject: str, pin: str, custom_id=None) -> str:
    """
    Registers a new student or teacher. Hash PIN and saves securely.
    Returns the assigned User ID.
    """
    # Auto-generate custom ID if none is supplied or if it requires the format match
    if role.strip().lower() == 'student':
        user_id = calculate_student_id(full_name, class_or_subject)
    else:
        user_id = calculate_teacher_id(full_name, class_or_subject)
        
    # If a specific valid override input ID is offered by user and matches structure, 
    # and doesn't exist, we can use it. But auto-generating guarantees correctness!
    if custom_id and len(custom_id.split('-')) == 4:
        user_id = custom_id.strip().upper()

    conn = get_db_connection()
    cursor = conn.cursor()
    
    # Verify ID is unique in SQLite
    cursor.execute("SELECT id FROM users WHERE id = ?", (user_id,))
    if cursor.fetchone() is not None:
        conn.close()
        raise ValueError(f"User ID {user_id} already exists! Try again.")
        
    # Save user row
    cursor.execute(
        "INSERT INTO users (id, full_name, role, class_or_subject, pin_hash) VALUES (?, ?, ?, ?, ?)",
        (user_id, full_name.strip(), role.strip().lower(), class_or_subject.strip(), hash_pin(pin))
    )
    conn.commit()
    conn.close()
    return user_id

def authenticate_user(user_id: str, pin: str):
    """
    Verifies user authentication against secure PIN hashes.
    Returns the user datatable dictionary if successful, or None.
    """
    clean_id = user_id.strip().upper()
    conn = get_db_connection()
    cursor = conn.cursor()
    
    cursor.execute("SELECT * FROM users WHERE id = ?", (clean_id,))
    user = cursor.fetchone()
    conn.close()
    
    if user:
        if user['pin_hash'] == hash_pin(pin):
            return dict(user)
            
    return None