pig.py/database.py

#!/usr/bin/env python3
"""
SQLite database for neopig metadata.

Stores:
- Crawl jobs (target, keywords, timestamps, stats)
- Media records (md5_hash, source URLs, metadata)
- Analysis results (Qwen 3 VL outputs)
"""

import json
import logging
from datetime import datetime, timezone
from typing import List, Dict, Any, Optional, Set

import aiosqlite

logger = logging.getLogger(__name__)


class Database:
    """
    Async SQLite database for neopig metadata.
    """

    def __init__(self, db_path: str = "neopig.db"):
        self.db_path = db_path
        self._initialized = False

    async def init(self) -> None:
        """Initialize database schema."""
        if self._initialized:
            return

        async with aiosqlite.connect(self.db_path) as db:
            # Crawl jobs table
            await db.execute("""
                CREATE TABLE IF NOT EXISTS crawl_jobs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    target_uri TEXT NOT NULL,
                    keywords TEXT,  -- JSON array
                    mode TEXT DEFAULT 'images',
                    status TEXT DEFAULT 'running',
                    started_at TEXT NOT NULL,
                    completed_at TEXT,
                    stats TEXT  -- JSON object
                )
            """)

            # Media records table (main deduped storage)
            await db.execute("""
                CREATE TABLE IF NOT EXISTS media (
                    md5_hash TEXT PRIMARY KEY,
                    media_type TEXT,  -- 'image', 'video', 'audio'
                    mime_type TEXT,
                    file_size INTEGER,
                    keywords TEXT,  -- JSON array
                    alt_text TEXT,
                    title TEXT,
                    first_seen_at TEXT NOT NULL,
                    analysis_status TEXT DEFAULT 'pending',  -- 'pending', 'analyzed', 'invalid', 'error'
                    analysis_result TEXT  -- JSON from Qwen 3 VL
                )
            """)

            # Media sources table (tracks all contexts where media was found)
            await db.execute("""
                CREATE TABLE IF NOT EXISTS media_sources (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    md5_hash TEXT NOT NULL,
                    media_uri TEXT NOT NULL,
                    page_uri TEXT,
                    page_title TEXT,
                    page_description TEXT,
                    page_keywords TEXT,
                    page_content TEXT,
                    alt_text TEXT,
                    link_text TEXT,
                    detail_page_uri TEXT,
                    detail_title TEXT,
                    detail_content TEXT,
                    searchable_text TEXT,  -- Combined metadata for full-text search
                    crawl_job_id INTEGER,
                    discovered_at TEXT NOT NULL,
                    FOREIGN KEY (md5_hash) REFERENCES media(md5_hash),
                    FOREIGN KEY (crawl_job_id) REFERENCES crawl_jobs(id),
                    UNIQUE(md5_hash, media_uri, page_uri)
                )
            """)

            # Pages table for full-text search
            await db.execute("""
                CREATE TABLE IF NOT EXISTS pages (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    uri TEXT NOT NULL UNIQUE,
                    path TEXT,
                    title TEXT,
                    content TEXT,
                    markdown TEXT,
                    raw_html TEXT,
                    crawl_job_id INTEGER,
                    crawled_at TEXT NOT NULL,
                    FOREIGN KEY (crawl_job_id) REFERENCES crawl_jobs(id)
                )
            """)

            # Add columns if they don't exist (migration for existing DBs)
            for col in ['raw_html', 'markdown']:
                try:
                    await db.execute(f"ALTER TABLE pages ADD COLUMN {col} TEXT")
                except Exception:
                    pass  # Column already exists

            # FTS5 virtual table for page search
            await db.execute("""
                CREATE VIRTUAL TABLE IF NOT EXISTS pages_fts USING fts5(
                    title, content, uri, path,
                    content='pages',
                    content_rowid='id'
                )
            """)

            # Triggers to keep FTS in sync
            await db.execute("""
                CREATE TRIGGER IF NOT EXISTS pages_ai AFTER INSERT ON pages BEGIN
                    INSERT INTO pages_fts(rowid, title, content, uri, path)
                    VALUES (new.id, new.title, new.content, new.uri, new.path);
                END
            """)

            await db.execute("""
                CREATE TRIGGER IF NOT EXISTS pages_ad AFTER DELETE ON pages BEGIN
                    INSERT INTO pages_fts(pages_fts, rowid, title, content, uri, path)
                    VALUES ('delete', old.id, old.title, old.content, old.uri, old.path);
                END
            """)

            await db.execute("""
                CREATE TRIGGER IF NOT EXISTS pages_au AFTER UPDATE ON pages BEGIN
                    INSERT INTO pages_fts(pages_fts, rowid, title, content, uri, path)
                    VALUES ('delete', old.id, old.title, old.content, old.uri, old.path);
                    INSERT INTO pages_fts(rowid, title, content, uri, path)
                    VALUES (new.id, new.title, new.content, new.uri, new.path);
                END
            """)

            # Indexes
            await db.execute("CREATE INDEX IF NOT EXISTS idx_media_type ON media(media_type)")
            await db.execute("CREATE INDEX IF NOT EXISTS idx_media_analysis ON media(analysis_status)")
            await db.execute("CREATE INDEX IF NOT EXISTS idx_sources_hash ON media_sources(md5_hash)")
            await db.execute("CREATE INDEX IF NOT EXISTS idx_sources_job ON media_sources(crawl_job_id)")
            await db.execute("CREATE INDEX IF NOT EXISTS idx_sources_media_uri ON media_sources(media_uri)")
            await db.execute("CREATE INDEX IF NOT EXISTS idx_sources_page_uri ON media_sources(page_uri)")
            await db.execute("CREATE INDEX IF NOT EXISTS idx_pages_uri ON pages(uri)")

            await db.commit()

        self._initialized = True
        logger.info(f"Database initialized: {self.db_path}")

    async def create_crawl_job(
        self,
        target_uri: str,
        keywords: List[str] = None,
        mode: str = "images"
    ) -> int:
        """Create a new crawl job and return its ID."""
        async with aiosqlite.connect(self.db_path) as db:
            cursor = await db.execute(
                """
                INSERT INTO crawl_jobs (target_uri, keywords, mode, started_at)
                VALUES (?, ?, ?, ?)
                """,
                (
                    target_uri,
                    json.dumps(keywords or []),
                    mode,
                    datetime.now(timezone.utc).isoformat()
                )
            )
            await db.commit()
            return cursor.lastrowid

    async def complete_crawl_job(self, job_id: int, stats: Dict[str, Any]) -> None:
        """Mark a crawl job as complete."""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute(
                """
                UPDATE crawl_jobs
                SET status = 'completed', completed_at = ?, stats = ?
                WHERE id = ?
                """,
                (
                    datetime.now(timezone.utc).isoformat(),
                    json.dumps(stats),
                    job_id
                )
            )
            await db.commit()

    async def create_media_record(
        self,
        md5_hash: str,
        media_uri: str,
        page_uri: str,
        crawl_job_id: int,
        media_type: str = "image",
        mime_type: str = "",
        file_size: int = 0,
        page_title: str = "",
        page_description: str = "",
        page_keywords: str = "",
        page_content: str = "",
        alt_text: str = "",
        link_text: str = "",
        detail_page_uri: str = "",
        detail_title: str = "",
        detail_content: str = "",
        searchable_text: str = "",
    ) -> None:
        """Create a new media record and add source context.

        Skeleton key approach: stores both embedding context (page_title, page_content from listing)
        and detail context (detail_title, detail_content from detail page) for maximum searchability.

        For blogs: page_content contains the post text so images are searchable by post content.
        For galleries: detail_content contains the detail page text for richer metadata.
        """
        now = datetime.now(timezone.utc).isoformat()

        async with aiosqlite.connect(self.db_path) as db:
            # Insert or ignore media record (just the hash and basic info)
            await db.execute(
                """
                INSERT OR IGNORE INTO media
                (md5_hash, media_type, mime_type, file_size, first_seen_at)
                VALUES (?, ?, ?, ?, ?)
                """,
                (md5_hash, media_type, mime_type, file_size, now)
            )

            # Always add source with full context (unique on media_uri + page_uri)
            await db.execute(
                """
                INSERT OR IGNORE INTO media_sources
                (md5_hash, media_uri, page_uri, page_title, page_description, page_keywords, page_content, alt_text, link_text, detail_page_uri, detail_title, detail_content, searchable_text, crawl_job_id, discovered_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (md5_hash, media_uri, page_uri, page_title, page_description, page_keywords, page_content, alt_text, link_text, detail_page_uri, detail_title, detail_content, searchable_text, crawl_job_id, now)
            )

            await db.commit()

    async def add_media_source(
        self,
        md5_hash: str,
        media_uri: str,
        page_uri: str,
        page_title: str = "",
        page_description: str = "",
        page_keywords: str = "",
        page_content: str = "",
        alt_text: str = "",
        link_text: str = "",
        detail_page_uri: str = "",
        detail_title: str = "",
        detail_content: str = "",
        searchable_text: str = "",
        crawl_job_id: int = None
    ) -> None:
        """Add another source context for an existing media hash.

        Skeleton key approach: stores both embedding context (page_title, page_content from listing)
        and detail context (detail_title, detail_content from detail page) for maximum searchability.

        For blogs: page_content contains the post text so images are searchable by post content.
        For galleries: detail_content contains the detail page text for richer metadata.
        """
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute(
                """
                INSERT OR IGNORE INTO media_sources
                (md5_hash, media_uri, page_uri, page_title, page_description, page_keywords, page_content, alt_text, link_text, detail_page_uri, detail_title, detail_content, searchable_text, crawl_job_id, discovered_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (md5_hash, media_uri, page_uri, page_title, page_description, page_keywords, page_content, alt_text, link_text, detail_page_uri, detail_title, detail_content, searchable_text, crawl_job_id, datetime.now(timezone.utc).isoformat())
            )
            await db.commit()

    async def get_pending_analysis(self, limit: int = 100) -> List[Dict[str, Any]]:
        """Get media items pending analysis."""
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            cursor = await db.execute(
                """
                SELECT md5_hash, media_type, mime_type, keywords, alt_text, title
                FROM media
                WHERE analysis_status = 'pending'
                LIMIT ?
                """,
                (limit,)
            )
            rows = await cursor.fetchall()
            return [dict(row) for row in rows]

    async def update_analysis(
        self,
        md5_hash: str,
        status: str,
        result: Dict[str, Any] = None
    ) -> None:
        """Update analysis status and result for a media item."""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute(
                """
                UPDATE media
                SET analysis_status = ?, analysis_result = ?
                WHERE md5_hash = ?
                """,
                (status, json.dumps(result) if result else None, md5_hash)
            )
            await db.commit()

    async def get_media_by_keyword(self, keyword: str, limit: int = 100) -> List[Dict[str, Any]]:
        """Search media by keyword."""
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            # Search in keywords JSON array and alt_text/title
            cursor = await db.execute(
                """
                SELECT m.*, GROUP_CONCAT(ms.source_url) as source_urls
                FROM media m
                LEFT JOIN media_sources ms ON m.md5_hash = ms.md5_hash
                WHERE m.keywords LIKE ?
                   OR m.alt_text LIKE ?
                   OR m.title LIKE ?
                GROUP BY m.md5_hash
                LIMIT ?
                """,
                (f'%{keyword}%', f'%{keyword}%', f'%{keyword}%', limit)
            )
            rows = await cursor.fetchall()
            return [dict(row) for row in rows]

    async def check_media_uri_exists(self, media_uri: str, page_uri: str) -> Optional[str]:
        """
        Check if a media URI from a specific page has already been crawled.

        Returns:
            The md5_hash if exists, None otherwise
        """
        async with aiosqlite.connect(self.db_path) as db:
            cursor = await db.execute(
                "SELECT md5_hash FROM media_sources WHERE media_uri = ? AND page_uri = ?",
                (media_uri, page_uri)
            )
            row = await cursor.fetchone()
            return row[0] if row else None

    async def get_crawled_media_uris(self) -> Set[str]:
        """Get all media URIs that have been crawled."""
        async with aiosqlite.connect(self.db_path) as db:
            cursor = await db.execute(
                "SELECT DISTINCT media_uri FROM media_sources"
            )
            rows = await cursor.fetchall()
            return {row[0] for row in rows}

    async def store_page(
        self,
        uri: str,
        title: str,
        content: str,
        path: str = "",
        markdown: str = "",
        raw_html: str = "",
        crawl_job_id: int = None
    ) -> None:
        """Store a page for full-text search and phantom site recreation."""
        async with aiosqlite.connect(self.db_path) as db:
            await db.execute(
                """
                INSERT OR REPLACE INTO pages (uri, path, title, content, markdown, raw_html, crawl_job_id, crawled_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (uri, path, title, content[:100000], markdown, raw_html, crawl_job_id, datetime.now(timezone.utc).isoformat())
            )
            await db.commit()

    async def search_pages(self, query: str, limit: int = 50) -> List[Dict[str, Any]]:
        """Search pages using FTS5 with LIKE fallback."""
        async with aiosqlite.connect(self.db_path) as db:
            db.row_factory = aiosqlite.Row
            results = []

            # Try FTS5 with prefix matching
            try:
                fts_query = ' '.join(f'"{word}"*' for word in query.split())
                cursor = await db.execute("""
                    SELECT p.uri, p.path, p.title,
                           snippet(pages_fts, 1, '<mark>', '</mark>', '...', 40) as snippet
                    FROM pages_fts
                    JOIN pages p ON pages_fts.rowid = p.id
                    WHERE pages_fts MATCH ?
                    ORDER BY rank
                    LIMIT ?
                """, (fts_query, limit))
                results = [dict(row) for row in await cursor.fetchall()]
            except Exception:
                pass

            # Fallback to LIKE
            if not results:
                try:
                    like_q = f'%{query}%'
                    cursor = await db.execute("""
                        SELECT uri, path, title, substr(content, 1, 200) as snippet
                        FROM pages
                        WHERE title LIKE ? COLLATE NOCASE
                           OR content LIKE ? COLLATE NOCASE
                        LIMIT ?
                    """, (like_q, like_q, limit))
                    results = [dict(row) for row in await cursor.fetchall()]
                except Exception:
                    pass

            return results

    async def get_stats(self) -> Dict[str, Any]:
        """Get database statistics."""
        async with aiosqlite.connect(self.db_path) as db:
            stats = {}

            # Total media
            cursor = await db.execute("SELECT COUNT(*) FROM media")
            stats['total_media'] = (await cursor.fetchone())[0]

            # By type
            cursor = await db.execute(
                "SELECT media_type, COUNT(*) FROM media GROUP BY media_type"
            )
            stats['by_type'] = {row[0]: row[1] for row in await cursor.fetchall()}

            # By analysis status
            cursor = await db.execute(
                "SELECT analysis_status, COUNT(*) FROM media GROUP BY analysis_status"
            )
            stats['by_analysis'] = {row[0]: row[1] for row in await cursor.fetchall()}

            # Total sources
            cursor = await db.execute("SELECT COUNT(*) FROM media_sources")
            stats['total_sources'] = (await cursor.fetchone())[0]

            # Crawl jobs
            cursor = await db.execute("SELECT COUNT(*) FROM crawl_jobs")
            stats['total_jobs'] = (await cursor.fetchone())[0]

            return stats