Update core/database.py

core/database.py

@@ -1,6 +1,7 @@
 """
 core/database.py
 SQLite persistence for Second Brain.
+Tables: users, life_areas, goals, tasks, user_context (AI memory)
 """
 
 import sqlite3
@@ -12,6 +13,8 @@ from typing import Optional
 DB_PATH = "second_brain.db"
 
 
+# ── Connection ────────────────────────────────────────────────────────────────
+
 def get_db() -> sqlite3.Connection:
     conn = sqlite3.connect(DB_PATH)
     conn.row_factory = sqlite3.Row
@@ -20,6 +23,7 @@ def get_db() -> sqlite3.Connection:
 
 
 def init_db():
+    """Create all tables on first run."""
     conn = get_db()
     c = conn.cursor()
 
@@ -75,15 +79,7 @@ def init_db():
         )
     """)
 
-    # Migrate existing DBs — add deadline_date if missing
-    try:
-        c.execute("ALTER TABLE tasks ADD COLUMN deadline_date TEXT DEFAULT ''")
-    except Exception:
-        pass  # column already exists
-
-    # Migrate existing DBs — clear default date so old tasks aren't shown as today
-    # (only affects new installs; existing data keeps its dates)
-
+    # AI-learned context stored as a JSON blob per user
     c.execute("""
         CREATE TABLE IF NOT EXISTS user_context (
             user_id    INTEGER PRIMARY KEY,
@@ -93,6 +89,12 @@ def init_db():
         )
     """)
 
+    # Migrate: add new columns to existing databases without breaking them
+    for _col in ["deadline_date", "scheduled_date"]:
+        try:
+            conn.execute(f"ALTER TABLE tasks ADD COLUMN {_col} TEXT DEFAULT ''")
+        except Exception:
+            pass  # column already exists
     conn.commit()
     conn.close()
 
@@ -100,6 +102,7 @@ def init_db():
 # ── Auth ──────────────────────────────────────────────────────────────────────
 
 def register_user(username: str, password: str) -> tuple:
+    """Returns (user_id, message). user_id is None on failure."""
     username = username.strip().lower()
     if not username or not password:
         return None, "Username and password cannot be empty."
@@ -108,7 +111,10 @@ def register_user(username: str, password: str) -> tuple:
     conn = get_db()
     try:
         pw_hash = bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()
-        conn.execute("INSERT INTO users (username, password_hash) VALUES (?, ?)", (username, pw_hash))
+        conn.execute(
+            "INSERT INTO users (username, password_hash) VALUES (?, ?)",
+            (username, pw_hash)
+        )
         conn.commit()
         row = conn.execute("SELECT id FROM users WHERE username = ?", (username,)).fetchone()
         return row["id"], "Account created!"
@@ -119,11 +125,14 @@ def register_user(username: str, password: str) -> tuple:
 
 
 def login_user(username: str, password: str) -> tuple:
+    """Returns (user_id, message). user_id is None on failure."""
     username = username.strip().lower()
     if not username or not password:
         return None, "Please enter your credentials."
     conn = get_db()
-    row = conn.execute("SELECT id, password_hash FROM users WHERE username = ?", (username,)).fetchone()
+    row = conn.execute(
+        "SELECT id, password_hash FROM users WHERE username = ?", (username,)
+    ).fetchone()
     conn.close()
     if not row:
         return None, "Username not found."
@@ -142,41 +151,68 @@ def get_username(user_id: int) -> str:
 # ── Life Areas ────────────────────────────────────────────────────────────────
 
 DEFAULT_AREAS = [
-    ("Work", "#4F8EF7"), ("Health", "#4CAF87"), ("Finance", "#F7A84F"),
-    ("Learning", "#A855F7"), ("Personal", "#EC4899"), ("Family", "#F59E0B"),
+    ("Work",     "#4F8EF7"),
+    ("Health",   "#4CAF87"),
+    ("Finance",  "#F7A84F"),
+    ("Learning", "#A855F7"),
+    ("Personal", "#EC4899"),
+    ("Family",   "#F59E0B"),
 ]
 
+
 def create_default_life_areas(user_id: int):
     conn = get_db()
     for name, color in DEFAULT_AREAS:
-        conn.execute("INSERT INTO life_areas (user_id, name, color) VALUES (?, ?, ?)", (user_id, name, color))
+        conn.execute(
+            "INSERT INTO life_areas (user_id, name, color) VALUES (?, ?, ?)",
+            (user_id, name, color)
+        )
     conn.commit()
     conn.close()
 
+
 def get_life_areas(user_id: int) -> list:
     conn = get_db()
-    rows = conn.execute("SELECT id, name, color FROM life_areas WHERE user_id = ? ORDER BY id", (user_id,)).fetchall()
+    rows = conn.execute(
+        "SELECT id, name, color FROM life_areas WHERE user_id = ? ORDER BY id",
+        (user_id,)
+    ).fetchall()
     conn.close()
     return [dict(r) for r in rows]
 
+
 def get_life_area_names(user_id: int) -> list:
     return [a["name"] for a in get_life_areas(user_id)]
 
+
 def add_life_area(user_id: int, name: str, color: str = "#6366f1") -> tuple:
     name = name.strip()
-    if not name: return False, "Name cannot be empty."
+    if not name:
+        return False, "Name cannot be empty."
     conn = get_db()
-    exists = conn.execute("SELECT id FROM life_areas WHERE user_id = ? AND LOWER(name) = LOWER(?)", (user_id, name)).fetchone()
+    exists = conn.execute(
+        "SELECT id FROM life_areas WHERE user_id = ? AND LOWER(name) = LOWER(?)",
+        (user_id, name)
+    ).fetchone()
     if exists:
-        conn.close(); return False, f'"{name}" already exists.'
-    conn.execute("INSERT INTO life_areas (user_id, name, color) VALUES (?, ?, ?)", (user_id, name, color))
-    conn.commit(); conn.close()
+        conn.close()
+        return False, f'"{name}" already exists.'
+    conn.execute(
+        "INSERT INTO life_areas (user_id, name, color) VALUES (?, ?, ?)",
+        (user_id, name, color)
+    )
+    conn.commit()
+    conn.close()
     return True, f'"{name}" added.'
 
+
 def delete_life_area(user_id: int, name: str) -> tuple:
     conn = get_db()
-    conn.execute("DELETE FROM life_areas WHERE user_id = ? AND name = ?", (user_id, name))
-    conn.commit(); conn.close()
+    conn.execute(
+        "DELETE FROM life_areas WHERE user_id = ? AND name = ?", (user_id, name)
+    )
+    conn.commit()
+    conn.close()
     return True, f'"{name}" removed.'
 
 
@@ -188,23 +224,27 @@ def save_goals(user_id: int, goals_text: str):
     for line in goals_text.strip().splitlines():
         line = line.strip("•- ").strip()
         if line:
-            conn.execute("INSERT INTO user_goals (user_id, goal_text) VALUES (?, ?)", (user_id, line))
-    conn.commit(); conn.close()
+            conn.execute(
+                "INSERT INTO user_goals (user_id, goal_text) VALUES (?, ?)",
+                (user_id, line)
+            )
+    conn.commit()
+    conn.close()
+
 
 def get_goals(user_id: int) -> list:
     conn = get_db()
-    rows = conn.execute("SELECT goal_text FROM user_goals WHERE user_id = ? ORDER BY id", (user_id,)).fetchall()
+    rows = conn.execute(
+        "SELECT goal_text FROM user_goals WHERE user_id = ? ORDER BY id",
+        (user_id,)
+    ).fetchall()
     conn.close()
     return [r["goal_text"] for r in rows]
 
 
 # ── Tasks ─────────────────────────────────────────────────────────────────────
 
-def save_task(user_id: int, task: dict, scheduled_date: str = None) -> int:
-    """
-    Save a task. scheduled_date=None means unscheduled (awaiting AI assignment).
-    Pass scheduled_date="" explicitly to also leave unscheduled.
-    """
+def save_task(user_id: int, task: dict, scheduled_date: str = "") -> int:
     conn = get_db()
     cursor = conn.execute("""
         INSERT INTO tasks
