db.py

import os
import sqlite3
from typing import Dict, List, Optional, Tuple


class Database:
    """Simple SQLite wrapper for caching emails and codes locally."""

    def __init__(self, db_path: Optional[str] = None) -> None:
        self.db_path = db_path or os.path.join(os.path.dirname(__file__), 'emails.db')
        self._connection = sqlite3.connect(self.db_path, check_same_thread=False)
        self._connection.row_factory = sqlite3.Row
        self._ensure_schema()

    def _ensure_schema(self) -> None:
        cursor = self._connection.cursor()
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS accounts (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT NOT NULL UNIQUE,
                host TEXT
            );
            """
        )
        cursor.execute(
            """
            CREATE TABLE IF NOT EXISTS emails (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                account_username TEXT NOT NULL,
                message_uid TEXT,
                message_id TEXT,
                subject TEXT,
                sender TEXT,
                recipient TEXT,
                date TEXT,
                body_snippet TEXT,
                code TEXT,
                is_unread INTEGER DEFAULT 1,
                created_at TEXT DEFAULT (datetime('now')),
                UNIQUE(account_username, message_uid)
            );
            """
        )
        try:
            cursor.execute("ALTER TABLE emails ADD COLUMN body_full TEXT")
        except sqlite3.OperationalError:
            pass
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_emails_account_date ON emails(account_username, date DESC);")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_emails_code ON emails(code);")
        self._connection.commit()

    def upsert_account(self, username: str, host: Optional[str]) -> None:
        cursor = self._connection.cursor()
        cursor.execute(
            """
            INSERT INTO accounts (username, host)
            VALUES (?, ?)
            ON CONFLICT(username) DO UPDATE SET host=excluded.host
            """,
            (username, host),
        )
        self._connection.commit()

    def insert_email(self, email_record: Dict[str, Optional[str]]) -> bool:
        """Insert new email if not exists; update existing otherwise.
        Returns True only when a NEW row is inserted (used to gate UI 'new_email' signal).
        """
        cursor = self._connection.cursor()
        # First, attempt insert-or-ignore using UID as unique key
        cursor.execute(
            """
            INSERT OR IGNORE INTO emails (
                account_username, message_uid, message_id, subject, sender, recipient, date,
                body_snippet, code, is_unread, body_full
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                email_record.get('account_username'),
                email_record.get('message_uid'),
                email_record.get('message_id'),
                email_record.get('subject'),
                email_record.get('sender'),
                email_record.get('recipient'),
                email_record.get('date'),
                email_record.get('body_snippet'),
                email_record.get('code'),
                1 if email_record.get('is_unread', True) else 0,
                email_record.get('body_full'),
            ),
        )
        inserted = cursor.rowcount == 1
        if not inserted:
            # On conflict, update key fields but do not flip code if new is None
            cursor.execute(
                """
                UPDATE emails
                SET subject=?, sender=?, recipient=?, date=?, body_snippet=?, body_full=?,
                    code=COALESCE(?, code), is_unread=?
                WHERE account_username=? AND message_uid=?
                """,
                (
                    email_record.get('subject'),
                    email_record.get('sender'),
                    email_record.get('recipient'),
                    email_record.get('date'),
                    email_record.get('body_snippet'),
                    email_record.get('body_full'),
                    email_record.get('code'),
                    1 if email_record.get('is_unread', True) else 0,
                    email_record.get('account_username'),
                    email_record.get('message_uid'),
                ),
            )
        self._connection.commit()
        return inserted

    def mark_email_read(self, account_username: str, message_uid: Optional[str], message_id: Optional[str]) -> None:
        cursor = self._connection.cursor()
        if message_uid:
            cursor.execute(
                "UPDATE emails SET is_unread=0 WHERE account_username=? AND message_uid=?",
                (account_username, message_uid),
            )
        elif message_id:
            cursor.execute(
                "UPDATE emails SET is_unread=0 WHERE account_username=? AND message_id=?",
                (account_username, message_id),
            )
        self._connection.commit()

    def mark_all_read(self, account_username: Optional[str] = None) -> None:
        cursor = self._connection.cursor()
        if account_username:
            cursor.execute("UPDATE emails SET is_unread=0 WHERE account_username=?", (account_username,))
        else:
            cursor.execute("UPDATE emails SET is_unread=0")
        self._connection.commit()

    def fetch_recent_emails(self, limit: int = 200, account_username: Optional[str] = None) -> List[Dict[str, str]]:
        cursor = self._connection.cursor()
        if account_username:
            cursor.execute(
                """
                SELECT account_username, message_uid, message_id, subject, sender, recipient, date, body_snippet, code, is_unread, created_at, body_full
                FROM emails
                WHERE account_username = ?
                ORDER BY datetime(date) DESC, id DESC
                LIMIT ?
                """,
                (account_username, limit),
            )
        else:
            cursor.execute(
                """
                SELECT account_username, message_uid, message_id, subject, sender, recipient, date, body_snippet, code, is_unread, created_at, body_full
                FROM emails
                ORDER BY datetime(date) DESC, id DESC
                LIMIT ?
                """,
                (limit,),
            )
        rows = cursor.fetchall()
        return [dict(row) for row in rows]

    def fetch_accounts(self) -> List[Tuple[str, Optional[str]]]:
        cursor = self._connection.cursor()
        cursor.execute("SELECT username, host FROM accounts ORDER BY username ASC")
        rows = cursor.fetchall()
        return [(row[0], row[1]) for row in rows]

    def close(self) -> None:
        try:
            self._connection.close()
        except Exception:
            pass


__all__ = ["Database"]