Update core/database.py

core/database.py

@@ -1,7 +1,6 @@
 """
 core/database.py
 SQLite persistence for Second Brain.
-Tables: users, life_areas, goals, tasks, user_context (AI memory)
 """
 
 import sqlite3
@@ -13,8 +12,6 @@ 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
@@ -23,7 +20,6 @@ def get_db() -> sqlite3.Connection:
 
 
 def init_db():
-    """Create all tables on first run."""
     conn = get_db()
     c = conn.cursor()
 
@@ -67,7 +63,8 @@ def init_db():
             importance      TEXT DEFAULT 'Important',
             state_of_mind   TEXT DEFAULT 'Easy',
             time_estimate   INTEGER DEFAULT 30,
-            scheduled_date  TEXT DEFAULT (date('now')),
+            scheduled_date  TEXT DEFAULT '',
+            deadline_date   TEXT DEFAULT '',
             is_completed    INTEGER DEFAULT 0,
             actual_duration INTEGER,
             is_habit        INTEGER DEFAULT 0,
@@ -78,7 +75,15 @@ def init_db():
         )
     """)
 
-    # AI-learned context stored as a JSON blob per user
+    # 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)
+
     c.execute("""
         CREATE TABLE IF NOT EXISTS user_context (
             user_id    INTEGER PRIMARY KEY,
@@ -95,7 +100,6 @@ 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."
@@ -104,10 +108,7 @@ 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!"
@@ -118,14 +119,11 @@ 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."
@@ -144,68 +142,41 @@ 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.'
 
 
@@ -217,20 +188,12 @@ 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]
 
@@ -238,12 +201,16 @@ def get_goals(user_id: int) -> list:
 # ── 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.
+    """
     conn = get_db()
     cursor = conn.execute("""
         INSERT INTO tasks
             (user_id, title, life_area, urgency, importance, state_of_mind,
-             time_estimate, scheduled_date, raw_input, is_habit, habit_interval)
-        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+             time_estimate, scheduled_date, deadline_date, raw_input, is_habit, habit_interval)
+        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
     """, (
         user_id,
         task.get("title", "Untitled"),
@@ -252,31 +219,41 @@ 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 or str(date.today()),
+        scheduled_date if scheduled_date is not None else "",   # "" = unscheduled
+        task.get("deadline_date", ""),
         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) -> list:
+              include_completed: bool = True, only_unscheduled: bool = False) -> 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()))
+        q += " AND scheduled_date = ?"; params.append(str(date.today()))
+    if only_unscheduled:
+        q += " AND (scheduled_date = '' OR scheduled_date IS NULL)"
     if not include_completed:
         q += " AND is_completed = 0"
-    q += " ORDER BY is_completed ASC, created_at DESC"
+    q += " ORDER BY is_completed ASC, CASE WHEN deadline_date = '' THEN '9999' ELSE deadline_date END ASC, created_at DESC"
     rows = conn.execute(q, params).fetchall()
     conn.close()
     return [dict(r) for r in rows]
@@ -284,92 +261,62 @@ 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 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
-        # Only spawn if original was scheduled before today
-        if h["scheduled_date"] and h["scheduled_date"] >= today:
-            continue
+        if existing: continue
+        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
 
 
@@ -378,9 +325,6 @@ 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()