app/storage/database.py

"""
SQLite operations for news indexing.
Uses aiosqlite for async operations compatible with FastAPI.
"""
import aiosqlite
import logging
from pathlib import Path
from datetime import datetime, timedelta
from typing import List, Optional

logger = logging.getLogger(__name__)

DB_PATH = Path(__file__).parent.parent.parent / "data" / "news_index.db"

DEFAULT_WATCHLIST = ["AAPL", "MSFT", "GOOGL", "NVDA", "AMZN"]


async def init_db():
    """Initialize database schema."""
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)

    async with aiosqlite.connect(DB_PATH) as db:
        await db.executescript("""
            CREATE TABLE IF NOT EXISTS articles (
                id TEXT PRIMARY KEY,
                symbol TEXT NOT NULL,
                headline TEXT NOT NULL,
                summary TEXT,
                created_at DATETIME NOT NULL,
                url TEXT,
                source TEXT,
                indexed_at DATETIME DEFAULT CURRENT_TIMESTAMP
            );

            CREATE TABLE IF NOT EXISTS index_metadata (
                symbol TEXT PRIMARY KEY,
                last_indexed DATETIME,
                article_count INTEGER DEFAULT 0,
                status TEXT DEFAULT 'pending'
            );

            CREATE TABLE IF NOT EXISTS watchlist (
                symbol TEXT PRIMARY KEY,
                added_at DATETIME DEFAULT CURRENT_TIMESTAMP
            );

            -- Chat history (for future persistent conversation sessions)
            CREATE TABLE IF NOT EXISTS chat_sessions (
                id         TEXT PRIMARY KEY,
                symbol     TEXT NOT NULL,
                title      TEXT,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
            );

            CREATE TABLE IF NOT EXISTS chat_messages (
                id         TEXT PRIMARY KEY,
                session_id TEXT NOT NULL REFERENCES chat_sessions(id) ON DELETE CASCADE,
                role       TEXT NOT NULL CHECK(role IN ('user', 'assistant', 'system')),
                content    TEXT NOT NULL,
                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
            );

            -- Key-value store for user preferences
            CREATE TABLE IF NOT EXISTS user_preferences (
                key        TEXT PRIMARY KEY,
                value      TEXT NOT NULL,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
            );

            -- Price alerts (future notification system)
            CREATE TABLE IF NOT EXISTS price_alerts (
                id           TEXT PRIMARY KEY,
                symbol       TEXT NOT NULL,
                condition    TEXT NOT NULL CHECK(condition IN ('above', 'below', 'change_pct')),
                threshold    REAL NOT NULL,
                active       INTEGER NOT NULL DEFAULT 1,
                created_at   DATETIME DEFAULT CURRENT_TIMESTAMP,
                triggered_at DATETIME
            );

            CREATE INDEX IF NOT EXISTS idx_articles_symbol ON articles(symbol);
            CREATE INDEX IF NOT EXISTS idx_articles_created_at ON articles(created_at);
            CREATE INDEX IF NOT EXISTS idx_articles_symbol_created ON articles(symbol, created_at);
            CREATE INDEX IF NOT EXISTS idx_chat_sessions_symbol ON chat_sessions(symbol);
            CREATE INDEX IF NOT EXISTS idx_chat_messages_session ON chat_messages(session_id);
            CREATE INDEX IF NOT EXISTS idx_alerts_symbol_active ON price_alerts(symbol, active);
        """)

        # Enable WAL mode for better concurrency
        await db.execute("PRAGMA journal_mode=WAL")
        await db.commit()

    logger.info(f"Database initialized at {DB_PATH}")


async def get_watchlist() -> List[str]:
    """Get current watchlist symbols."""
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute("SELECT symbol FROM watchlist ORDER BY added_at")
        rows = await cursor.fetchall()
        return [row[0] for row in rows]


async def save_watchlist(symbols: List[str]):
    """Replace watchlist with new symbols."""
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("DELETE FROM watchlist")
        for symbol in symbols:
            await db.execute(
                "INSERT OR IGNORE INTO watchlist (symbol) VALUES (?)",
                (symbol.upper(),)
            )
        await db.commit()


