database.py

"""
Designed and Created by Bahamut 0 - Ernest Hardison
"""
import sqlite3
from datetime import datetime
from typing import List, Optional, Any

class Database:
    """
    Handles all interactions with the SQLite database.
    Manages connections, migrations, and CRUD operations for habits and logs.
    """
    def __init__(self, db_name="habits.db"):
        """Initialize the database manager with a specific file name."""
        self.db_name = db_name
        self.init_db()

    def get_connection(self):
        """Creates and returns a thread-safe connection to the SQLite database."""
        return sqlite3.connect(self.db_name)

    def init_db(self):
        """
        Initializes the database schema.
        Creates 'habits' and 'logs' tables if they don't exist.
        Performs simple migrations to add new columns (start_date, end_date) to existing tables.
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            # Habits table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS habits (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    type TEXT DEFAULT 'simple', -- 'simple' or 'quantitative'
                    unit TEXT,
                    goal REAL DEFAULT 1,
                    start_date TEXT,
                    end_date TEXT,
                    archived BOOLEAN DEFAULT 0
                )
            """)

            # Simple Migration for existing tables (Try/Except to add columns)
            try:
                cursor.execute("ALTER TABLE habits ADD COLUMN start_date TEXT")
            except sqlite3.OperationalError:
                pass # Already exists
            try:
                cursor.execute("ALTER TABLE habits ADD COLUMN end_date TEXT")
            except sqlite3.OperationalError:
                pass

            # Logs table (for tracking progress)
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS logs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    habit_id INTEGER NOT NULL,
                    date TEXT NOT NULL, -- YYYY-MM-DD
                    value REAL NOT NULL,
                    FOREIGN KEY (habit_id) REFERENCES habits (id)
                )
            """)
            
            # Index for performance on date lookups (Crucial for 2-year queries)
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_logs_habit_date ON logs (habit_id, date)")
            # User Profile table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS user_profile (
                    id INTEGER PRIMARY KEY CHECK (id = 1), -- Singleton row
                    name TEXT DEFAULT '',
                    birthdate TEXT DEFAULT '',
                    phone TEXT DEFAULT '',
                    email TEXT DEFAULT '',
                    goals TEXT DEFAULT ''
                )
            """)
            conn.commit()

    def get_profile(self) -> dict:
        """Retrieves the user profile data. Returns default dict if empty."""
        with self.get_connection() as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM user_profile WHERE id = 1")
            row = cursor.fetchone()
            if row:
                return dict(row)
            return {"name": "", "birthdate": "", "phone": "", "email": "", "goals": ""}

    def update_profile(self, name: str, birthdate: str, phone: str, email: str, goals: str):
        """Updates or inserts the singleton user profile."""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""
                INSERT INTO user_profile (id, name, birthdate, phone, email, goals)
                VALUES (1, ?, ?, ?, ?, ?)
                ON CONFLICT(id) DO UPDATE SET
                    name=excluded.name,
                    birthdate=excluded.birthdate,
                    phone=excluded.phone,
                    email=excluded.email,
                    goals=excluded.goals
            """, (name, birthdate, phone, email, goals))

    def add_habit(self, name: str, type: str = 'simple', unit: str = '', goal: float = 1, start_date: str = None, end_date: str = None) -> int:
        """
        Inserts a new habit into the database.
        
        Args:
            name: The name of the habit.
            type: 'simple' (yes/no) or 'quantitative' (numeric).
            unit: Unit of measurement for quantitative habits (e.g., 'ml', 'min').
            goal: Daily target value.
            start_date: Date to start tracking (YYYY-MM-DD). Defaults to today.
            end_date: Optional end date (YYYY-MM-DD).
            
        Returns:
            int: The ID of the newly created habit.
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            created_at = datetime.now().isoformat()
            if not start_date:
                start_date = datetime.now().strftime("%Y-%m-%d")
                
            cursor.execute(
                "INSERT INTO habits (name, created_at, type, unit, goal, start_date, end_date) VALUES (?, ?, ?, ?, ?, ?, ?)",
                (name, created_at, type, unit, goal, start_date, end_date)
            )
            return cursor.lastrowid

    def get_all_habits(self) -> List[Any]:
        """
        Retrieves all active (non-archived) habits, ordered by creation date (newest first).
        
        Returns:
            List[sqlite3.Row]: A list of row objects representing the habits.
        """
        with self.get_connection() as conn:
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM habits WHERE archived = 0 ORDER BY created_at DESC")
            return [dict(row) for row in cursor.fetchall()]

    def log_habit(self, habit_id: int, date: str, value: float):
        """
        Upsert logic: If a log exists for this date, update it. If not, insert it.
        For simple habits (bool), value usually toggles, but here we enforce value.
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            # Check existing
            cursor.execute("SELECT id, value FROM logs WHERE habit_id = ? AND date = ?", (habit_id, date))
            existing = cursor.fetchone()
            
            if existing:
                # Update
                new_value = existing[1] + value # Additive for quantitative
                if value == 1 and existing[1] == 1: # Toggle logic for simple items handled in UI usually, but DB just stores
                    pass 
                
                # For this implementation, we will treat 'log' as 'set value' or 'add value'
                # Let's assume the UI sends the DELTA or the FINAL value.
                # To be safely robust, let's implement 'set_log' and 'add_log'.
                # This function will be 'add_log' behavior for quantitative, 'set' for simple?
                # Let's simplify: log_habit sets the value for specific logs, or increments.
                # Actually, simpler: DELETE if 0, INSERT/UPDATE otherwise.
                pass

    def toggle_simple_habit(self, habit_id: int, date: str) -> bool:
        """Toggles a simple habit for a date. Returns new state (True/False)."""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT id FROM logs WHERE habit_id = ? AND date = ?", (habit_id, date))
            existing = cursor.fetchone()
            
            if existing:
                cursor.execute("DELETE FROM logs WHERE id = ?", (existing[0],))
                return False
            else:
                cursor.execute("INSERT INTO logs (habit_id, date, value) VALUES (?, ?, ?)", (habit_id, date, 1.0))
                return True

    def update_log_value(self, habit_id: int, date: str, value: float):
        """Sets a specific value for a date (for quantitative habits)."""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            if value <= 0:
                cursor.execute("DELETE FROM logs WHERE habit_id = ? AND date = ?", (habit_id, date))
            else:
                cursor.execute("""
                    INSERT INTO logs (habit_id, date, value) 
                    VALUES (?, ?, ?)
                    ON CONFLICT(id) DO UPDATE SET value = ? -- Wait, no unique constraint on (habit_id, date) yet?
                """, (habit_id, date, value, value)) 
                # SQLite upsert requires unique index/constraint.
                # Let's fix the schema or handle logic manually. 
                # Manual is safer across sqlite versions if unsure.
                
                cursor.execute("SELECT id FROM logs WHERE habit_id = ? AND date = ?", (habit_id, date))
                row = cursor.fetchone()
                if row:
                    cursor.execute("UPDATE logs SET value = ? WHERE id = ?", (value, row[0]))
                else:
                    cursor.execute("INSERT INTO logs (habit_id, date, value) VALUES (?, ?, ?)", (habit_id, date, value))

    def get_logs_for_habits(self, habit_ids: List[int], start_date: str, end_date: str) -> dict:
        """
        Efficient batch fetch for dashboard.
        Returns entries: { habit_id: { date: value } }
        """
        if not habit_ids:
            return {}
            
        placeholders = ','.join('?' * len(habit_ids))
        with self.get_connection() as conn:
            cursor = conn.cursor()
            query = f"""
                SELECT habit_id, date, value 
                FROM logs 
                WHERE habit_id IN ({placeholders}) 
                AND date BETWEEN ? AND ?
            """
            params = habit_ids + [start_date, end_date]
            cursor.execute(query, params)
            
            result = {hid: {} for hid in habit_ids}
            for row in cursor.fetchall():
                h_id, date, val = row
                result[h_id][date] = val
            return result

    def delete_habit(self, habit_id: int):
        """
        Soft-deletes a habit by setting its 'archived' flag to 1.
        
        Args:
            habit_id: The ID of the habit to archive.
        """
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("UPDATE habits SET archived = 1 WHERE id = ?", (habit_id,))

db = Database()