core/database.py

"""
core/database.py
SQLite persistence for Second Brain.
Tables: users, life_areas, goals, tasks, user_context (AI memory)
"""

import sqlite3
import json
import bcrypt
from datetime import datetime, date
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
    conn.execute("PRAGMA foreign_keys = ON")
    return conn


def init_db():
    """Create all tables on first run."""
    conn = get_db()
    c = conn.cursor()

    c.execute("""
        CREATE TABLE IF NOT EXISTS users (
            id            INTEGER PRIMARY KEY AUTOINCREMENT,
            username      TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            created_at    TEXT DEFAULT (datetime('now'))
        )
    """)

    c.execute("""
        CREATE TABLE IF NOT EXISTS life_areas (
            id         INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id    INTEGER NOT NULL,
            name       TEXT NOT NULL,
            color      TEXT DEFAULT '#6366f1',
            created_at TEXT DEFAULT (datetime('now')),
            FOREIGN KEY (user_id) REFERENCES users(id)
        )
    """)

    c.execute("""
        CREATE TABLE IF NOT EXISTS user_goals (
            id         INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id    INTEGER NOT NULL,
            goal_text  TEXT NOT NULL,
            created_at TEXT DEFAULT (datetime('now')),
            FOREIGN KEY (user_id) REFERENCES users(id)
        )
    """)

    c.execute("""
        CREATE TABLE IF NOT EXISTS tasks (
            id              INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id         INTEGER NOT NULL,
            title           TEXT NOT NULL,
            life_area       TEXT DEFAULT '',
            urgency         TEXT DEFAULT 'Not Urgent',
            importance      TEXT DEFAULT 'Important',
            state_of_mind   TEXT DEFAULT 'Easy',
            time_estimate   INTEGER DEFAULT 30,
            scheduled_date  TEXT DEFAULT '',
            deadline_date   TEXT DEFAULT '',
            is_completed    INTEGER DEFAULT 0,
            actual_duration INTEGER,
            is_habit        INTEGER DEFAULT 0,
            habit_interval  TEXT DEFAULT '',
            raw_input       TEXT DEFAULT '',
            created_at      TEXT DEFAULT (datetime('now')),
            FOREIGN KEY (user_id) REFERENCES users(id)
        )
    """)

    # AI-learned context stored as a JSON blob per user
    c.execute("""
        CREATE TABLE IF NOT EXISTS user_context (
            user_id    INTEGER PRIMARY KEY,
            context    TEXT NOT NULL,
            updated_at TEXT DEFAULT (datetime('now')),
            FOREIGN KEY (user_id) REFERENCES users(id)
        )
    """)

    # 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()


# ── 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."
    if len(password) < 6:
        return None, "Password must be at least 6 characters."
    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.commit()
        row = conn.execute("SELECT id FROM users WHERE username = ?", (username,)).fetchone()
        return row["id"], "Account created!"
    except sqlite3.IntegrityError:
        return None, "Username already taken."
    finally:
        conn.close()


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()
    conn.close()
    if not row:
        return None, "Username not found."
    if not bcrypt.checkpw(password.encode(), row["password_hash"].encode()):
        return None, "Incorrect password."
    return row["id"], f"Welcome back, {username}!"


def get_username(user_id: int) -> str:
    conn = get_db()
    row = conn.execute("SELECT username FROM users WHERE id = ?", (user_id,)).fetchone()
    conn.close()
    return row["username"].capitalize() if row else "User"


# ── Life Areas ────────────────────────────────────────────────────────────────

DEFAULT_AREAS = [
    ("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.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()
    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."
    conn = get_db()
    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()
    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()
    return True, f'"{name}" removed.'


# ── Goals ─────────────────────────────────────────────────────────────────────

def save_goals(user_id: int, goals_text: str):
    conn = get_db()
    conn.execute("DELETE FROM user_goals WHERE user_id = ?", (user_id,))
    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()


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()
    conn.close()
    return [r["goal_text"] for r in rows]


# ── Tasks ─────────────────────────────────────────────────────────────────────

def save_task(user_id: int, task: dict, scheduled_date: str = "") -> int:
    conn = get_db()
    cursor = conn.execute("""
        INSERT INTO tasks
            (user_id, title, life_area, urgency, importance, state_of_mind,
             time_estimate, scheduled_date, deadline_date, raw_input, is_habit, habit_interval)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    """, (
        user_id,
        task.get("title", "Untitled"),
        task.get("life_area", ""),
        task.get("urgency", "Not Urgent"),
        task.get("importance", "Important"),
        task.get("state_of_mind", "Easy"),
        int(task.get("time_estimate") or 30),
        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()
    return task_id


def get_tasks(user_id: int, filter_area: str = "All", only_today: bool = False,
              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)
    if only_today:
        q += " AND scheduled_date = ?"
        params.append(str(date.today()))
    if not include_completed:
        q += " AND is_completed = 0"
    q += " ORDER BY is_completed ASC, created_at DESC"
    rows = conn.execute(q, params).fetchall()
    conn.close()
    return [dict(r) for r in rows]


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()
    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)
            )
        else:
            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()


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"])
    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()

    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
        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()


# ── 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()
    conn.close()
    if row:
        try:
            return json.loads(row["context"])
        except Exception:
            return None
    return None


def save_user_context(user_id: int, context: dict):
    conn = get_db()
    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
    """, (user_id, json.dumps(context)))
    conn.commit()
    conn.close()