async def add_to_watchlist(symbol: str) -> bool:
    """Add a symbol to watchlist. Returns True if added (not already present)."""
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute(
            "INSERT OR IGNORE INTO watchlist (symbol) VALUES (?)",
            (symbol.upper(),)
        )
        await db.commit()
        return cursor.rowcount > 0


async def remove_from_watchlist(symbol: str):
    """Remove a symbol from watchlist."""
    async with aiosqlite.connect(DB_PATH) as db:
        await db.execute("DELETE FROM watchlist WHERE symbol = ?", (symbol.upper(),))
        await db.commit()


async def upsert_article(
    article_id: str,
    symbol: str,
    headline: str,
    summary: Optional[str],
    created_at: datetime,
    url: Optional[str],
    source: Optional[str]
) -> bool:
    """Insert article if not exists. Returns True if inserted."""
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute(
            """INSERT OR IGNORE INTO articles
               (id, symbol, headline, summary, created_at, url, source, indexed_at)
               VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
            (article_id, symbol.upper(), headline, summary, created_at, url, source, datetime.utcnow())
        )
        await db.commit()
        return cursor.rowcount > 0


async def get_article_count(symbol: str) -> int:
    """Get count of indexed articles for a symbol."""
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute(
            "SELECT COUNT(*) FROM articles WHERE symbol = ?",
            (symbol.upper(),)
        )
        row = await cursor.fetchone()
        return row[0] if row else 0


async def get_articles_by_symbol(symbol: str, limit: int = 50) -> List[dict]:
    """Get articles for a symbol, ordered by most recent."""
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        cursor = await db.execute(
            """SELECT id, symbol, headline, summary, created_at, url, source
               FROM articles
               WHERE symbol = ?
               ORDER BY created_at DESC
               LIMIT ?""",
            (symbol.upper(), limit)
        )
        rows = await cursor.fetchall()
        return [dict(row) for row in rows]


async def get_old_article_ids(days: int = 30) -> List[str]:
    """Get IDs of articles older than retention period."""
    cutoff = datetime.utcnow() - timedelta(days=days)
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute(
            "SELECT id FROM articles WHERE created_at < ?",
            (cutoff,)
        )
        rows = await cursor.fetchall()
        return [row[0] for row in rows]


async def cleanup_old_articles(days: int = 30) -> int:
    """Delete articles older than retention period. Returns count deleted."""
    cutoff = datetime.utcnow() - timedelta(days=days)
    async with aiosqlite.connect(DB_PATH) as db:
        cursor = await db.execute(
            "DELETE FROM articles WHERE created_at < ?",
            (cutoff,)
        )
        await db.commit()
        return cursor.rowcount


async def get_index_metadata(symbol: str) -> Optional[dict]:
    """Get indexing metadata for a symbol."""
    async with aiosqlite.connect(DB_PATH) as db:
        db.row_factory = aiosqlite.Row
        cursor = await db.execute(
            "SELECT * FROM index_metadata WHERE symbol = ?",
            (symbol.upper(),)
        )
        row = await cursor.fetchone()
        return dict(row) if row else None


async def update_index_metadata(
    symbol: str,
    article_count: Optional[int] = None,
    status: Optional[str] = None
):
    """Update indexing metadata for a symbol (atomic upsert)."""
    async with aiosqlite.connect(DB_PATH) as db:
        now = datetime.utcnow()

        # Use INSERT ... ON CONFLICT for atomic upsert (avoids race conditions)
        await db.execute(
            """INSERT INTO index_metadata (symbol, last_indexed, article_count, status)
               VALUES (?, ?, ?, ?)
               ON CONFLICT(symbol) DO UPDATE SET
                   last_indexed = excluded.last_indexed,
                   article_count = COALESCE(excluded.article_count, article_count),
                   status = COALESCE(excluded.status, status)""",
            (symbol.upper(), now, article_count, status)
        )
        await db.commit()


async def is_index_stale(symbol: str, max_age_minutes: int = 30) -> bool:
    """Check if index for a symbol needs refresh."""
    metadata = await get_index_metadata(symbol)
    if not metadata or not metadata.get("last_indexed"):
        return True

    last_indexed = metadata["last_indexed"]
    if isinstance(last_indexed, str):
        last_indexed = datetime.fromisoformat(last_indexed)

    age = datetime.utcnow() - last_indexed
    return age > timedelta(minutes=max_age_minutes)