database.py

from __future__ import annotations

import sqlite3
from datetime import date, datetime
from pathlib import Path
from typing import Optional

DATETIME_FMT = "%Y-%m-%d %H:%M:%S"


class TimeTrackerDB:
    def __init__(self, db_path: Path) -> None:
        self.db_path = db_path
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        self.conn = sqlite3.connect(self.db_path)
        self.conn.row_factory = sqlite3.Row
        self.conn.execute("PRAGMA foreign_keys = ON;")
        self._init_schema()
        self._ensure_default_project()

    def _init_schema(self) -> None:
        self.conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS projects (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL UNIQUE,
                created_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS entries (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE RESTRICT,
                description TEXT NOT NULL DEFAULT '',
                start_time TEXT NOT NULL,
                end_time TEXT NOT NULL,
                duration_seconds INTEGER NOT NULL CHECK (duration_seconds >= 0),
                created_at TEXT NOT NULL
            );

            CREATE INDEX IF NOT EXISTS idx_entries_start_time ON entries(start_time);
            CREATE INDEX IF NOT EXISTS idx_entries_project_id ON entries(project_id);

            CREATE TABLE IF NOT EXISTS active_session (
                id INTEGER PRIMARY KEY CHECK (id = 1),
                project_id INTEGER NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
                description TEXT NOT NULL DEFAULT '',
                start_time TEXT NOT NULL
            );
            """
        )
        self.conn.commit()

    def _ensure_default_project(self) -> None:
        count = self.conn.execute("SELECT COUNT(*) AS cnt FROM projects").fetchone()["cnt"]
        if count == 0:
            now = datetime.now().strftime(DATETIME_FMT)
            self.conn.execute(
                "INSERT INTO projects(name, created_at) VALUES(?, ?)",
                ("General", now),
            )
            self.conn.commit()

    def close(self) -> None:
        self.conn.close()

    def list_projects(self) -> list[sqlite3.Row]:
        rows = self.conn.execute(
            "SELECT id, name FROM projects ORDER BY lower(name) ASC"
        ).fetchall()
        return list(rows)

    def add_project(self, name: str) -> None:
        clean = name.strip()
        if not clean:
            raise ValueError("Project name cannot be empty.")
        now = datetime.now().strftime(DATETIME_FMT)
        self.conn.execute(
            "INSERT INTO projects(name, created_at) VALUES(?, ?)",
            (clean, now),
        )
        self.conn.commit()

    def delete_project(self, project_id: int) -> None:
        has_entries = self.conn.execute(
            "SELECT COUNT(*) AS cnt FROM entries WHERE project_id = ?",
            (project_id,),
        ).fetchone()["cnt"]
        if has_entries:
            raise ValueError("Project cannot be deleted because it has time entries.")

        self.conn.execute("DELETE FROM active_session WHERE project_id = ?", (project_id,))
        cur = self.conn.execute("DELETE FROM projects WHERE id = ?", (project_id,))
        self.conn.commit()
        if cur.rowcount == 0:
            raise ValueError("Project not found.")

    def create_entry(
        self,
        project_id: int,
        description: str,
        start_dt: datetime,
        end_dt: datetime,
    ) -> None:
        duration = int((end_dt - start_dt).total_seconds())
        if duration <= 0:
            raise ValueError("End time must be later than start time.")

        now = datetime.now().strftime(DATETIME_FMT)
        self.conn.execute(
            """
            INSERT INTO entries(project_id, description, start_time, end_time, duration_seconds, created_at)
            VALUES(?, ?, ?, ?, ?, ?)
            """,
            (
                project_id,
                description.strip(),
                start_dt.strftime(DATETIME_FMT),
                end_dt.strftime(DATETIME_FMT),
                duration,
                now,
            ),
        )
        self.conn.commit()

    def list_entries(
        self,
        from_date: Optional[date] = None,
        to_date: Optional[date] = None,
    ) -> list[sqlite3.Row]:
        query = """
            SELECT
                e.id,
                e.project_id,
                p.name AS project_name,
                e.description,
                e.start_time,
                e.end_time,
                e.duration_seconds
            FROM entries e
            JOIN projects p ON p.id = e.project_id
        """
        params: list[str] = []
        conditions = []

        if from_date is not None:
            conditions.append("date(e.start_time) >= ?")
            params.append(from_date.isoformat())
        if to_date is not None:
            conditions.append("date(e.start_time) <= ?")
            params.append(to_date.isoformat())

        if conditions:
            query += " WHERE " + " AND ".join(conditions)

        query += " ORDER BY e.start_time DESC"

        rows = self.conn.execute(query, params).fetchall()
        return list(rows)

    def delete_entry(self, entry_id: int) -> None:
        cur = self.conn.execute("DELETE FROM entries WHERE id = ?", (entry_id,))
        self.conn.commit()
        if cur.rowcount == 0:
            raise ValueError("Entry not found.")

    def update_entry(
        self,
        entry_id: int,
        project_id: int,
        description: str,
        start_dt: datetime,
        end_dt: datetime,
    ) -> None:
        duration = int((end_dt - start_dt).total_seconds())
        if duration <= 0:
            raise ValueError("End time must be later than start time.")

        cur = self.conn.execute(
            """
            UPDATE entries
            SET
                project_id = ?,
                description = ?,
                start_time = ?,
                end_time = ?,
                duration_seconds = ?
            WHERE id = ?
            """,
            (
                project_id,
                description.strip(),
                start_dt.strftime(DATETIME_FMT),
                end_dt.strftime(DATETIME_FMT),
                duration,
                entry_id,
            ),
        )
        self.conn.commit()
        if cur.rowcount == 0:
            raise ValueError("Entry not found.")

    def set_active_session(
        self,
        project_id: int,
        description: str,
        start_dt: datetime,
    ) -> None:
        self.conn.execute(
            """
            INSERT INTO active_session(id, project_id, description, start_time)
            VALUES(1, ?, ?, ?)
            ON CONFLICT(id) DO UPDATE SET
                project_id = excluded.project_id,
                description = excluded.description,
                start_time = excluded.start_time
            """,
            (project_id, description.strip(), start_dt.strftime(DATETIME_FMT)),
        )
        self.conn.commit()

    def clear_active_session(self) -> None:
        self.conn.execute("DELETE FROM active_session WHERE id = 1")
        self.conn.commit()

    def get_active_session(self) -> Optional[sqlite3.Row]:
        row = self.conn.execute(
            """
            SELECT
                s.project_id,
                p.name AS project_name,
                s.description,
                s.start_time
            FROM active_session s
            JOIN projects p ON p.id = s.project_id
            WHERE s.id = 1
            """
        ).fetchone()
        return row

    def total_seconds_between(self, start_dt: datetime, end_dt: datetime) -> int:
        row = self.conn.execute(
            """
            SELECT COALESCE(SUM(duration_seconds), 0) AS total
            FROM entries
            WHERE start_time >= ? AND start_time < ?
            """,
            (start_dt.strftime(DATETIME_FMT), end_dt.strftime(DATETIME_FMT)),
        ).fetchone()
        return int(row["total"])