@@ -219,41 +259,32 @@ def save_task(user_id: int, task: dict, scheduled_date: str = None) -> int:
         task.get("importance", "Important"),
         task.get("state_of_mind", "Easy"),
         int(task.get("time_estimate") or 30),
-        scheduled_date if scheduled_date is not None else "",   # "" = unscheduled
-        task.get("deadline_date", ""),
+        scheduled_date or "",
+        task.get("deadline_date", "") or "",
         task.get("raw_input", ""),
         1 if task.get("is_habit") else 0,
         task.get("habit_interval", ""),
     ))
     task_id = cursor.lastrowid
-    conn.commit(); conn.close()
+    conn.commit()
+    conn.close()
     return task_id
 
 
-def assign_task_date(task_id: int, user_id: int, scheduled_date: str):
-    """Update a task's scheduled_date (called by the AI scheduler)."""
-    conn = get_db()
-    conn.execute(
-        "UPDATE tasks SET scheduled_date = ? WHERE id = ? AND user_id = ?",
-        (scheduled_date, task_id, user_id)
-    )
-    conn.commit(); conn.close()
-
-
 def get_tasks(user_id: int, filter_area: str = "All", only_today: bool = False,
-              include_completed: bool = True, only_unscheduled: bool = False) -> list:
+              include_completed: bool = True) -> list:
     conn = get_db()
     q = "SELECT * FROM tasks WHERE user_id = ?"
     params = [user_id]
     if filter_area and filter_area != "All":
