db.py

import pymysql
import json
import os
import time
from dotenv import load_dotenv
from werkzeug.security import generate_password_hash, check_password_hash

load_dotenv(override=True)

DB_HOST = os.environ.get("DB_HOST", "localhost")
DB_USER = os.environ.get("DB_USER", "root")
DB_PASS = os.environ.get("DB_PASS", "")
DB_NAME = os.environ.get("DB_NAME", "vorniity")

def get_db_connection():
    return pymysql.connect(
        host=DB_HOST,
        user=DB_USER,
        password=DB_PASS,
        database=DB_NAME,
        cursorclass=pymysql.cursors.DictCursor
    )

def init_db():
    """Initialize the MySQL database for caching results."""
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS results_cache_v2 (
                    usn VARCHAR(20),
                    url VARCHAR(255),
                    data JSON,
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                    PRIMARY KEY (usn, url)
                )
            ''')
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS subject_credits (
                    subject_code VARCHAR(50) PRIMARY KEY,
                    credits INT NOT NULL
                )
            ''')
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS classes (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    name VARCHAR(255) NOT NULL,
                    start_usn VARCHAR(20) NULL,
                    end_usn VARCHAR(20) NULL,
                    lateral_start VARCHAR(20) NULL,
                    lateral_end VARCHAR(20) NULL,
                    usn_list MEDIUMTEXT NULL
                )
            ''')
            try:
                cursor.execute("ALTER TABLE classes ADD COLUMN lateral_start VARCHAR(20) NULL AFTER end_usn")
                cursor.execute("ALTER TABLE classes ADD COLUMN lateral_end VARCHAR(20) NULL AFTER lateral_start")
            except Exception:
                pass
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS scrape_history (
                    id VARCHAR(255) PRIMARY KEY,
                    start_usn VARCHAR(20),
                    end_usn VARCHAR(20),
                    total_usns INT,
                    completed INT,
                    time_taken FLOAT,
                    status VARCHAR(50),
                    timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS users (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    name VARCHAR(255) NOT NULL,
                    college VARCHAR(255) NOT NULL,
                    email VARCHAR(255) UNIQUE NOT NULL,
                    phone VARCHAR(50),
                    password_hash VARCHAR(255) NOT NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
        conn.commit()
        conn.close()
    except Exception as e:
        print(f"[DB ERROR] Failed to initialize DB: {e}")

def seed_credits_if_empty(hardcoded_map):
    """Ensure hardcoded subject credits exist in the DB without overwriting user changes."""
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            for code, credit in hardcoded_map.items():
                cursor.execute(
                    """
                    INSERT INTO subject_credits (subject_code, credits) 
                    VALUES (%s, %s)
                    ON DUPLICATE KEY UPDATE credits = IF(credits = 0, VALUES(credits), credits)
                    """,
                    (code, credit)
                )
            conn.commit()
        conn.close()
    except Exception as e:
        print(f"[DB ERROR] Failed to seed credits: {e}")

_CREDITS_CACHE = {}
_CREDITS_LAST_FETCH = 0
_CACHE_TTL = 300 # 5 minutes cache to prevent Hostinger max_connections_per_hour (500) errors

def get_all_credits(force_refresh=False):
    """Retrieve all subject credits as a dictionary with caching."""
    global _CREDITS_CACHE, _CREDITS_LAST_FETCH
    
    current_time = time.time()
    if not force_refresh and _CREDITS_CACHE and (current_time - _CREDITS_LAST_FETCH < _CACHE_TTL):
        return _CREDITS_CACHE

    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute("SELECT subject_code, credits FROM subject_credits")
            rows = cursor.fetchall()
        conn.close()
        
        _CREDITS_CACHE = {row['subject_code']: row['credits'] for row in rows}
        _CREDITS_LAST_FETCH = current_time
        return _CREDITS_CACHE
    except Exception as e:
        print(f"[DB ERROR] Failed to fetch credits: {e}")
        return _CREDITS_CACHE if _CREDITS_CACHE else {}

def save_credit(subject_code, credits):
    """Save or update a subject credit."""
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute(
                "REPLACE INTO subject_credits (subject_code, credits) VALUES (%s, %s)",
                (subject_code.upper(), int(credits))
            )
        conn.commit()
        conn.close()
        
        # Instantly update cache
        global _CREDITS_CACHE
        if _CREDITS_CACHE is not None:
            _CREDITS_CACHE[subject_code.upper()] = int(credits)
            
        return True
    except Exception as e:
        print(f"[DB ERROR] Failed to save credit: {e}")
        return False

def save_bulk_credits(credits_list):
    """Save multiple subject credits efficiently."""
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            # credits_list is a list of dicts: [{'subject_code': 'ABC', 'credits': 3}, ...]
            values = [(item['subject_code'].upper(), int(item['credits'])) for item in credits_list]
            cursor.executemany(
                "REPLACE INTO subject_credits (subject_code, credits) VALUES (%s, %s)",
                values
            )
        conn.commit()
        conn.close()
        
        # Invalidate cache so next fetch gets everything
        global _CREDITS_CACHE
        _CREDITS_CACHE = None
        
        return True
    except Exception as e:
        print(f"[DB ERROR] Failed to save bulk credits: {e}")
        return False

def delete_credit(subject_code):
    """Delete a subject credit."""
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute("DELETE FROM subject_credits WHERE subject_code = %s", (subject_code.upper(),))
        conn.commit()
        conn.close()
        return True
    except Exception as e:
        print(f"[DB ERROR] Failed to delete credit: {e}")
        return False

# --- Classes Logic ---

def create_class(name, start_usn=None, end_usn=None, lateral_start=None, lateral_end=None, usn_list=None):
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute(
                "INSERT INTO classes (name, start_usn, end_usn, lateral_start, lateral_end, usn_list) VALUES (%s, %s, %s, %s, %s, %s)",
                (name, start_usn.upper() if start_usn else None, end_usn.upper() if end_usn else None, lateral_start.upper() if lateral_start else None, lateral_end.upper() if lateral_end else None, usn_list)
            )
        conn.commit()
        conn.close()
        return True
    except Exception as e:
        print(f"[DB ERROR] Failed to create class: {e}")
        return False

def get_all_classes():
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute("SELECT id, name, start_usn, end_usn, lateral_start, lateral_end, usn_list FROM classes ORDER BY id DESC")
            rows = cursor.fetchall()
        conn.close()
        return rows
    except Exception as e:
        print(f"[DB ERROR] Failed to fetch classes: {e}")
        return []

def delete_class(class_id):
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute("DELETE FROM classes WHERE id = %s", (class_id,))
        conn.commit()
        conn.close()
        return True
    except Exception as e:
        print(f"[DB ERROR] Failed to delete class: {e}")
        return False

def update_class(class_id, name, start_usn=None, end_usn=None, lateral_start=None, lateral_end=None, usn_list=None):
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute(
                "UPDATE classes SET name = %s, start_usn = %s, end_usn = %s, lateral_start = %s, lateral_end = %s, usn_list = %s WHERE id = %s",
                (name, start_usn.upper() if start_usn else None, end_usn.upper() if end_usn else None, lateral_start.upper() if lateral_start else None, lateral_end.upper() if lateral_end else None, usn_list, class_id)
            )
        conn.commit()
        conn.close()
        return True
    except Exception as e:
        print(f"[DB ERROR] Failed to update class: {e}")
        return False

# --- Scrape History Logic ---

def save_scrape_history(job_id, start_usn, end_usn, total_usns, completed, time_taken, status):
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute(
                "REPLACE INTO scrape_history (id, start_usn, end_usn, total_usns, completed, time_taken, status) VALUES (%s, %s, %s, %s, %s, %s, %s)",
                (job_id, start_usn, end_usn, total_usns, completed, time_taken, status)
            )
        conn.commit()
        conn.close()
        return True
    except Exception as e:
        print(f"[DB ERROR] Failed to save scrape history: {e}")
        return False

def get_scrape_history():
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute("SELECT id, start_usn, end_usn, total_usns, completed, time_taken, status, timestamp FROM scrape_history ORDER BY timestamp DESC")
            rows = cursor.fetchall()
        conn.close()
        
        # Format timestamp to ISO string
        for row in rows:
            if row['timestamp']:
                row['timestamp'] = row['timestamp'].isoformat()
        return rows
    except Exception as e:
        print(f"[DB ERROR] Failed to fetch scrape history: {e}")
        return []

def delete_student(usn):
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute("DELETE FROM results_cache_v2 WHERE usn = %s", (usn.upper(),))
        conn.commit()
        conn.close()
        return True
    except Exception as e:
        print(f"[DB ERROR] Failed to delete student: {e}")
        return False

def get_cached_result(usn, url):
    """Retrieve a cached result for a given USN and URL if it exists."""
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute("SELECT data FROM results_cache_v2 WHERE usn = %s AND url = %s", (usn.upper(), url))
            row = cursor.fetchone()
        conn.close()
        
        if row:
            if isinstance(row['data'], str):
                return json.loads(row['data'])
            return row['data']
    except Exception as e:
        print(f"[CACHE ERROR] Failed to read cache for {usn}: {e}")
    return None

def save_cached_result(usn, url, result_dict):
    """Save a successfully parsed result to the cache."""
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            # REPLACE INTO handles updating the record if it already exists
            cursor.execute(
                "REPLACE INTO results_cache_v2 (usn, url, data) VALUES (%s, %s, %s)",
                (usn.upper(), url, json.dumps(result_dict))
            )
        conn.commit()
        conn.close()
    except Exception as e:
        print(f"[CACHE ERROR] Failed to save cache for {usn}: {e}")

def clear_database():
    """Wipe all student scraped data, results, and history."""
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute("SET FOREIGN_KEY_CHECKS = 0;")
            cursor.execute("TRUNCATE TABLE results_cache_v2;")
            cursor.execute("TRUNCATE TABLE scrape_history;")
            cursor.execute("SET FOREIGN_KEY_CHECKS = 1;")
        conn.commit()
        conn.close()
        return True, "Database cleared successfully"
    except Exception as e:
        print(f"[DB ERROR] Failed to clear database: {e}")
        return False, str(e)

# --- Auth Logic ---
def create_user(name, college, email, phone, password):
    try:
        conn = get_db_connection()
        password_hash = generate_password_hash(password)
        with conn.cursor() as cursor:
            cursor.execute(
                "INSERT INTO users (name, college, email, phone, password_hash) VALUES (%s, %s, %s, %s, %s)",
                (name, college, email, phone, password_hash)
            )
        conn.commit()
        conn.close()
        return True, "User created successfully"
    except pymysql.err.IntegrityError:
        return False, "Email already exists"
    except Exception as e:
        print(f"[DB ERROR] Failed to create user: {e}")
        return False, str(e)

def get_user_by_email(email):
    try:
        conn = get_db_connection()
        with conn.cursor() as cursor:
            cursor.execute("SELECT * FROM users WHERE email = %s", (email,))
            user = cursor.fetchone()
        conn.close()
        return user
    except Exception as e:
        print(f"[DB ERROR] Failed to get user: {e}")
        return None