#!/usr/bin/env python3
"""
PostgreSQL database module for production microservice.
Stores job metadata and reviews as JSONB.
"""
import asyncpg
import json
from datetime import datetime
from typing import Optional, List, Dict, Any
from uuid import UUID, uuid4
import logging

from core.enums import JobStatus

log = logging.getLogger(__name__)


class DatabaseManager:
    """PostgreSQL database manager with connection pooling"""

    def __init__(self, database_url: str):
        """
        Initialize database manager.

        Args:
            database_url: PostgreSQL connection URL
                         Format: postgresql://user:password@host:port/database
        """
        self.database_url = database_url
        self.pool: Optional[asyncpg.Pool] = None

    async def connect(self):
        """Create connection pool"""
        log.info("Connecting to PostgreSQL database...")
        self.pool = await asyncpg.create_pool(
            self.database_url,
            min_size=5,
            max_size=20,
            command_timeout=60
        )
        log.info("Database connection pool created")

    async def disconnect(self):
        """Close connection pool"""
        if self.pool:
            await self.pool.close()
            log.info("Database connection pool closed")

    async def initialize_schema(self):
        """Create database schema if it doesn't exist"""
        async with self.pool.acquire() as conn:
            # Create jobs table
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS jobs (
                    job_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    status VARCHAR(20) NOT NULL DEFAULT 'pending',
                    url TEXT NOT NULL,
                    webhook_url TEXT,
                    webhook_secret TEXT,

                    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
                    started_at TIMESTAMP,
                    completed_at TIMESTAMP,
                    updated_at TIMESTAMP,

                    reviews_count INTEGER,
                    total_reviews INTEGER,
                    reviews_data JSONB,
                    scrape_time REAL,

                    error_message TEXT,
                    metadata JSONB,
                    scrape_logs JSONB,

                    CONSTRAINT valid_status CHECK (status IN ('pending', 'running', 'completed', 'failed', 'cancelled', 'partial'))
                );
            """)

            # Add scrape_logs column if it doesn't exist (for existing databases)
            await conn.execute("""
                ALTER TABLE jobs ADD COLUMN IF NOT EXISTS scrape_logs JSONB;
            """)

            # Add updated_at column if it doesn't exist (for incremental progress tracking)
            await conn.execute("""
                ALTER TABLE jobs ADD COLUMN IF NOT EXISTS updated_at TIMESTAMP;
            """)

            # Add review_topics column if it doesn't exist (extracted topic filters with mention counts)
            await conn.execute("""
                ALTER TABLE jobs ADD COLUMN IF NOT EXISTS review_topics JSONB;
            """)

            # Update constraint to include 'partial' status (for existing databases)
            await conn.execute("""
                ALTER TABLE jobs DROP CONSTRAINT IF EXISTS valid_status;
            """)
            await conn.execute("""
                ALTER TABLE jobs ADD CONSTRAINT valid_status CHECK (status IN ('pending', 'running', 'completed', 'failed', 'cancelled', 'partial'));
            """)

            # Create indexes
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_jobs_status ON jobs(status);
            """)
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_jobs_created_at ON jobs(created_at DESC);
            """)
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_jobs_webhook ON jobs(webhook_url) WHERE webhook_url IS NOT NULL;
            """)

            # Create canary results table
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS canary_results (
                    id SERIAL PRIMARY KEY,
                    timestamp TIMESTAMP NOT NULL DEFAULT NOW(),
                    success BOOLEAN NOT NULL,
                    reviews_count INTEGER,
                    scrape_time REAL,
                    error_message TEXT,
                    metadata JSONB
                );
            """)

            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_canary_timestamp ON canary_results(timestamp DESC);
            """)

            # Create webhook attempts table (for retry tracking)
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS webhook_attempts (
                    id SERIAL PRIMARY KEY,
                    job_id UUID NOT NULL REFERENCES jobs(job_id) ON DELETE CASCADE,
                    attempt_number INTEGER NOT NULL,
                    timestamp TIMESTAMP NOT NULL DEFAULT NOW(),
                    success BOOLEAN NOT NULL,
                    status_code INTEGER,
                    error_message TEXT,
                    response_time_ms REAL
                );
            """)

            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_webhook_job_id ON webhook_attempts(job_id);
            """)

            # Add session_fingerprint and metrics_history columns to jobs table
            await conn.execute("""
                ALTER TABLE jobs ADD COLUMN IF NOT EXISTS session_fingerprint JSONB;
            """)
            await conn.execute("""
                ALTER TABLE jobs ADD COLUMN IF NOT EXISTS metrics_history JSONB;
            """)

            # Create crash_reports table
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS crash_reports (
                    crash_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    job_id UUID REFERENCES jobs(job_id) ON DELETE CASCADE,
                    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
                    crash_type VARCHAR(50) NOT NULL,
                    error_message TEXT,
                    state JSONB NOT NULL,
                    metrics_history JSONB,
                    logs_before_crash JSONB,
                    analysis JSONB,
                    screenshot_url TEXT,
                    dom_snapshot_id UUID
                );
            """)

            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_crash_reports_job ON crash_reports(job_id);
            """)
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_crash_reports_type ON crash_reports(crash_type);
            """)
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_crash_reports_created ON crash_reports(created_at DESC);
            """)

            log.info("Database schema initialized")

    async def run_migrations(self, migrations_dir: str = "migrations/versions"):
        """
        Run versioned migrations from SQL files.

        Args:
            migrations_dir: Path to directory containing .sql migration files.
                           Files are run in sorted order.

        Returns:
            Number of migrations applied.
        """
        from pathlib import Path

        migrations_path = Path(migrations_dir)
        if not migrations_path.exists():
            log.warning(f"Migrations directory not found: {migrations_dir}")
            return 0

        async with self.pool.acquire() as conn:
            # Create migrations tracking table
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS _migrations (
                    id SERIAL PRIMARY KEY,
                    filename VARCHAR(255) UNIQUE NOT NULL,
                    applied_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                )
            """)

            # Get already applied migrations
            applied = await conn.fetch("SELECT filename FROM _migrations")
            applied_set = {r["filename"] for r in applied}

            # Find and run pending migrations
            migration_files = sorted(migrations_path.glob("*.sql"))
            migrations_run = 0

            for migration_file in migration_files:
                filename = migration_file.name
                if filename in applied_set:
                    continue

                log.info(f"Running migration: {filename}")

                async with conn.transaction():
                    try:
                        sql = migration_file.read_text()
                        await conn.execute(sql)
                        await conn.execute(
                            "INSERT INTO _migrations (filename) VALUES ($1)",
                            filename,
                        )
                        migrations_run += 1
                        log.info(f"Migration {filename} applied successfully")
                    except Exception as e:
                        log.error(f"Migration {filename} failed: {e}")
                        raise

            log.info(f"Ran {migrations_run} migrations")
            return migrations_run

    # ==================== Job Operations ====================

    async def create_job(
        self,
        url: str,
        webhook_url: Optional[str] = None,
        webhook_secret: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
        # Phase 2: New fields for requester tracking
        requester_client_id: Optional[str] = None,
        requester_source: Optional[str] = None,
        scrape_purpose: Optional[str] = None,
        requester_metadata: Optional[Dict[str, Any]] = None,
        # Phase 2: Batch support
        batch_id: Optional[UUID] = None,
        batch_index: Optional[int] = None,
        # Phase 2: Job configuration
        job_type: str = 'google_reviews',
        priority: int = 0,
        callback_url: Optional[str] = None,
        # Phase 2: Scraper versioning
        scraper_version: Optional[str] = None,
        scraper_variant: Optional[str] = None
    ) -> UUID:
        """
        Create a new scraping job.

        Args:
            url: Google Maps URL to scrape
            webhook_url: Optional webhook URL for notifications
            webhook_secret: Optional secret for webhook signature
            metadata: Optional additional metadata
            requester_client_id: Client ID of the requester (for tracking)
            requester_source: Source of the request (e.g., 'api', 'web', 'batch')
            scrape_purpose: Purpose of the scrape (e.g., 'competitor_analysis')
            requester_metadata: Additional requester-specific metadata
            batch_id: ID of the batch this job belongs to
            batch_index: Index of this job within the batch
            job_type: Type of job (default: 'google_reviews')
            priority: Job priority (higher = more urgent, default: 0)
            callback_url: URL to call when job completes
            scraper_version: Version of the scraper to use
            scraper_variant: Variant of the scraper (e.g., 'stealth', 'fast')

        Returns:
            UUID of created job
        """
        async with self.pool.acquire() as conn:
            job_id = await conn.fetchval("""
                INSERT INTO jobs (
                    url, webhook_url, webhook_secret, metadata,
                    requester_client_id, requester_source, scrape_purpose, requester_metadata,
                    batch_id, batch_index,
                    job_type, priority, callback_url,
                    scraper_version, scraper_variant
                )
                VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15)
                RETURNING job_id
            """,
                url, webhook_url, webhook_secret,
                json.dumps(metadata) if metadata else None,
                requester_client_id, requester_source, scrape_purpose,
                json.dumps(requester_metadata) if requester_metadata else None,
                batch_id, batch_index,
                job_type, priority, callback_url,
                scraper_version, scraper_variant
            )

            log.info(f"Created job {job_id} for URL: {url[:80]}...")
            return job_id

    async def get_job(self, job_id: UUID) -> Optional[Dict[str, Any]]:
        """
        Get job by ID.

        Args:
            job_id: Job UUID

        Returns:
            Job dictionary or None if not found
        """
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow("""
                SELECT
                    job_id,
                    status,
                    url,
                    webhook_url,
                    created_at,
                    started_at,
                    completed_at,
                    updated_at,
                    reviews_count,
                    total_reviews,
                    reviews_data,
                    scrape_time,
                    error_message,
                    metadata,
                    scrape_logs,
                    review_topics,
                    requester_client_id,
                    requester_source,
                    scrape_purpose,
                    requester_metadata,
                    batch_id,
                    batch_index,
                    job_type,
                    priority,
                    callback_url,
                    callback_status,
                    callback_attempts,
                    scraper_version,
                    scraper_variant,
                    business_name,
                    business_category,
                    business_address,
                    business_rating
                FROM jobs
                WHERE job_id = $1
            """, job_id)

            if not row:
                return None

            return dict(row)

    async def get_job_reviews(self, job_id: UUID, include_partial: bool = True) -> Optional[List[Dict[str, Any]]]:
        """
        Get reviews for a specific job.

        Args:
            job_id: Job UUID
            include_partial: If True, also return reviews for running and partial jobs

        Returns:
            List of reviews or None if not found/no reviews
        """
        async with self.pool.acquire() as conn:
            if include_partial:
                # Return reviews for completed, running, or partial jobs
                reviews_data = await conn.fetchval("""
                    SELECT reviews_data
                    FROM jobs
                    WHERE job_id = $1 AND status IN ('completed', 'running', 'partial')
                """, job_id)
            else:
                # Only return reviews for completed jobs
                reviews_data = await conn.fetchval("""
                    SELECT reviews_data
                    FROM jobs
                    WHERE job_id = $1 AND status = 'completed'
                """, job_id)

            if not reviews_data:
                return None

            # asyncpg returns JSONB as string, need to parse it
            if isinstance(reviews_data, str):
                return json.loads(reviews_data)

            return reviews_data

    async def update_job_status(
        self,
        job_id: UUID,
        status: JobStatus,
        **kwargs
    ):
        """
        Update job status and optional fields.

        Args:
            job_id: Job UUID
            status: New status
            **kwargs: Additional fields to update (started_at, completed_at, error_message, etc.)
        """
        # Build dynamic UPDATE query
        set_clauses = ["status = $2"]
        params = [job_id, status.value]
        param_idx = 3

        if status == JobStatus.RUNNING and 'started_at' not in kwargs:
            kwargs['started_at'] = datetime.now()
        elif status in [JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED] and 'completed_at' not in kwargs:
            kwargs['completed_at'] = datetime.now()

        for key, value in kwargs.items():
            # Handle JSONB fields specially
            if key == 'scrape_logs' and value is not None:
                set_clauses.append(f"{key} = ${param_idx}::jsonb")
                params.append(json.dumps(value) if not isinstance(value, str) else value)
            else:
                set_clauses.append(f"{key} = ${param_idx}")
                params.append(value)
            param_idx += 1

        query = f"""
            UPDATE jobs
            SET {', '.join(set_clauses)}
            WHERE job_id = $1
        """

        async with self.pool.acquire() as conn:
            await conn.execute(query, *params)

    async def save_job_result(
        self,
        job_id: UUID,
        reviews: List[Dict[str, Any]],
        scrape_time: float,
        total_reviews: Optional[int] = None,
        scrape_logs: Optional[List[Dict[str, Any]]] = None,
        review_topics: Optional[List[Dict[str, Any]]] = None
    ):
        """
        Save scraping results to database.

        Args:
            job_id: Job UUID
            reviews: List of review dictionaries
            scrape_time: Time taken to scrape in seconds
            total_reviews: Total reviews available (from page counter)
            scrape_logs: List of log entries from the scraper
            review_topics: List of topic filter dictionaries with topic and count
        """
        async with self.pool.acquire() as conn:
            # If reviews list is empty, check if job already has reviews from incremental saves
            # This happens when flush_callback was used during scraping
            if not reviews:
                existing = await conn.fetchval(
                    "SELECT reviews_count FROM jobs WHERE job_id = $1", job_id
                )
                if existing and existing > 0:
                    # Job has reviews from incremental saves, don't overwrite reviews_data
                    await conn.execute("""
                        UPDATE jobs
                        SET
                            status = 'completed',
                            completed_at = NOW(),
                            total_reviews = COALESCE($2, total_reviews),
                            scrape_time = $3,
                            scrape_logs = $4::jsonb,
                            review_topics = $5::jsonb
                        WHERE job_id = $1
                    """, job_id, total_reviews, scrape_time,
                        json.dumps(scrape_logs) if scrape_logs else None,
                        json.dumps(review_topics) if review_topics else None)
                    log.info(f"Completed job {job_id} with {existing} reviews (from incremental saves)")
                    return

            await conn.execute("""
                UPDATE jobs
                SET
                    status = 'completed',
                    completed_at = NOW(),
                    reviews_count = $2,
                    total_reviews = $3,
                    reviews_data = $4::jsonb,
                    scrape_time = $5,
                    scrape_logs = $6::jsonb,
                    review_topics = $7::jsonb
                WHERE job_id = $1
            """, job_id, len(reviews), total_reviews, json.dumps(reviews), scrape_time,
                json.dumps(scrape_logs) if scrape_logs else None,
                json.dumps(review_topics) if review_topics else None)

            log.info(f"Saved {len(reviews)} reviews for job {job_id}")

    async def save_reviews_incremental(
        self,
        job_id: UUID,
        reviews: List[Dict[str, Any]],
        total_reviews: Optional[int] = None
    ):
        """
        Save reviews incrementally during scraping.
        Called on each flush to preserve progress in case of crash.

        Args:
            job_id: Job UUID
            reviews: ALL reviews collected so far (not just new ones)
            total_reviews: Total reviews available (from page counter)
        """
        async with self.pool.acquire() as conn:
            await conn.execute("""
                UPDATE jobs
                SET
                    reviews_count = $2,
                    total_reviews = COALESCE($3, total_reviews),
                    reviews_data = $4::jsonb,
                    updated_at = NOW()
                WHERE job_id = $1 AND status = 'running'
            """, job_id, len(reviews), total_reviews, json.dumps(reviews))

            log.debug(f"Incremental save: {len(reviews)} reviews for job {job_id}")

    async def update_session_fingerprint(
        self,
        job_id: UUID,
        session_fingerprint: Dict[str, Any]
    ):
        """
        Update the session fingerprint for a job.

        This should be called early in the scraping process after the browser
        fingerprint is captured, to record browser characteristics for
        bot detection analysis.

        Args:
            job_id: Job UUID
            session_fingerprint: Dictionary containing browser fingerprint data:
                - user_agent: Browser user agent string
                - platform: OS platform
                - language: Primary language
                - languages: List of accepted languages
                - timezone: Timezone string
                - screen: {width, height, colorDepth}
                - viewport: {width, height}
                - webgl_vendor: WebGL vendor string
                - webgl_renderer: WebGL renderer string
                - canvas_fingerprint: Canvas fingerprint hash
                - hardware_concurrency: Number of CPU cores
                - device_memory: Device memory in GB
                - bot_detection_tests: {webdriver_hidden, chrome_runtime, permissions_query}
                - captured_at: ISO timestamp when fingerprint was captured
        """
        async with self.pool.acquire() as conn:
            await conn.execute("""
                UPDATE jobs
                SET
                    session_fingerprint = $2::jsonb,
                    updated_at = NOW()
                WHERE job_id = $1
            """, job_id, json.dumps(session_fingerprint))

            log.debug(f"Updated session fingerprint for job {job_id}")

    async def update_job_metadata(
        self,
        job_id: UUID,
        metadata_updates: Dict[str, Any]
    ):
        """
        Update specific fields in job metadata without overwriting existing data.

        Args:
            job_id: Job UUID
            metadata_updates: Dictionary of metadata fields to update/add
                - bot_detected: True if sort button was hidden (bot detection)
                - initial_sort_used: Sort order used for scraping
                - sort_orders_attempted: List of all sort orders tried
                - multi_sort: Multi-sort completion info
        """
        async with self.pool.acquire() as conn:
            # Merge new metadata with existing metadata using JSONB concatenation
            await conn.execute("""
                UPDATE jobs
                SET
                    metadata = COALESCE(metadata, '{}'::jsonb) || $2::jsonb,
                    updated_at = NOW()
                WHERE job_id = $1
            """, job_id, json.dumps(metadata_updates))

            log.debug(f"Updated job metadata for job {job_id}: {list(metadata_updates.keys())}")

    async def update_business_info(
        self,
        job_id: UUID,
        business_name: Optional[str] = None,
        business_category: Optional[str] = None,
        business_address: Optional[str] = None,
        business_rating: Optional[float] = None
    ):
        """
        Update business info columns for a job.

        These are dedicated columns (not JSONB) for queryable business data
        captured from the Google Maps page during scraping.

        Args:
            job_id: Job UUID
            business_name: Business name from Google Maps
            business_category: Business category (e.g., "Restaurant", "Toy store")
            business_address: Full address from Google Maps
            business_rating: Aggregate rating at time of scrape (e.g., 4.5)
        """
        async with self.pool.acquire() as conn:
            await conn.execute("""
                UPDATE jobs
                SET
                    business_name = COALESCE($2, business_name),
                    business_category = COALESCE($3, business_category),
                    business_address = COALESCE($4, business_address),
                    business_rating = COALESCE($5, business_rating),
                    updated_at = NOW()
                WHERE job_id = $1
            """, job_id, business_name, business_category, business_address, business_rating)

            log.debug(f"Updated business info for job {job_id}: name={business_name}, category={business_category}")

    async def mark_job_partial(
        self,
        job_id: UUID,
        error_message: str,
        scrape_logs: Optional[List[Dict[str, Any]]] = None
    ):
        """
        Mark a job as partial (crashed but has some reviews saved).

        Args:
            job_id: Job UUID
            error_message: Error that caused the crash
            scrape_logs: Log entries from the scraper
        """
        async with self.pool.acquire() as conn:
            await conn.execute("""
                UPDATE jobs
                SET
                    status = 'partial',
                    completed_at = NOW(),
                    error_message = $2,
                    scrape_logs = $3::jsonb
                WHERE job_id = $1
            """, job_id, error_message, json.dumps(scrape_logs) if scrape_logs else None)

            log.info(f"Marked job {job_id} as partial due to: {error_message}")

    async def list_jobs(
        self,
        status: Optional[JobStatus] = None,
        limit: int = 100,
        offset: int = 0,
        requester_client_id: Optional[str] = None,
        batch_id: Optional[UUID] = None
    ) -> List[Dict[str, Any]]:
        """
        List jobs with optional filtering.

        Args:
            status: Optional status filter
            limit: Maximum number of jobs to return
            offset: Number of jobs to skip
            requester_client_id: Optional filter by requester client ID
            batch_id: Optional filter by batch ID

        Returns:
            List of job dictionaries
        """
        async with self.pool.acquire() as conn:
            # Build dynamic WHERE clause
            conditions = []
            params = []
            param_idx = 1

            if status:
                conditions.append(f"status = ${param_idx}")
                params.append(status.value)
                param_idx += 1

            if requester_client_id:
                conditions.append(f"requester_client_id = ${param_idx}")
                params.append(requester_client_id)
                param_idx += 1

            if batch_id:
                conditions.append(f"batch_id = ${param_idx}")
                params.append(batch_id)
                param_idx += 1

            where_clause = f"WHERE {' AND '.join(conditions)}" if conditions else ""

            # Add limit and offset params
            params.extend([limit, offset])

            query = f"""
                SELECT
                    job_id,
                    status,
                    url,
                    created_at,
                    completed_at,
                    reviews_count,
                    total_reviews,
                    scrape_time,
                    error_message,
                    metadata,
                    review_topics,
                    requester_client_id,
                    requester_source,
                    scrape_purpose,
                    requester_metadata,
                    batch_id,
                    batch_index,
                    job_type,
                    priority,
                    callback_url,
                    callback_status,
                    callback_attempts,
                    scraper_version,
                    scraper_variant,
                    business_name,
                    business_category,
                    business_address,
                    business_rating
                FROM jobs
                {where_clause}
                ORDER BY created_at DESC
                LIMIT ${param_idx} OFFSET ${param_idx + 1}
            """

            rows = await conn.fetch(query, *params)
            return [dict(row) for row in rows]

    async def update_job_callback(
        self,
        job_id: UUID,
        status: str,
        attempts: Optional[int] = None
    ):
        """
        Update callback status for a job.

        Args:
            job_id: Job UUID
            status: Callback status ('pending', 'success', 'failed', 'skipped')
            attempts: Number of callback attempts (if not provided, increments by 1)
        """
        async with self.pool.acquire() as conn:
            if attempts is not None:
                await conn.execute("""
                    UPDATE jobs
                    SET callback_status = $2, callback_attempts = $3, updated_at = NOW()
                    WHERE job_id = $1
                """, job_id, status, attempts)
            else:
                await conn.execute("""
                    UPDATE jobs
                    SET callback_status = $2,
                        callback_attempts = COALESCE(callback_attempts, 0) + 1,
                        updated_at = NOW()
                    WHERE job_id = $1
                """, job_id, status)

            log.debug(f"Updated callback status for job {job_id}: {status}")

    # ==================== Batch Operations ====================

    async def create_batch(
        self,
        name: str,
        requester_client_id: Optional[str] = None,
        requester_source: Optional[str] = None,
        scrape_purpose: Optional[str] = None,
        total_jobs: int = 0,
        callback_url: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> UUID:
        """
        Create a new batch for grouping jobs.

        Args:
            name: Batch name/description
            requester_client_id: Client ID of the requester
            requester_source: Source of the batch request
            scrape_purpose: Purpose of the scrape
            total_jobs: Expected total number of jobs in batch
            callback_url: URL to call when batch completes
            metadata: Additional batch metadata

        Returns:
            UUID of created batch
        """
        async with self.pool.acquire() as conn:
            batch_id = await conn.fetchval("""
                INSERT INTO batches (
                    name, requester_client_id, requester_source, scrape_purpose,
                    total_jobs, callback_url, metadata
                )
                VALUES ($1, $2, $3, $4, $5, $6, $7)
                RETURNING batch_id
            """,
                name, requester_client_id, requester_source, scrape_purpose,
                total_jobs, callback_url,
                json.dumps(metadata) if metadata else None
            )

            log.info(f"Created batch {batch_id}: {name} ({total_jobs} jobs)")
            return batch_id

    async def get_batch(self, batch_id: UUID) -> Optional[Dict[str, Any]]:
        """
        Get batch by ID with job counts.

        Args:
            batch_id: Batch UUID

        Returns:
            Batch dictionary with job counts or None if not found
        """
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow("""
                SELECT
                    b.batch_id,
                    b.name,
                    b.status,
                    b.requester_client_id,
                    b.requester_source,
                    b.scrape_purpose,
                    b.total_jobs,
                    b.completed_jobs,
                    b.failed_jobs,
                    b.callback_url,
                    b.callback_status,
                    b.metadata,
                    b.created_at,
                    b.updated_at,
                    b.completed_at,
                    COUNT(j.job_id) FILTER (WHERE j.status = 'pending') as pending_jobs,
                    COUNT(j.job_id) FILTER (WHERE j.status = 'running') as running_jobs,
                    COUNT(j.job_id) as actual_total_jobs
                FROM batches b
                LEFT JOIN jobs j ON j.batch_id = b.batch_id
                WHERE b.batch_id = $1
                GROUP BY b.batch_id
            """, batch_id)

            if not row:
                return None

            return dict(row)

    async def update_batch_progress(self, batch_id: UUID):
        """
        Recalculate and update batch progress from jobs table.

        Args:
            batch_id: Batch UUID
        """
        async with self.pool.acquire() as conn:
            # Calculate job counts
            counts = await conn.fetchrow("""
                SELECT
                    COUNT(*) FILTER (WHERE status = 'completed') as completed,
                    COUNT(*) FILTER (WHERE status = 'failed') as failed,
                    COUNT(*) FILTER (WHERE status = 'partial') as partial,
                    COUNT(*) as total
                FROM jobs
                WHERE batch_id = $1
            """, batch_id)

            completed = counts['completed'] or 0
            failed = counts['failed'] or 0
            partial = counts['partial'] or 0
            total = counts['total'] or 0

            # Determine batch status
            if total == 0:
                status = 'pending'
            elif completed + failed + partial >= total:
                status = 'completed'
            elif completed > 0 or failed > 0 or partial > 0:
                status = 'running'
            else:
                status = 'pending'

            # Update batch
            await conn.execute("""
                UPDATE batches
                SET
                    completed_jobs = $2,
                    failed_jobs = $3,
                    status = $4,
                    updated_at = NOW(),
                    completed_at = CASE WHEN $4 = 'completed' THEN NOW() ELSE completed_at END
                WHERE batch_id = $1
            """, batch_id, completed, failed, status)

            log.debug(f"Updated batch {batch_id} progress: {completed}/{total} completed, {failed} failed")

    async def get_batches(
        self,
        requester_client_id: Optional[str] = None,
        status: Optional[str] = None,
        limit: int = 50
    ) -> List[Dict[str, Any]]:
        """
        List batches with optional filtering.

        Args:
            requester_client_id: Optional filter by requester client ID
            status: Optional filter by batch status
            limit: Maximum number of batches to return

        Returns:
            List of batch dictionaries
        """
        async with self.pool.acquire() as conn:
            # Build dynamic WHERE clause
            conditions = []
            params = []
            param_idx = 1

            if requester_client_id:
                conditions.append(f"requester_client_id = ${param_idx}")
                params.append(requester_client_id)
                param_idx += 1

            if status:
                conditions.append(f"status = ${param_idx}")
                params.append(status)
                param_idx += 1

            where_clause = f"WHERE {' AND '.join(conditions)}" if conditions else ""
            params.append(limit)

            query = f"""
                SELECT
                    batch_id,
                    name,
                    status,
                    requester_client_id,
                    requester_source,
                    scrape_purpose,
                    total_jobs,
                    completed_jobs,
                    failed_jobs,
                    callback_url,
                    callback_status,
                    metadata,
                    created_at,
                    updated_at,
                    completed_at
                FROM batches
                {where_clause}
                ORDER BY created_at DESC
                LIMIT ${param_idx}
            """

            rows = await conn.fetch(query, *params)
            return [dict(row) for row in rows]

    async def get_pending_jobs_with_webhooks(self, limit: int = 100) -> List[Dict[str, Any]]:
        """
        Get completed jobs that have webhooks pending delivery.

        Args:
            limit: Maximum number of jobs to return

        Returns:
            List of job dictionaries with webhook info
        """
        async with self.pool.acquire() as conn:
            rows = await conn.fetch("""
                SELECT
                    job_id,
                    status,
                    url,
                    webhook_url,
                    webhook_secret,
                    reviews_count,
                    scrape_time,
                    error_message,
                    completed_at
                FROM jobs
                WHERE webhook_url IS NOT NULL
                  AND status IN ('completed', 'failed')
                  AND job_id NOT IN (
                      SELECT job_id
                      FROM webhook_attempts
                      WHERE success = true
                  )
                ORDER BY completed_at ASC
                LIMIT $1
            """, limit)

            return [dict(row) for row in rows]

    async def delete_job(self, job_id: UUID) -> bool:
        """
        Delete a job from the database.

        Args:
            job_id: Job UUID

        Returns:
            True if deleted, False if not found
        """
        async with self.pool.acquire() as conn:
            result = await conn.execute("""
                DELETE FROM jobs WHERE job_id = $1
            """, job_id)

            deleted = result.split()[-1] == "1"
            if deleted:
                log.info(f"Deleted job {job_id}")
            return deleted

    async def cleanup_old_jobs(self, max_age_days: int = 30):
        """
        Delete old completed/failed jobs.

        Args:
            max_age_days: Maximum age in days before deletion
        """
        async with self.pool.acquire() as conn:
            result = await conn.execute("""
                DELETE FROM jobs
                WHERE status IN ('completed', 'failed', 'cancelled')
                  AND completed_at < NOW() - INTERVAL '%s days'
            """, max_age_days)

            deleted_count = int(result.split()[-1])
            if deleted_count > 0:
                log.info(f"Cleaned up {deleted_count} old jobs")

    # ==================== Statistics ====================

    async def get_stats(self) -> Dict[str, Any]:
        """
        Get job statistics.

        Returns:
            Statistics dictionary
        """
        async with self.pool.acquire() as conn:
            stats = await conn.fetchrow("""
                SELECT
                    COUNT(*) as total_jobs,
                    COUNT(*) FILTER (WHERE status = 'pending') as pending,
                    COUNT(*) FILTER (WHERE status = 'running') as running,
                    COUNT(*) FILTER (WHERE status = 'completed') as completed,
                    COUNT(*) FILTER (WHERE status = 'failed') as failed,
                    COUNT(*) FILTER (WHERE status = 'cancelled') as cancelled,
                    AVG(scrape_time) FILTER (WHERE status = 'completed') as avg_scrape_time,
                    SUM(reviews_count) FILTER (WHERE status = 'completed') as total_reviews
                FROM jobs
            """)

            return dict(stats)

    # ==================== Canary Operations ====================

    async def save_canary_result(
        self,
        success: bool,
        reviews_count: Optional[int] = None,
        scrape_time: Optional[float] = None,
        error_message: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None
    ):
        """
        Save canary test result.

        Args:
            success: Whether canary test succeeded
            reviews_count: Number of reviews scraped
            scrape_time: Time taken in seconds
            error_message: Error message if failed
            metadata: Additional metadata
        """
        async with self.pool.acquire() as conn:
            await conn.execute("""
                INSERT INTO canary_results (success, reviews_count, scrape_time, error_message, metadata)
                VALUES ($1, $2, $3, $4, $5)
            """, success, reviews_count, scrape_time, error_message, json.dumps(metadata) if metadata else None)

    async def get_canary_history(self, limit: int = 100) -> List[Dict[str, Any]]:
        """
        Get canary test history.

        Args:
            limit: Maximum number of results to return

        Returns:
            List of canary result dictionaries
        """
        async with self.pool.acquire() as conn:
            rows = await conn.fetch("""
                SELECT
                    timestamp,
                    success,
                    reviews_count,
                    scrape_time,
                    error_message
                FROM canary_results
                ORDER BY timestamp DESC
                LIMIT $1
            """, limit)

            return [dict(row) for row in rows]

    # ==================== Webhook Attempts ====================

    async def log_webhook_attempt(
        self,
        job_id: UUID,
        attempt_number: int,
        success: bool,
        status_code: Optional[int] = None,
        error_message: Optional[str] = None,
        response_time_ms: Optional[float] = None
    ):
        """
        Log a webhook delivery attempt.

        Args:
            job_id: Job UUID
            attempt_number: Attempt number (1, 2, 3...)
            success: Whether delivery succeeded
            status_code: HTTP status code
            error_message: Error message if failed
            response_time_ms: Response time in milliseconds
        """
        async with self.pool.acquire() as conn:
            await conn.execute("""
                INSERT INTO webhook_attempts (job_id, attempt_number, success, status_code, error_message, response_time_ms)
                VALUES ($1, $2, $3, $4, $5, $6)
            """, job_id, attempt_number, success, status_code, error_message, response_time_ms)

    # ==================== Crash Reports ====================

    async def save_crash_report(self, job_id: str, crash_data: dict) -> str:
        """
        Save a crash report and return the crash_id.

        Args:
            job_id: Job UUID as string
            crash_data: Dictionary containing crash report data:
                - crash_type: Type of crash (required)
                - error_message: Error message (optional)
                - state: Current state at crash time (required)
                - metrics_history: Historical metrics (optional)
                - logs_before_crash: Log entries before crash (optional)
                - analysis: Crash analysis data (optional)
                - screenshot_url: URL to screenshot (optional)
                - dom_snapshot_id: UUID of DOM snapshot (optional)

        Returns:
            UUID of created crash report as string
        """
        async with self.pool.acquire() as conn:
            # Convert job_id string to UUID
            job_uuid = UUID(job_id) if isinstance(job_id, str) else job_id

            crash_id = await conn.fetchval("""
                INSERT INTO crash_reports (
                    job_id,
                    crash_type,
                    error_message,
                    state,
                    metrics_history,
                    logs_before_crash,
                    analysis,
                    screenshot_url,
                    dom_snapshot_id
                )
                VALUES ($1, $2, $3, $4::jsonb, $5::jsonb, $6::jsonb, $7::jsonb, $8, $9)
                RETURNING crash_id
            """,
                job_uuid,
                crash_data.get('crash_type'),
                crash_data.get('error_message'),
                json.dumps(crash_data.get('state', {})),
                json.dumps(crash_data.get('metrics_history')) if crash_data.get('metrics_history') else None,
                json.dumps(crash_data.get('logs_before_crash')) if crash_data.get('logs_before_crash') else None,
                json.dumps(crash_data.get('analysis')) if crash_data.get('analysis') else None,
                crash_data.get('screenshot_url'),
                UUID(crash_data['dom_snapshot_id']) if crash_data.get('dom_snapshot_id') else None
            )

            log.info(f"Saved crash report {crash_id} for job {job_id}, type: {crash_data.get('crash_type')}")
            return str(crash_id)

    async def get_crash_report(self, job_id: str) -> Optional[dict]:
        """
        Get crash report for a job, if any.

        Args:
            job_id: Job UUID as string

        Returns:
            Crash report dictionary or None if not found
        """
        async with self.pool.acquire() as conn:
            job_uuid = UUID(job_id) if isinstance(job_id, str) else job_id

            row = await conn.fetchrow("""
                SELECT
                    crash_id,
                    job_id,
                    created_at,
                    crash_type,
                    error_message,
                    state,
                    metrics_history,
                    logs_before_crash,
                    analysis,
                    screenshot_url,
                    dom_snapshot_id
                FROM crash_reports
                WHERE job_id = $1
                ORDER BY created_at DESC
                LIMIT 1
            """, job_uuid)

            if not row:
                return None

            result = dict(row)
            # Convert UUIDs to strings for JSON serialization
            result['crash_id'] = str(result['crash_id'])
            result['job_id'] = str(result['job_id'])
            if result.get('dom_snapshot_id'):
                result['dom_snapshot_id'] = str(result['dom_snapshot_id'])

            return result

    async def get_crash_stats(self, days: int = 7) -> dict:
        """
        Get crash statistics for the last N days.

        Args:
            days: Number of days to look back (default: 7)

        Returns:
            Dictionary with:
                - total: Total number of crashes
                - by_type: Dict mapping crash type to count
                - by_day: List of dicts with date and count
        """
        async with self.pool.acquire() as conn:
            # Get total count
            total = await conn.fetchval("""
                SELECT COUNT(*)
                FROM crash_reports
                WHERE created_at >= NOW() - INTERVAL '%s days'
            """, days)

            # Get counts by type
            type_rows = await conn.fetch("""
                SELECT crash_type, COUNT(*) as count
                FROM crash_reports
                WHERE created_at >= NOW() - INTERVAL '%s days'
                GROUP BY crash_type
                ORDER BY count DESC
            """, days)

            by_type = {row['crash_type']: row['count'] for row in type_rows}

            # Get counts by day
            day_rows = await conn.fetch("""
                SELECT DATE(created_at) as date, COUNT(*) as count
                FROM crash_reports
                WHERE created_at >= NOW() - INTERVAL '%s days'
                GROUP BY DATE(created_at)
                ORDER BY date DESC
            """, days)

            by_day = [{'date': str(row['date']), 'count': row['count']} for row in day_rows]

            return {
                'total': total or 0,
                'by_type': by_type,
                'by_day': by_day
            }

    # ==================== API Key Operations ====================

    async def initialize_api_keys_schema(self):
        """Create api_keys table if it doesn't exist."""
        async with self.pool.acquire() as conn:
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS api_keys (
                    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
                    key_hash VARCHAR(64) NOT NULL UNIQUE,
                    key_prefix VARCHAR(8) NOT NULL,
                    name VARCHAR(255) NOT NULL,
                    client_id VARCHAR(255) NOT NULL,
                    scopes TEXT[] DEFAULT '{}',
                    rate_limit_rpm INTEGER DEFAULT 60,
                    is_active BOOLEAN DEFAULT true,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    last_used_at TIMESTAMP,
                    expires_at TIMESTAMP,
                    metadata JSONB
                );
            """)

            # Create indexes
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_api_keys_key_hash ON api_keys (key_hash);
            """)
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_api_keys_client_id ON api_keys (client_id);
            """)
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_api_keys_active ON api_keys (is_active) WHERE is_active = true;
            """)

            log.info("API keys schema initialized")

    async def get_api_key_by_hash(self, key_hash: str) -> Optional[Dict[str, Any]]:
        """
        Look up API key by its SHA-256 hash.

        This is the primary authentication lookup method. The hash is computed
        from the API key provided in the request header.

        Args:
            key_hash: SHA-256 hash of the API key (64 hex characters)

        Returns:
            API key record dictionary or None if not found:
            {
                "id": UUID,
                "key_prefix": "riq_a1b2",
                "name": "Production Key",
                "client_id": "veritas_123",
                "scopes": ["jobs:read", "jobs:write"],
                "rate_limit_rpm": 60,
                "is_active": True,
                "created_at": datetime,
                "last_used_at": datetime or None,
                "expires_at": datetime or None,
                "metadata": dict or None
            }
        """
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow("""
                SELECT
                    id,
                    key_prefix,
                    name,
                    client_id,
                    scopes,
                    rate_limit_rpm,
                    is_active,
                    created_at,
                    last_used_at,
                    expires_at,
                    metadata
                FROM api_keys
                WHERE key_hash = $1
            """, key_hash)

            if not row:
                return None

            result = dict(row)
            # Convert scopes from PostgreSQL array to Python list
            result['scopes'] = list(result['scopes']) if result['scopes'] else []
            return result

    async def create_api_key(
        self,
        client_id: str,
        name: str,
        scopes: List[str],
        rate_limit_rpm: int = 60,
        expires_at: Optional[datetime] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> tuple:
        """
        Create a new API key for a client.

        IMPORTANT: This method returns the plain text API key exactly once.
        After this, only the hash is stored - the key cannot be recovered.
        Make sure to display or securely transmit this key to the user.

        Args:
            client_id: External client identifier (e.g., "veritas_client_123")
            name: Human-readable name for the key (e.g., "Production API Key")
            scopes: List of permission scopes (e.g., ["jobs:read", "jobs:write"])
            rate_limit_rpm: Maximum requests per minute (default: 60)
            expires_at: Optional expiration datetime (None = never expires)
            metadata: Optional additional metadata dict

        Returns:
            Tuple of (plain_api_key, key_id):
                - plain_api_key: The full API key to give to the user (str)
                - key_id: UUID of the created key record

        Security Note:
            The plain_api_key is ONLY returned here. After creation, only
            the SHA-256 hash is stored. Never log or persist the plain key.
        """
        # Import here to avoid circular dependency
        from api.middleware.auth import generate_api_key, APIKeyAuth

        # Generate secure random key
        plain_api_key = generate_api_key()

        # Hash for storage
        key_hash = APIKeyAuth.hash_api_key(plain_api_key)

        # Extract prefix for identification
        key_prefix = APIKeyAuth.get_key_prefix(plain_api_key)

        async with self.pool.acquire() as conn:
            key_id = await conn.fetchval("""
                INSERT INTO api_keys (
                    key_hash,
                    key_prefix,
                    name,
                    client_id,
                    scopes,
                    rate_limit_rpm,
                    expires_at,
                    metadata
                )
                VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
                RETURNING id
            """,
                key_hash,
                key_prefix,
                name,
                client_id,
                scopes,
                rate_limit_rpm,
                expires_at,
                json.dumps(metadata) if metadata else None
            )

            # Log creation with only the prefix (never log full key)
            log.info(
                f"Created API key {key_prefix}... for client {client_id} "
                f"with scopes {scopes}"
            )

            return (plain_api_key, key_id)

    async def update_api_key_last_used(self, key_id: UUID):
        """
        Update the last_used_at timestamp for an API key.

        Called after each successful authentication to track key usage.
        This is non-blocking and failures are logged but not raised.

        Args:
            key_id: UUID of the API key record
        """
        async with self.pool.acquire() as conn:
            await conn.execute("""
                UPDATE api_keys
                SET last_used_at = NOW()
                WHERE id = $1
            """, key_id)

    async def revoke_api_key(self, key_id: UUID) -> bool:
        """
        Revoke an API key by setting is_active to false.

        This is preferred over deletion as it preserves audit history.
        The key will immediately become invalid for authentication.

        Args:
            key_id: UUID of the API key to revoke

        Returns:
            True if key was found and revoked, False if not found
        """
        async with self.pool.acquire() as conn:
            result = await conn.execute("""
                UPDATE api_keys
                SET is_active = false
                WHERE id = $1 AND is_active = true
            """, key_id)

            revoked = result.split()[-1] == "1"
            if revoked:
                log.info(f"Revoked API key {key_id}")
            return revoked

    async def delete_api_key(self, key_id: UUID) -> bool:
        """
        Permanently delete an API key.

        Use revoke_api_key instead if you want to preserve audit history.
        Deletion is permanent and cannot be undone.

        Args:
            key_id: UUID of the API key to delete

        Returns:
            True if deleted, False if not found
        """
        async with self.pool.acquire() as conn:
            result = await conn.execute("""
                DELETE FROM api_keys WHERE id = $1
            """, key_id)

            deleted = result.split()[-1] == "1"
            if deleted:
                log.info(f"Deleted API key {key_id}")
            return deleted

    async def list_api_keys_for_client(
        self,
        client_id: str,
        include_inactive: bool = False
    ) -> List[Dict[str, Any]]:
        """
        List all API keys for a specific client.

        Note: This returns key metadata only, never the actual keys
        (since we only store hashes).

        Args:
            client_id: Client identifier to filter by
            include_inactive: Whether to include revoked keys (default: False)

        Returns:
            List of API key records (without key_hash for security)
        """
        async with self.pool.acquire() as conn:
            if include_inactive:
                rows = await conn.fetch("""
                    SELECT
                        id,
                        key_prefix,
                        name,
                        client_id,
                        scopes,
                        rate_limit_rpm,
                        is_active,
                        created_at,
                        last_used_at,
                        expires_at
                    FROM api_keys
                    WHERE client_id = $1
                    ORDER BY created_at DESC
                """, client_id)
            else:
                rows = await conn.fetch("""
                    SELECT
                        id,
                        key_prefix,
                        name,
                        client_id,
                        scopes,
                        rate_limit_rpm,
                        is_active,
                        created_at,
                        last_used_at,
                        expires_at
                    FROM api_keys
                    WHERE client_id = $1 AND is_active = true
                    ORDER BY created_at DESC
                """, client_id)

            results = []
            for row in rows:
                record = dict(row)
                record['scopes'] = list(record['scopes']) if record['scopes'] else []
                results.append(record)

            return results

    async def get_api_key_by_id(self, key_id: UUID) -> Optional[Dict[str, Any]]:
        """
        Get API key metadata by its ID.

        Note: This returns key metadata only, never the actual key
        (since we only store hashes).

        Args:
            key_id: UUID of the API key

        Returns:
            API key record dictionary or None if not found
        """
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow("""
                SELECT
                    id,
                    key_prefix,
                    name,
                    client_id,
                    scopes,
                    rate_limit_rpm,
                    is_active,
                    created_at,
                    last_used_at,
                    expires_at,
                    metadata
                FROM api_keys
                WHERE id = $1
            """, key_id)

            if not row:
                return None

            result = dict(row)
            result['scopes'] = list(result['scopes']) if result['scopes'] else []
            return result

    async def update_api_key_scopes(
        self,
        key_id: UUID,
        scopes: List[str]
    ) -> bool:
        """
        Update the scopes for an API key.

        Args:
            key_id: UUID of the API key
            scopes: New list of permission scopes

        Returns:
            True if updated, False if key not found
        """
        async with self.pool.acquire() as conn:
            result = await conn.execute("""
                UPDATE api_keys
                SET scopes = $2
                WHERE id = $1
            """, key_id, scopes)

            updated = result.split()[-1] == "1"
            if updated:
                log.info(f"Updated scopes for API key {key_id}: {scopes}")
            return updated