-        q += " AND life_area = ?"; params.append(filter_area)
+        q += " AND life_area = ?"
+        params.append(filter_area)
     if only_today:
-        q += " AND scheduled_date = ?"; params.append(str(date.today()))
-    if only_unscheduled:
-        q += " AND (scheduled_date = '' OR scheduled_date IS NULL)"
+        q += " AND scheduled_date = ?"
+        params.append(str(date.today()))
     if not include_completed:
         q += " AND is_completed = 0"
-    q += " ORDER BY is_completed ASC, CASE WHEN deadline_date = '' THEN '9999' ELSE deadline_date END ASC, created_at DESC"
+    q += " ORDER BY is_completed ASC, created_at DESC"
     rows = conn.execute(q, params).fetchall()
     conn.close()
     return [dict(r) for r in rows]
@@ -261,62 +292,103 @@ def get_tasks(user_id: int, filter_area: str = "All", only_today: bool = False,
 
 def toggle_task_complete(task_id: int, user_id: int, actual_duration: int = None):
     conn = get_db()
-    task = conn.execute("SELECT is_completed FROM tasks WHERE id = ? AND user_id = ?", (task_id, user_id)).fetchone()
+    task = conn.execute(
+        "SELECT is_completed FROM tasks WHERE id = ? AND user_id = ?", (task_id, user_id)
+    ).fetchone()
     if task:
         new_status = 1 - task["is_completed"]
         if actual_duration and new_status == 1:
-            conn.execute("UPDATE tasks SET is_completed = ?, actual_duration = ? WHERE id = ? AND user_id = ?",
-                         (new_status, actual_duration, task_id, user_id))
+            conn.execute(
+                "UPDATE tasks SET is_completed = ?, actual_duration = ? WHERE id = ? AND user_id = ?",
+                (new_status, actual_duration, task_id, user_id)
+            )
         else:
-            conn.execute("UPDATE tasks SET is_completed = ? WHERE id = ? AND user_id = ?",
-                         (new_status, task_id, user_id))
-    conn.commit(); conn.close()
+            conn.execute(
+                "UPDATE tasks SET is_completed = ? WHERE id = ? AND user_id = ?",
+                (new_status, task_id, user_id)
+            )
+    conn.commit()
+    conn.close()
 
 
 def delete_task(task_id: int, user_id: int):
     conn = get_db()
     conn.execute("DELETE FROM tasks WHERE id = ? AND user_id = ?", (task_id, user_id))
-    conn.commit(); conn.close()
+    conn.commit()
+    conn.close()
+
+
+def assign_task_date(task_id: int, user_id: int, scheduled_date: str):
+    """Assign a scheduled date to an existing task (called by the AI planner)."""
+    conn = get_db()
+    conn.execute(
+        "UPDATE tasks SET scheduled_date = ? WHERE id = ? AND user_id = ?",
+        (scheduled_date, task_id, user_id)
+    )
+    conn.commit()
+    conn.close()
 
 
 def get_today_stats(user_id: int) -> dict:
     tasks = get_tasks(user_id, only_today=True)
-    total = len(tasks); done = sum(1 for t in tasks if t["is_completed"])
+    total = len(tasks)
+    done  = sum(1 for t in tasks if t["is_completed"])
     return {"total": total, "done": done, "remaining": total - done}
 
 
 # ── Habit recurrence ──────────────────────────────────────────────────────────
 
 def spawn_due_habits(user_id: int):
+    """
+    Check all habit tasks. If a habit's scheduled_date < today and
+    it's not already scheduled for today, create a fresh copy for today.
+    Called on login / tab load.
+    """
     today = str(date.today())
-    conn  = get_db()
-    habits = conn.execute("SELECT * FROM tasks WHERE user_id = ? AND is_habit = 1", (user_id,)).fetchall()
+    conn = get_db()
+    habits = conn.execute(
+        "SELECT * FROM tasks WHERE user_id = ? AND is_habit = 1",
+        (user_id,)
+    ).fetchall()
+
     for h in habits:
+        # Check if already exists today
         existing = conn.execute(
             "SELECT id FROM tasks WHERE user_id = ? AND title = ? AND is_habit = 1 AND scheduled_date = ?",
             (user_id, h["title"], today)
         ).fetchone()
-        if existing: continue
-        if h["scheduled_date"] and h["scheduled_date"] >= today: continue
+        if existing:
+            continue
+        # Only spawn if original was scheduled before today
+        if h["scheduled_date"] and h["scheduled_date"] >= today:
+            continue
         conn.execute("""
             INSERT INTO tasks (user_id, title, life_area, urgency, importance,
                                state_of_mind, time_estimate, scheduled_date, is_habit,
                                habit_interval, raw_input)
             VALUES (?, ?, ?, ?, ?, ?, ?, ?, 1, ?, ?)
-        """, (user_id, h["title"], h["life_area"], "Habit", h["importance"],
-              h["state_of_mind"], h["time_estimate"], today, h["habit_interval"], h["raw_input"]))
-    conn.commit(); conn.close()
+        """, (
+            user_id, h["title"], h["life_area"], "Habit",
+            h["importance"], h["state_of_mind"], h["time_estimate"],
+            today, h["habit_interval"], h["raw_input"]
+        ))
+    conn.commit()
+    conn.close()
 
 
 # ── AI Context ────────────────────────────────────────────────────────────────
 
 def load_user_context(user_id: int) -> Optional[dict]:
     conn = get_db()
-    row  = conn.execute("SELECT context FROM user_context WHERE user_id = ?", (user_id,)).fetchone()
+    row = conn.execute(
+        "SELECT context FROM user_context WHERE user_id = ?", (user_id,)
+    ).fetchone()
     conn.close()
     if row:
-        try: return json.loads(row["context"])
-        except Exception: return None
+        try:
+            return json.loads(row["context"])
+        except Exception:
+            return None
     return None
 
 
@@ -325,6 +397,9 @@ def save_user_context(user_id: int, context: dict):
     conn.execute("""
         INSERT INTO user_context (user_id, context, updated_at)
         VALUES (?, ?, datetime('now'))
-        ON CONFLICT(user_id) DO UPDATE SET context = excluded.context, updated_at = excluded.updated_at
+        ON CONFLICT(user_id) DO UPDATE SET
+            context    = excluded.context,
+            updated_at = excluded.updated_at
     """, (user_id, json.dumps(context)))
-    conn.commit(); conn.close()
+    conn.commit()
+    conn.close()