app/database.py

import json
import sys
from contextlib import asynccontextmanager
from pathlib import Path
from typing import AsyncGenerator

import aiosqlite

_DDL = """
PRAGMA journal_mode=WAL;
PRAGMA foreign_keys=ON;

CREATE TABLE IF NOT EXISTS subscriptions (
    id               TEXT PRIMARY KEY,
    folder_id        TEXT NOT NULL,
    folder_name      TEXT NOT NULL,
    expiry_datetime  TEXT NOT NULL,
    notification_url TEXT NOT NULL,
    client_state     TEXT NOT NULL,
    is_active        INTEGER NOT NULL DEFAULT 1,
    created_at       TEXT NOT NULL DEFAULT (datetime('now')),
    last_renewed_at  TEXT
);

CREATE TABLE IF NOT EXISTS notifications (
    id                INTEGER PRIMARY KEY AUTOINCREMENT,
    received_at       TEXT NOT NULL DEFAULT (datetime('now')),
    subscription_id   TEXT,
    change_type       TEXT,
    resource          TEXT,
    message_id        TEXT,
    email_identifier  TEXT,
    processing_status TEXT NOT NULL DEFAULT 'pending',
    error_message     TEXT,
    storage_path      TEXT
);

CREATE TABLE IF NOT EXISTS routing_results (
    id                  INTEGER PRIMARY KEY AUTOINCREMENT,
    notification_id     INTEGER NOT NULL,
    processed_at        TEXT NOT NULL DEFAULT (datetime('now')),
    email_identifier    TEXT,
    attachment_filename TEXT,
    document_type       TEXT,
    action              TEXT,
    destination_path    TEXT,
    confidence_scores   TEXT,
    error_message       TEXT
);

CREATE TABLE IF NOT EXISTS daily_brief (
    id                  INTEGER PRIMARY KEY AUTOINCREMENT,
    created_at          TEXT NOT NULL DEFAULT (datetime('now')),
    email_identifier    TEXT NOT NULL,
    notification_id     INTEGER NOT NULL,
    sender_email        TEXT,
    subject             TEXT,
    attachment_filename TEXT,
    reason              TEXT NOT NULL,
    confidence_scores   TEXT,
    error_message       TEXT,
    included_in_report  INTEGER NOT NULL DEFAULT 0,
    report_sent_at      TEXT
);

CREATE TABLE IF NOT EXISTS client_domains (
    domain        TEXT PRIMARY KEY,
    client_name   TEXT NOT NULL,
    created_at    TEXT NOT NULL DEFAULT (datetime('now')),
    last_seen_at  TEXT,
    notes         TEXT
);
"""


async def init_db(db_path: Path) -> None:
    db_path.parent.mkdir(parents=True, exist_ok=True)
    async with aiosqlite.connect(db_path) as conn:
        await conn.executescript(_DDL)
        await conn.commit()


@asynccontextmanager
async def get_db(db_path: Path) -> AsyncGenerator[aiosqlite.Connection, None]:
    async with aiosqlite.connect(db_path) as conn:
        conn.row_factory = aiosqlite.Row
        await conn.execute("PRAGMA journal_mode=WAL")
        await conn.execute("PRAGMA foreign_keys=ON")
        yield conn


async def upsert_subscription(
    conn: aiosqlite.Connection,
    sub_id: str,
    folder_id: str,
    folder_name: str,
    expiry_datetime: str,
    notification_url: str,
    client_state: str,
) -> None:
    await conn.execute(
        """
        INSERT INTO subscriptions (id, folder_id, folder_name, expiry_datetime,
            notification_url, client_state, is_active)
        VALUES (?, ?, ?, ?, ?, ?, 1)
        ON CONFLICT(id) DO UPDATE SET
            expiry_datetime  = excluded.expiry_datetime,
            last_renewed_at  = datetime('now'),
            is_active        = 1
        """,
        (sub_id, folder_id, folder_name, expiry_datetime, notification_url, client_state),
    )
    await conn.commit()


async def get_active_subscription(conn: aiosqlite.Connection) -> dict | None:
    cursor = await conn.execute(
        "SELECT * FROM subscriptions WHERE is_active = 1 ORDER BY created_at DESC LIMIT 1"
    )
    row = await cursor.fetchone()
    return dict(row) if row else None


async def mark_subscription_inactive(conn: aiosqlite.Connection, sub_id: str) -> None:
    await conn.execute(
        "UPDATE subscriptions SET is_active = 0 WHERE id = ?", (sub_id,)
    )
    await conn.commit()


