core/database.py

"""
core/database.py
SQLite persistence for Second Brain.
"""

import sqlite3
import json
import bcrypt
from datetime import datetime, date
from typing import Optional

DB_PATH = "second_brain.db"


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

    # 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,
            context    TEXT NOT NULL,
            updated_at TEXT DEFAULT (datetime('now')),
            FOREIGN KEY (user_id) REFERENCES users(id)
        )
    """)

    conn.commit()
    conn.close()


# ── Auth ──────────────────────────────────────────────────────────────────────

def register_user(username: str, password: str) -> tuple:
    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:
    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 = 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, 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 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()
    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:
    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 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, 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]


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