db_schema.py

"""
Database connection and data operations for Postgres.

Handles connections, upserts for food entries and exercises.
"""

import psycopg2
from psycopg2.extras import RealDictCursor

try:
    import config
except ImportError:
    print("ERROR: config.py not found!")
    print("Please copy config.example.py to config.py and fill in your settings.")
    raise


def get_connection():
    """
    Create and return a database connection.
    
    Returns:
        psycopg2.connection: Database connection
    """
    return psycopg2.connect(**config.POSTGRES_CONFIG)


def delete_records_for_date(target_date, table_name):
    """
    Delete all records for a specific date from a table.
    
    Args:
        target_date: date object
        table_name: Name of the table to delete from
        
    Returns:
        int: Number of rows deleted
    """
    conn = get_connection()
    cur = conn.cursor()
    
    try:
        sql = f"DELETE FROM {table_name} WHERE date = %s"
        cur.execute(sql, (target_date,))
        deleted_count = cur.rowcount
        conn.commit()
        return deleted_count
    except Exception as e:
        conn.rollback()
        raise e
    finally:
        cur.close()
        conn.close()


def replace_day_data(target_date, food_entries, exercises):
    """
    Atomically replace a day's food and exercise data.
    
    Deletes existing records and inserts new ones in a single transaction.
    If anything fails, the entire operation is rolled back (old data preserved).
    
    Args:
        target_date: date object
        food_entries: list of food entry dicts (from parse_food_entry)
        exercises: list of exercise dicts (from parse_exercise)
        
    Returns:
        tuple: (deleted_food, deleted_exercise, inserted_food, inserted_exercise)
    """
    conn = get_connection()
    cur = conn.cursor()
    
    try:
        # Delete existing records for the day
        cur.execute(f"DELETE FROM {config.FOOD_ENTRIES_TABLE} WHERE date = %s", (target_date,))
        deleted_food = cur.rowcount
        
        cur.execute(f"DELETE FROM {config.EXERCISES_TABLE} WHERE date = %s", (target_date,))
        deleted_exercise = cur.rowcount
        
        # Insert all new food entries
        food_sql = f"""
            INSERT INTO {config.FOOD_ENTRIES_TABLE} (
                date, meal_type, food_name, short_name, quantity, unit,
                calories, carbohydrates, fat, protein, sodium, sugar
            ) VALUES (
                %(date)s, %(meal_type)s, %(food_name)s, %(short_name)s, %(quantity)s, %(unit)s,
                %(calories)s, %(carbohydrates)s, %(fat)s, %(protein)s, %(sodium)s, %(sugar)s
            )
        """
        for data in food_entries:
            cur.execute(food_sql, data)
        
        # Insert all exercises
        exercise_sql = f"""
            INSERT INTO {config.EXERCISES_TABLE} (
                date, exercise_type, exercise_name, calories_burned, duration_minutes,
                sets, reps_per_set, weight_per_set
            ) VALUES (
                %(date)s, %(exercise_type)s, %(exercise_name)s, %(calories_burned)s, %(duration_minutes)s,
                %(sets)s, %(reps_per_set)s, %(weight_per_set)s
            )
        """
        for data in exercises:
            cur.execute(exercise_sql, data)
        
        # Commit entire transaction (all or nothing)
        conn.commit()
        return (deleted_food, deleted_exercise, len(food_entries), len(exercises))
        
    except Exception as e:
        conn.rollback()  # Preserve original data on any failure
        raise e
    finally:
        cur.close()
        conn.close()


def upsert_food_entry(data):
    """
    Insert or update a food entry in the database.
    
    Args:
        data: Dictionary with food entry fields from parse_food_entry()
    """
    conn = get_connection()
    cur = conn.cursor()
    
    try:
        # Only include fields that are actually available from MyFitnessPal API
        sql = f"""
            INSERT INTO {config.FOOD_ENTRIES_TABLE} (
                date, meal_type, food_name, short_name, quantity, unit,
                calories, carbohydrates, fat, protein, sodium, sugar
            ) VALUES (
                %(date)s, %(meal_type)s, %(food_name)s, %(short_name)s, %(quantity)s, %(unit)s,
                %(calories)s, %(carbohydrates)s, %(fat)s, %(protein)s, %(sodium)s, %(sugar)s
            )
            ON CONFLICT (date, meal_type, food_name, quantity, unit)
            DO UPDATE SET
                short_name = EXCLUDED.short_name,
                calories = EXCLUDED.calories,
                carbohydrates = EXCLUDED.carbohydrates,
                fat = EXCLUDED.fat,
                protein = EXCLUDED.protein,
                sodium = EXCLUDED.sodium,
                sugar = EXCLUDED.sugar;
        """
        
        cur.execute(sql, data)
        conn.commit()
        
    except Exception as e:
        conn.rollback()
        raise e
    finally:
        cur.close()
        conn.close()


def upsert_exercise(data):
    """
    Insert or update an exercise entry in the database.
    
    Args:
        data: Dictionary with exercise fields from parse_exercise()
    """
    conn = get_connection()
    cur = conn.cursor()
    
    try:
        sql = f"""
            INSERT INTO {config.EXERCISES_TABLE} (
                date, exercise_type, exercise_name, calories_burned, duration_minutes,
                sets, reps_per_set, weight_per_set
            ) VALUES (
                %(date)s, %(exercise_type)s, %(exercise_name)s, %(calories_burned)s, %(duration_minutes)s,
                %(sets)s, %(reps_per_set)s, %(weight_per_set)s
            )
            ON CONFLICT ON CONSTRAINT exercises_pkey
            DO UPDATE SET
                sets = EXCLUDED.sets,
                reps_per_set = EXCLUDED.reps_per_set,
                weight_per_set = EXCLUDED.weight_per_set;
        """
        
        cur.execute(sql, data)
        conn.commit()
        
    except Exception as e:
        conn.rollback()
        raise e
    finally:
        cur.close()
        conn.close()


def upsert_weight_measurement(data):
    """
    Insert or update a weight measurement in the database.
    
    Args:
        data: Dictionary with weight measurement fields from parse_weight_measurement()
    """
    conn = get_connection()
    cur = conn.cursor()
    
    try:
        sql = f"""
            INSERT INTO {config.WEIGHT_MEASUREMENTS_TABLE} (
                date, weight_lbs
            ) VALUES (
                %(date)s, %(weight_lbs)s
            )
            ON CONFLICT (date)
            DO UPDATE SET
                weight_lbs = EXCLUDED.weight_lbs;
        """
        
        cur.execute(sql, data)
        conn.commit()
        
    except Exception as e:
        conn.rollback()
        raise e
    finally:
        cur.close()
        conn.close()