#!/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,
                    alt_text TEXT,
                    link_text TEXT,
                    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)
                )
            """)

            # 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.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 = "",
        alt_text: str = "",
        link_text: str = "",
    ) -> None:
        """Create a new media record and add source context."""
        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, alt_text, link_text, crawl_job_id, discovered_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (md5_hash, media_uri, page_uri, page_title, page_description, page_keywords, alt_text, link_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 = "",
        alt_text: str = "",
        link_text: str = "",
        crawl_job_id: int = None
    ) -> None:
        """Add another source context for an existing media hash."""
        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, alt_text, link_text, crawl_job_id, discovered_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (md5_hash, media_uri, page_uri, page_title, page_description, page_keywords, alt_text, link_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 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