async def log_notification(
    conn: aiosqlite.Connection,
    subscription_id: str | None,
    change_type: str | None,
    resource: str | None,
    message_id: str | None,
    email_identifier: str | None = None,
) -> int:
    cursor = await conn.execute(
        """
        INSERT INTO notifications (subscription_id, change_type, resource, message_id, email_identifier)
        VALUES (?, ?, ?, ?, ?)
        """,
        (subscription_id, change_type, resource, message_id, email_identifier or message_id),
    )
    await conn.commit()
    return cursor.lastrowid


async def update_subscription_expiry(
    conn: aiosqlite.Connection, sub_id: str, expiry_datetime: str
) -> None:
    await conn.execute(
        """
        UPDATE subscriptions
        SET expiry_datetime = ?, last_renewed_at = datetime('now')
        WHERE id = ?
        """,
        (expiry_datetime, sub_id),
    )
    await conn.commit()


async def get_notification_by_email_identifier(
    conn: aiosqlite.Connection, email_identifier: str
) -> dict | None:
    cursor = await conn.execute(
        "SELECT * FROM notifications WHERE email_identifier = ? LIMIT 1",
        (email_identifier,),
    )
    row = await cursor.fetchone()
    return dict(row) if row else None


async def update_notification_status(
    conn: aiosqlite.Connection,
    notif_id: int,
    status: str,
    error_message: str | None = None,
    storage_path: str | None = None,
) -> None:
    await conn.execute(
        """
        UPDATE notifications
        SET processing_status = ?, error_message = ?, storage_path = ?
        WHERE id = ?
        """,
        (status, error_message, storage_path, notif_id),
    )
    await conn.commit()


async def get_fetched_notifications(conn: aiosqlite.Connection) -> list[dict]:
    cursor = await conn.execute(
        "SELECT * FROM notifications WHERE processing_status = 'fetched' ORDER BY received_at"
    )
    rows = await cursor.fetchall()
    return [dict(r) for r in rows]


async def set_notification_routing_status(
    conn: aiosqlite.Connection, notif_id: int, status: str
) -> None:
    await conn.execute(
        "UPDATE notifications SET processing_status = ? WHERE id = ?",
        (status, notif_id),
    )
    await conn.commit()


async def log_routing_result(
    conn: aiosqlite.Connection,
    notification_id: int,
    filename: str | None,
    doc_type: str | None,
    action: str,
    email_identifier: str | None = None,
    destination_path: str | None = None,
    confidence_scores: dict | None = None,
    error_message: str | None = None,
) -> None:
    await conn.execute(
        """
        INSERT INTO routing_results
            (notification_id, email_identifier, attachment_filename, document_type, action,
             destination_path, confidence_scores, error_message)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        """,
        (
            notification_id,
            email_identifier,
            filename,
            doc_type,
            action,
            destination_path,
            json.dumps(confidence_scores) if confidence_scores else None,
            error_message,
        ),
    )
    await conn.commit()


async def insert_daily_brief(
    conn: aiosqlite.Connection,
    email_identifier: str,
    notification_id: int,
    sender_email: str | None,
    subject: str | None,
    filename: str | None,
    reason: str,
    confidence_scores: dict | None = None,
    error_message: str | None = None,
) -> None:
    await conn.execute(
        """
        INSERT INTO daily_brief
            (email_identifier, notification_id, sender_email, subject,
             attachment_filename, reason, confidence_scores, error_message)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        """,
        (
            email_identifier,
            notification_id,
            sender_email,
            subject,
            filename,
            reason,
            json.dumps(confidence_scores) if confidence_scores else None,
            error_message,
        ),
    )
    await conn.commit()


async def get_pending_brief_items(conn: aiosqlite.Connection) -> list[dict]:
    cursor = await conn.execute(
        "SELECT * FROM daily_brief WHERE included_in_report = 0 ORDER BY created_at"
    )
    rows = await cursor.fetchall()
    return [dict(r) for r in rows]


async def mark_brief_items_sent(conn: aiosqlite.Connection, ids: list[int]) -> None:
    await conn.execute(
        f"UPDATE daily_brief SET included_in_report = 1, report_sent_at = datetime('now') "
        f"WHERE id IN ({','.join('?' * len(ids))})",
        ids,
    )
    await conn.commit()


async def run_retention_cleanup(conn: aiosqlite.Connection) -> None:
    await conn.executescript("""
        DELETE FROM notifications   WHERE received_at  < datetime('now', '-12 months');
        DELETE FROM routing_results WHERE processed_at < datetime('now', '-12 months');
        DELETE FROM daily_brief     WHERE created_at   < datetime('now', '-12 months');
        DELETE FROM subscriptions   WHERE is_active = 0 AND created_at < datetime('now', '-12 months');
    """)
    await conn.commit()


if __name__ == "__main__":
    import asyncio

    if "--init" in sys.argv:
        from app.config import settings

        asyncio.run(init_db(settings.database_path))
        print(f"Database initialized at {settings.database_path}")