batch_embed.py

#!/usr/bin/env python3
"""
Batch embedding pipeline for coomer_viewer content.

Generates embeddings for existing media content using:
- OpenCLIP ViT-L/14 (768-dim) for images and video thumbnails
- BGE-M3 (truncated to 768-dim via Matryoshka) for tag text

Usage:
    python batch_embed.py --mode image --batch-size 64 --limit 1000
    python batch_embed.py --mode video --batch-size 32                 # multi-frame (4 frames per video)
    python batch_embed.py --mode tags --batch-size 128
    python batch_embed.py --mode aggregate-posts                       # build clip_*_post rows from frame rows
    python batch_embed.py --mode all                                    # image -> video -> tags -> aggregate-posts
    python batch_embed.py --mode image --dry-run

Schema notes (post-migration 009):
    - Every embedding row carries an `aggregation` column ∈ {'frame','post'}.
    - 'frame' rows are per-piece-of-media (1 per image, N per video). Used as
      intermediate state for computing post-level vectors.
    - 'post' rows are the canonical per-post vectors used by the similarity APIs.
      Types: clip_image_post, clip_video_post, clip_comic_post, bge_tags.
    - The DiskANN index is partial on aggregation='post' — frame rows are
      unindexed for ANN.
"""

import argparse
import io
import json
import logging
import os
import subprocess
import sys
import time
from collections import defaultdict
from pathlib import Path
from typing import Optional

import numpy as np
import psycopg2
import psycopg2.extras
from PIL import Image

# Lazy-loaded heavy dependencies (torch, open_clip, sentence_transformers)
# to allow import without GPU libs in test environments.

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.FileHandler("batch_embed.log"),
        logging.StreamHandler(),
    ],
)
log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

DB_URL = os.environ.get("DB_URL", "postgresql://coomer:coomer_dev@localhost:5433/coomer")
NAS_BASE = os.environ.get(
    "NAS_BASE",
    "/mnt/nas/drive_2",
)
EXTERNAL_CONTENT_BASE = os.path.join(NAS_BASE, "onlyfans", "external_content")

_OLD_GVFS_PREFIX = "/run/user/1000/gvfs/smb-share:server=unas-pro-8.local,share=drive_2"
CLIP_MODEL_NAME = "ViT-L-14"
CLIP_PRETRAINED = "laion2b_s32b_b82k"
CLIP_MODEL_VERSION = "ViT-L/14"
BGE_MODEL_NAME = "BAAI/bge-m3"
BGE_MODEL_VERSION = "bge-m3"
EMBEDDING_DIM = 768

CHECKPOINT_FILE = "batch_embed_checkpoint.json"

# Progress logging interval
LOG_EVERY = 50

# Multi-frame video embedding
VIDEO_FRAMES_PER_POST = 4
VIDEO_FRAME_CANDIDATES = 10  # smart_thumbs samples 10 candidates, we keep top-4

# Comic post-level aggregation
COMIC_PAGE_SAMPLE_CAP = 20  # Comics with > N pages get a stratified sample of N

# Embedding-type values written by this pipeline
EMB_CLIP_IMAGE = "clip_image"
EMB_CLIP_VIDEO_FRAME = "clip_video_frame"
EMB_BGE_TAGS = "bge_tags"
EMB_CLIP_IMAGE_POST = "clip_image_post"
EMB_CLIP_VIDEO_POST = "clip_video_post"
EMB_CLIP_COMIC_POST = "clip_comic_post"

# Aggregation values
AGG_FRAME = "frame"
AGG_POST = "post"

# ---------------------------------------------------------------------------
# Lazy model singletons
# ---------------------------------------------------------------------------

_clip_model = None
_clip_preprocess = None
_clip_tokenizer = None
_tag_model = None
_device = None


def _get_device():
    """Get torch device (cuda if available, else cpu)."""
    global _device
    if _device is None:
        import torch
        _device = "cuda" if torch.cuda.is_available() else "cpu"
        log.info("Using device: %s", _device)
    return _device


def _load_clip():
    """Lazy-load OpenCLIP model and preprocessing.

    Tries CUDA first; falls back to CPU on OOM (e.g. when another GPU job is
    saturating VRAM). Once the device is chosen it sticks for the process —
    no per-batch device juggling.
    """
    global _clip_model, _clip_preprocess, _clip_tokenizer, _device
    if _clip_model is None:
        import torch
        import open_clip

        log.info("Loading OpenCLIP %s (%s)...", CLIP_MODEL_NAME, CLIP_PRETRAINED)
        model, _, preprocess = open_clip.create_model_and_transforms(
            CLIP_MODEL_NAME, pretrained=CLIP_PRETRAINED
        )
        model.eval()
        device = _get_device()
        if device == "cuda":
            try:
                model = model.cuda()
            except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
                log.warning("CUDA OOM loading CLIP, falling back to CPU: %s", e)
                _device = "cpu"
                model = model.cpu()
                try:
                    torch.cuda.empty_cache()
                except Exception:
                    pass
        _clip_model = model
        _clip_preprocess = preprocess
        _clip_tokenizer = open_clip.get_tokenizer(CLIP_MODEL_NAME)
        log.info("OpenCLIP loaded on %s.", _device)
    return _clip_model, _clip_preprocess, _clip_tokenizer


def _load_tag_model():
    """Lazy-load BGE-M3 sentence transformer for tag embeddings."""
    global _tag_model
    if _tag_model is None:
        from sentence_transformers import SentenceTransformer

        log.info("Loading %s...", BGE_MODEL_NAME)
        _tag_model = SentenceTransformer(BGE_MODEL_NAME)
        log.info("BGE-M3 loaded.")
    return _tag_model


# ---------------------------------------------------------------------------
# Path resolution
# ---------------------------------------------------------------------------

def resolve_media_path(local_path: str) -> Optional[str]:
    """Resolve a media file path, handling GVFS rewrites and relative paths.

    Returns the resolved path if the file exists, or None if not found.
    """
    if not local_path:
        return None

    # Rewrite old GVFS paths to current NAS_BASE
    if local_path.startswith(_OLD_GVFS_PREFIX):
        local_path = NAS_BASE + local_path[len(_OLD_GVFS_PREFIX):]

    # Check absolute path
    if os.path.isabs(local_path):
        if os.path.isfile(local_path):
            return local_path
        # Not found at absolute path
        return None

    # Try relative to external content base
    resolved = os.path.join(EXTERNAL_CONTENT_BASE, local_path)
    if os.path.isfile(resolved):
        return resolved

    # Try relative to NAS base
    resolved = os.path.join(NAS_BASE, local_path)
    if os.path.isfile(resolved):
        return resolved

    return None


# ---------------------------------------------------------------------------
# Video frame extraction
# ---------------------------------------------------------------------------

def extract_frame(video_path: str, position_pct: float = 0.25) -> Optional[Image.Image]:
    """Extract a single frame from a video at position_pct of duration.

    Uses ffprobe to get duration, then ffmpeg to extract the frame as PNG
    piped to stdout. Returns a PIL Image or None on failure.
    """
    # Get duration via ffprobe
    try:
        probe = subprocess.run(
            [
                "ffprobe", "-v", "error",
                "-show_entries", "format=duration",
                "-of", "json",
                video_path,
            ],
            capture_output=True,
            timeout=30,
        )
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return None

    if probe.returncode != 0:
        return None

    try:
        info = json.loads(probe.stdout)
        duration = float(info.get("format", {}).get("duration", 0))
    except (json.JSONDecodeError, ValueError, TypeError):
        return None

    if duration < 0.1:
        return None

    seek_time = duration * position_pct

    # Extract frame via ffmpeg, output PNG to stdout
    try:
        result = subprocess.run(
            [
                "ffmpeg", "-y",
                "-ss", str(seek_time),
                "-i", video_path,
                "-frames:v", "1",
                "-f", "image2pipe",
                "-vcodec", "png",
                "pipe:1",
            ],
            capture_output=True,
            timeout=30,
        )
    except (subprocess.TimeoutExpired, FileNotFoundError):
        return None

    if result.returncode != 0 or not result.stdout:
        return None

    try:
        return Image.open(io.BytesIO(result.stdout)).convert("RGB")
    except Exception:
        return None


# ---------------------------------------------------------------------------
# Embedding functions
# ---------------------------------------------------------------------------

def embed_image(image_path: str) -> Optional[np.ndarray]:
    """Encode a single image file with OpenCLIP. Returns 768-dim vector or None."""
    if not os.path.isfile(image_path):
        return None

    try:
        img = Image.open(image_path).convert("RGB")
    except Exception:
        return None

    return embed_image_pil(img)


def embed_image_pil(img: Image.Image) -> Optional[np.ndarray]:
    """Encode a PIL Image with OpenCLIP. Returns 768-dim vector or None."""
    import torch

    model, preprocess, _ = _load_clip()
    device = _get_device()

    try:
        image_tensor = preprocess(img).unsqueeze(0)
        if device == "cuda":
            image_tensor = image_tensor.cuda()

        with torch.no_grad():
            if device == "cuda":
                with torch.cuda.amp.autocast():
                    features = model.encode_image(image_tensor)
            else:
                features = model.encode_image(image_tensor)

        # Normalize
        features = features / features.norm(dim=-1, keepdim=True)
        embedding = features.cpu().numpy().astype(np.float32).flatten()

        if embedding.shape[0] != EMBEDDING_DIM:
            log.warning("Unexpected embedding dim %d, expected %d", embedding.shape[0], EMBEDDING_DIM)
            return None

        return embedding

    except Exception as e:
        log.warning("Failed to embed image: %s", e)
        return None


def embed_tags(text: str) -> Optional[np.ndarray]:
    """Encode tag text with BGE-M3. Returns 768-dim vector (Matryoshka truncation) or None."""
    if not text or not text.strip():
        return None

    model = _load_tag_model()

    try:
        # BGE-M3 produces 1024-dim embeddings; truncate to 768 via Matryoshka
        embedding = model.encode([text], normalize_embeddings=True)
        vec = embedding[0].astype(np.float32)

        # Matryoshka truncation: take first 768 dims and re-normalize
        if vec.shape[0] > EMBEDDING_DIM:
            vec = vec[:EMBEDDING_DIM]
            norm = np.linalg.norm(vec)
            if norm > 0:
                vec = vec / norm

        return vec

    except Exception as e:
        log.warning("Failed to embed tags: %s", e)
        return None


def format_tag_text(
    creator_name: str,
    title: str,
    tags: list[tuple[str, str]],
) -> str:
    """Format post metadata and tags into structured text for embedding.

    Args:
        creator_name: The creator/artist name.
        title: The post title.
        tags: List of (category, tag_name) tuples.

    Returns:
        Formatted text string like:
        "artist: X. title: Y. character: a, b. general: c, d, e."
    """
    parts = []

    if creator_name:
        parts.append(f"artist: {creator_name}")

    if title:
        parts.append(f"title: {title}")

    if tags:
        # Group tags by category
        grouped = defaultdict(list)
        for category, tag_name in tags:
            cat = category.strip().lower() if category else "general"
            grouped[cat].append(tag_name)

        # Output each category
        for category in sorted(grouped.keys()):
            tag_list = ", ".join(grouped[category])
            parts.append(f"{category}: {tag_list}")

    return ". ".join(parts) + "." if parts else ""


# ---------------------------------------------------------------------------
# Database queries
# ---------------------------------------------------------------------------

def query_unembedded_images(conn, limit: Optional[int] = None) -> list:
    """Query media_files for images that don't have clip_image embeddings yet.

    file_type may be plain 'image'/'photo' (set by coomer_nas_downloader._classify_file)
    OR a mime-prefixed value like 'image/jpg', 'image/webp' (set by older ingest paths).
    The LIKE catches both. Verified against media_files breakdown 2026-04-27:
    448K image, 73K image/jpg, 17K image/webp, 12K image/png, 4 image/gif.

    Returns list of (post_id, media_file_id, creator_id, local_path).
    """
    sql = """
        SELECT mf.post_id, mf.id AS media_file_id, mf.creator_id, mf.local_path
        FROM media_files mf
        LEFT JOIN content_embeddings ce
            ON ce.media_file_id = mf.id
            AND ce.embedding_type = 'clip_image'
            AND ce.model_version = %s
        WHERE (mf.file_type LIKE 'image%%' OR mf.file_type = 'photo')
          AND mf.local_path IS NOT NULL
          AND ce.id IS NULL
        ORDER BY mf.id DESC
    """
    params = [CLIP_MODEL_VERSION]

    if limit is not None:
        sql += " LIMIT %s"
        params.append(limit)

    with conn.cursor() as cur:
        cur.execute(sql, params)
        return cur.fetchall()


def query_unembedded_videos(conn, limit: Optional[int] = None) -> list:
    """Query media_files for videos that don't have a post-level clip_video_post row yet.

    Multi-frame backfill: a video is "embedded" once its post has a clip_video_post
    aggregate vector. We anti-join on the post-level row, not on individual frames,
    so a partially-embedded video (e.g. 2/4 frames written before a crash) gets
    fully re-embedded on resume — the upsert path handles dedup.
    """
    sql = """
        SELECT mf.post_id, mf.id AS media_file_id, mf.creator_id, mf.local_path
        FROM media_files mf
        LEFT JOIN content_embeddings ce
            ON ce.post_id = mf.post_id
            AND ce.embedding_type = %s
            AND ce.aggregation = 'post'
            AND ce.model_version = %s
        WHERE mf.file_type = 'video'
          AND mf.local_path IS NOT NULL
          AND ce.id IS NULL
        ORDER BY mf.id DESC
    """
    params = [EMB_CLIP_VIDEO_POST, CLIP_MODEL_VERSION]

    if limit is not None:
        sql += " LIMIT %s"
        params.append(limit)

    with conn.cursor() as cur:
        cur.execute(sql, params)
        return cur.fetchall()


def query_unembedded_tags(conn, limit: Optional[int] = None) -> list:
    """Query posts with their tags that don't have bge_tags embeddings yet.

    Returns list of (post_id, creator_id, creator_name, title, tag_category, tag_name)
    rows. Caller must group by post_id.
    """
    sql = """
        SELECT p.id AS post_id, p.creator_id,
               COALESCE(c.username, '') AS creator_name,
               COALESCE(p.title, '') AS title,
               COALESCE(t.category, 'general') AS tag_category,
               t.name AS tag_name
        FROM posts p
        JOIN post_tags pt ON pt.post_id = p.id
        JOIN tags t ON t.id = pt.tag_id
        LEFT JOIN creators c ON c.id = p.creator_id
        LEFT JOIN content_embeddings ce
            ON ce.post_id = p.id
            AND ce.embedding_type = 'bge_tags'
            AND ce.model_version = %s
        WHERE ce.id IS NULL
        ORDER BY p.id DESC
    """
    params = [BGE_MODEL_VERSION]

    if limit is not None:
        sql += " LIMIT %s"
        params.append(limit)

    with conn.cursor() as cur:
        cur.execute(sql, params)
        return cur.fetchall()


def count_unembedded(conn, mode: str) -> int:
    """Count items eligible for embedding (used by --dry-run).

    Args:
        mode: One of 'image', 'video', 'tags'.

    Returns:
        Count of unembedded items.
    """
    if mode == "image":
        sql = """
            SELECT COUNT(*)
            FROM media_files mf
            LEFT JOIN content_embeddings ce
                ON ce.media_file_id = mf.id
                AND ce.embedding_type = 'clip_image'
                AND ce.model_version = %s
            WHERE (mf.file_type LIKE 'image%%' OR mf.file_type = 'photo')
              AND mf.local_path IS NOT NULL
              AND ce.id IS NULL
        """
        params = [CLIP_MODEL_VERSION]
    elif mode == "video":
        sql = """
            SELECT COUNT(*)
            FROM media_files mf
            LEFT JOIN content_embeddings ce
                ON ce.post_id = mf.post_id
                AND ce.embedding_type = %s
                AND ce.aggregation = 'post'
                AND ce.model_version = %s
            WHERE mf.file_type = 'video'
              AND mf.local_path IS NOT NULL
              AND ce.id IS NULL
        """
        params = [EMB_CLIP_VIDEO_POST, CLIP_MODEL_VERSION]
    elif mode == "tags":
        sql = """
            SELECT COUNT(DISTINCT p.id)
            FROM posts p
            JOIN post_tags pt ON pt.post_id = p.id
            LEFT JOIN content_embeddings ce
                ON ce.post_id = p.id
                AND ce.embedding_type = 'bge_tags'
                AND ce.model_version = %s
            WHERE ce.id IS NULL
        """
        params = [BGE_MODEL_VERSION]
    else:
        return 0

    with conn.cursor() as cur:
        cur.execute(sql, params)
        row = cur.fetchone()
        return row[0] if row else 0


# ---------------------------------------------------------------------------
# Batch insert
# ---------------------------------------------------------------------------

def insert_embeddings(conn, batch: list) -> None:
    """Insert a batch of embeddings using execute_values with ON CONFLICT DO NOTHING.

    Each item in batch is a tuple of 9 fields:
        (post_id, media_file_id, creator_id, embedding_list, embedding_type,
         model_version, aggregation, frame_index, frame_timestamp_sec)

    The two partial unique indexes uq_embedding_frame and uq_embedding_post
    handle the dedup; we use a target-less ON CONFLICT DO NOTHING because
    psycopg2 can't infer a single conflict target across two partial indexes.
    """
    with conn.cursor() as cur:
        psycopg2.extras.execute_values(
            cur,
            """INSERT INTO content_embeddings
                   (post_id, media_file_id, creator_id, embedding, embedding_type,
                    model_version, aggregation, frame_index, frame_timestamp_sec)
               VALUES %s
               ON CONFLICT DO NOTHING""",
            batch,
            template="(%s, %s, %s, %s::vector, %s, %s, %s, %s, %s)",
        )
    conn.commit()


def upsert_embedding(conn, row: tuple) -> None:
    """Insert or update a single embedding row, replacing the embedding on conflict.

    Targets the appropriate partial unique index by column + WHERE clause —
    Postgres's ON CONFLICT (cols) WHERE <pred> resolves to the matching partial
    index. Constraint-name targeting (ON CONFLICT ON CONSTRAINT ...) does not
    work because uq_embedding_post / uq_embedding_frame are unique INDEXES, not
    table constraints.
    """
    aggregation = row[6]
    if aggregation == AGG_POST:
        on_conflict = (
            "ON CONFLICT (post_id, embedding_type, model_version) WHERE aggregation='post'"
        )
    else:
        on_conflict = (
            "ON CONFLICT (post_id, media_file_id, embedding_type, model_version, COALESCE(frame_index, -1)) "
            "WHERE aggregation='frame'"
        )

    with conn.cursor() as cur:
        cur.execute(
            f"""INSERT INTO content_embeddings
                   (post_id, media_file_id, creator_id, embedding, embedding_type,
                    model_version, aggregation, frame_index, frame_timestamp_sec)
               VALUES (%s, %s, %s, %s::vector, %s, %s, %s, %s, %s)
               {on_conflict} DO UPDATE
                  SET embedding = EXCLUDED.embedding,
                      frame_timestamp_sec = EXCLUDED.frame_timestamp_sec""",
            row,
        )
    conn.commit()


# ---------------------------------------------------------------------------
# Batch processing pipelines
# ---------------------------------------------------------------------------

def process_image_batch(conn, items: list, batch_size: int = 64) -> dict:
    """Process a batch of images: resolve path, embed, insert.

    Args:
        conn: psycopg2 connection.
        items: List of (post_id, media_file_id, creator_id, local_path).
        batch_size: Number of embeddings to insert at once.

    Returns:
        Dict with stats: {"embedded": int, "skipped": int, "errors": int}.
    """
    stats = {"embedded": 0, "skipped": 0, "errors": 0}
    pending = []

    for i, (post_id, media_file_id, creator_id, local_path) in enumerate(items):
        resolved = resolve_media_path(local_path)
        if resolved is None:
            stats["skipped"] += 1
            continue

        embedding = embed_image(resolved)
        if embedding is None:
            stats["skipped"] += 1
            continue

        emb_list = embedding.tolist()
        # Frame-level row.
        pending.append((
            post_id, media_file_id, creator_id, emb_list,
            EMB_CLIP_IMAGE, CLIP_MODEL_VERSION, AGG_FRAME, None, None,
        ))
        # Post-level row: for image posts, the post vector IS the single image
        # vector (duplicated so the partial DiskANN index can be cleanly
        # partitioned by aggregation='post'). Comic-only posts are aggregated
        # later in --mode aggregate-posts using mean-pool over pages.
        pending.append((
            post_id, None, creator_id, emb_list,
            EMB_CLIP_IMAGE_POST, CLIP_MODEL_VERSION, AGG_POST, None, None,
        ))

        if len(pending) >= batch_size:
            try:
                insert_embeddings(conn, pending)
                stats["embedded"] += len(pending)
            except Exception as e:
                log.error("Insert batch failed: %s", e)
                stats["errors"] += len(pending)
                conn.rollback()
            pending = []

        if (i + 1) % LOG_EVERY == 0:
            log.info("Images: %d/%d processed (%d embedded, %d skipped)",
                     i + 1, len(items), stats["embedded"], stats["skipped"])

    # Flush remaining
    if pending:
        try:
            insert_embeddings(conn, pending)
            stats["embedded"] += len(pending)
        except Exception as e:
            log.error("Insert final batch failed: %s", e)
            stats["errors"] += len(pending)
            conn.rollback()

    return stats


def process_video_batch(conn, items: list, batch_size: int = 32) -> dict:
    """Process a batch of videos with multi-frame embedding.

    For each video:
      1. Use smart_thumbs.select_top_k_frames to pick the top VIDEO_FRAMES_PER_POST
         (default 4) frames from VIDEO_FRAME_CANDIDATES (default 10) candidates.
      2. Embed each frame with OpenCLIP -> 4 clip_video_frame rows
         (frame_index 0..3, frame_timestamp_sec set).
      3. L2-normalized mean-pool the 4 frame vectors -> 1 clip_video_post row
         (aggregation='post', media_file_id=NULL).
      4. Clean up the temp frame files.

    items: List of (post_id, media_file_id, creator_id, local_path).
    """
    from smart_thumbs import select_top_k_frames

    stats = {"embedded": 0, "skipped": 0, "errors": 0, "videos_done": 0}
    pending = []

    for i, (post_id, media_file_id, creator_id, local_path) in enumerate(items):
        resolved = resolve_media_path(local_path)
        if resolved is None:
            stats["skipped"] += 1
            continue

        ranked = select_top_k_frames(
            resolved,
            k=VIDEO_FRAMES_PER_POST,
            num_frames=VIDEO_FRAME_CANDIDATES,
        )
        if not ranked:
            stats["skipped"] += 1
            continue

        frame_vectors = []
        frame_paths = [s["path"] for s in ranked]

        try:
            for rank, scored in enumerate(ranked):
                try:
                    img = Image.open(scored["path"]).convert("RGB")
                except Exception:
                    continue
                vec = embed_image_pil(img)
                if vec is None:
                    continue

                frame_vectors.append(vec)
                pending.append((
                    post_id, media_file_id, creator_id, vec.tolist(),
                    EMB_CLIP_VIDEO_FRAME, CLIP_MODEL_VERSION, AGG_FRAME,
                    rank, float(scored.get("frame_timestamp_sec", 0.0)),
                ))

            # Post-level: L2-normalized mean-pool of the frame vectors we got.
            if frame_vectors:
                stacked = np.stack(frame_vectors, axis=0)
                pooled = stacked.mean(axis=0)
                norm = np.linalg.norm(pooled)
                if norm > 0:
                    pooled = pooled / norm
                pending.append((
                    post_id, None, creator_id, pooled.tolist(),
                    EMB_CLIP_VIDEO_POST, CLIP_MODEL_VERSION, AGG_POST, None, None,
                ))
                stats["videos_done"] += 1
            else:
                stats["skipped"] += 1
        finally:
            # Always clean up extracted frame files.
            for p in frame_paths:
                try:
                    if os.path.isfile(p):
                        os.unlink(p)
                except OSError:
                    pass

        if len(pending) >= batch_size:
            try:
                insert_embeddings(conn, pending)
                stats["embedded"] += len(pending)
            except Exception as e:
                log.error("Insert batch failed: %s", e)
                stats["errors"] += len(pending)
                conn.rollback()
            pending = []

        if (i + 1) % LOG_EVERY == 0:
            log.info("Videos: %d/%d processed (%d embedded rows, %d videos done, %d skipped)",
                     i + 1, len(items), stats["embedded"], stats["videos_done"], stats["skipped"])

    if pending:
        try:
            insert_embeddings(conn, pending)
            stats["embedded"] += len(pending)
        except Exception as e:
            log.error("Insert final batch failed: %s", e)
            stats["errors"] += len(pending)
            conn.rollback()

    return stats


def process_tag_batch(conn, items: list, batch_size: int = 128) -> dict:
    """Process a batch of post tag data: format text, embed, insert.

    Args:
        conn: psycopg2 connection.
        items: List of dicts with keys:
            post_id, creator_id, creator_name, title,
            tags [(category, tag_name), ...].
        batch_size: Number of embeddings to insert at once.

    Returns:
        Dict with stats.
    """
    stats = {"embedded": 0, "skipped": 0, "errors": 0}
    pending = []

    for i, item in enumerate(items):
        text = format_tag_text(
            creator_name=item["creator_name"],
            title=item["title"],
            tags=item["tags"],
        )

        if not text:
            stats["skipped"] += 1
            continue

        embedding = embed_tags(text)
        if embedding is None:
            stats["skipped"] += 1
            continue

        pending.append((
            item["post_id"],
            None,  # no media_file_id for tag embeddings
            item["creator_id"],
            embedding.tolist(),
            EMB_BGE_TAGS,
            BGE_MODEL_VERSION,
            AGG_POST,
            None,
            None,
        ))

        if len(pending) >= batch_size:
            try:
                insert_embeddings(conn, pending)
                stats["embedded"] += len(pending)
            except Exception as e:
                log.error("Insert batch failed: %s", e)
                stats["errors"] += len(pending)
                conn.rollback()
            pending = []

        if (i + 1) % LOG_EVERY == 0:
            log.info("Tags: %d/%d processed (%d embedded, %d skipped)",
                     i + 1, len(items), stats["embedded"], stats["skipped"])

    if pending:
        try:
            insert_embeddings(conn, pending)
            stats["embedded"] += len(pending)
        except Exception as e:
            log.error("Insert final batch failed: %s", e)
            stats["errors"] += len(pending)
            conn.rollback()

    return stats


def process_aggregate_posts(conn, limit: Optional[int] = None) -> dict:
    """Build clip_*_post rows from existing frame-level embeddings.

    Steps:
      1. Find posts that have multiple clip_image (frame) rows -> COMIC posts.
         For each, mean-pool the page vectors with stratified-N cap (where
         N = COMIC_PAGE_SAMPLE_CAP), L2-normalize, upsert as clip_comic_post.
         Also DELETE any (incorrect) clip_image_post rows for those comic posts
         (the inline path in process_image_batch optimistically inserts one
         per page, but for comics that representation is wrong).
      3. Single-image posts already have a correct clip_image_post row from
         process_image_batch, so we skip them here.
      4. Video posts: clip_video_post is built inline by process_video_batch
         (mean-pool of the 4 frame vectors); nothing to do here.
      5. bge_tags: aggregation='post' is set at insert time; nothing to do.

    Returns dict with stats.
    """
    stats = {"comics_pooled": 0, "comic_image_post_deleted": 0, "skipped": 0, "errors": 0}

    # Step 1: identify comic posts (>= 2 clip_image rows under one post).
    # We pull the page embeddings for each, sample, mean-pool in Python.
    sql_comic_posts = """
        SELECT post_id, count(*) AS n_pages
          FROM content_embeddings
         WHERE embedding_type = %s
           AND aggregation = 'frame'
           AND model_version = %s
         GROUP BY post_id
        HAVING count(*) >= 2
         ORDER BY post_id
    """
    params = [EMB_CLIP_IMAGE, CLIP_MODEL_VERSION]
    if limit is not None:
        sql_comic_posts += " LIMIT %s"
        params.append(limit)

    with conn.cursor() as cur:
        cur.execute(sql_comic_posts, params)
        comic_post_ids = [r[0] for r in cur.fetchall()]

    log.info("Comic posts to aggregate: %d", len(comic_post_ids))

    for i, post_id in enumerate(comic_post_ids):
        try:
            with conn.cursor() as cur:
                # Pull page-level vectors ordered by media_file_id so the
                # stratified sample is deterministic across reruns.
                cur.execute(
                    """SELECT ce.embedding::text, ce.creator_id
                         FROM content_embeddings ce
                        WHERE ce.post_id = %s
                          AND ce.embedding_type = %s
                          AND ce.aggregation = 'frame'
                          AND ce.model_version = %s
                        ORDER BY ce.media_file_id""",
                    (post_id, EMB_CLIP_IMAGE, CLIP_MODEL_VERSION),
                )
                rows = cur.fetchall()

            if len(rows) < 2:
                stats["skipped"] += 1
                continue

            # Parse vectors (pgvector returns "[0.1,0.2,...]" as text).
            vecs = np.array(
                [[float(x) for x in r[0].strip("[]").split(",")] for r in rows],
                dtype=np.float32,
            )
            creator_id = rows[0][1]

            # Stratified-N sample if over the cap.
            n_pages = vecs.shape[0]
            if n_pages > COMIC_PAGE_SAMPLE_CAP:
                idxs = np.linspace(0, n_pages - 1, COMIC_PAGE_SAMPLE_CAP).astype(int)
                vecs = vecs[idxs]

            pooled = vecs.mean(axis=0)
            norm = np.linalg.norm(pooled)
            if norm > 0:
                pooled = pooled / norm

            # Drop the (incorrect) per-page clip_image_post rows for this comic
            # before upserting the canonical clip_comic_post.
            with conn.cursor() as cur:
                cur.execute(
                    """DELETE FROM content_embeddings
                        WHERE post_id = %s
                          AND embedding_type = %s
                          AND aggregation = 'post'""",
                    (post_id, EMB_CLIP_IMAGE_POST),
                )
                stats["comic_image_post_deleted"] += cur.rowcount

            upsert_embedding(conn, (
                post_id, None, creator_id, pooled.tolist(),
                EMB_CLIP_COMIC_POST, CLIP_MODEL_VERSION, AGG_POST, None, None,
            ))
            stats["comics_pooled"] += 1

            if (i + 1) % LOG_EVERY == 0:
                log.info("Comics: %d/%d aggregated", i + 1, len(comic_post_ids))

        except Exception as e:
            log.error("Failed to aggregate comic post %s: %s", post_id, e)
            stats["errors"] += 1
            conn.rollback()

    # Step 6: video-trumps-image rule. Posts with BOTH images and videos got
    # both a clip_image_post (or clip_comic_post) and a clip_video_post —
    # the video one is canonical (matches /api/atlas:2670-2677 post_kind logic),
    # so drop the redundant image variants.
    with conn.cursor() as cur:
        cur.execute(
            """DELETE FROM content_embeddings ce_img
                WHERE ce_img.aggregation = 'post'
                  AND ce_img.embedding_type IN ('clip_image_post', 'clip_comic_post')
                  AND EXISTS (
                      SELECT 1 FROM content_embeddings ce_vid
                       WHERE ce_vid.post_id = ce_img.post_id
                         AND ce_vid.embedding_type = 'clip_video_post'
                         AND ce_vid.aggregation = 'post'
                  )"""
        )
        n_dropped = cur.rowcount
        conn.commit()
    if n_dropped:
        log.info("video-trumps-image: dropped %d redundant clip_*_post rows", n_dropped)
    stats["mixed_media_resolved"] = n_dropped

    return stats


def _group_tag_rows(rows: list) -> list[dict]:
    """Group raw tag query rows by post_id into structured dicts.

    Input rows: (post_id, creator_id, creator_name, title, tag_category, tag_name)
    Output: list of dicts with post_id, creator_id, creator_name, title, tags.
    """
    grouped = {}
    for post_id, creator_id, creator_name, title, tag_category, tag_name in rows:
        if post_id not in grouped:
            grouped[post_id] = {
                "post_id": post_id,
                "creator_id": creator_id,
                "creator_name": creator_name,
                "title": title,
                "tags": [],
            }
        if tag_name:
            grouped[post_id]["tags"].append((tag_category, tag_name))

    return list(grouped.values())


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args(argv: Optional[list] = None) -> argparse.Namespace:
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(
        description="Batch embedding pipeline for coomer_viewer content.",
    )
    parser.add_argument(
        "--mode",
        required=True,
        choices=["image", "video", "tags", "aggregate-posts", "all"],
        help=(
            "Embedding mode. image: single-image posts (writes clip_image + clip_image_post). "
            "video: multi-frame video embedding (writes 4×clip_video_frame + clip_video_post). "
            "tags: post tag embedding (writes bge_tags). "
            "aggregate-posts: build clip_comic_post for multi-image posts. "
            "all: image -> video -> tags -> aggregate-posts."
        ),
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=64,
        help="Number of items per DB insert batch (default: 64).",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Max number of items to process (default: all).",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        default=False,
        help="Count eligible items without processing.",
    )
    parser.add_argument(
        "--db-url",
        default=DB_URL,
        help="Database connection URL.",
    )
    return parser.parse_args(argv)


def main(argv: Optional[list] = None) -> None:
    """Main entry point."""
    args = parse_args(argv)

    log.info("Connecting to database...")
    conn = psycopg2.connect(args.db_url)
    conn.autocommit = False

    try:
        modes = (
            ["image", "video", "tags", "aggregate-posts"]
            if args.mode == "all"
            else [args.mode]
        )

        for mode in modes:
            log.info("=" * 60)
            log.info("Mode: %s", mode)

            if args.dry_run and mode != "aggregate-posts":
                count = count_unembedded(conn, mode)
                log.info("[DRY RUN] %d %s items eligible for embedding.", count, mode)
                continue

            # Each worklist query below opens a read transaction (autocommit is
            # off). End it BEFORE the GPU embedding loop — otherwise the
            # connection sits 'idle in transaction' through the first batch of
            # CLIP encodes and Postgres' role-level
            # idle_in_transaction_session_timeout aborts it ("SSL connection
            # has been closed unexpectedly"). insert_embeddings commits each
            # batch, so subsequent inter-batch GPU windows are merely 'idle'
            # (safe). _end_read_txn() makes the first window safe too.
            def _end_read_txn():
                try:
                    conn.rollback()
                except Exception:
                    pass

            if mode == "image":
                items = query_unembedded_images(conn, limit=args.limit)
                _end_read_txn()
                log.info("Found %d unembedded images.", len(items))
                if items:
                    stats = process_image_batch(conn, items, batch_size=args.batch_size)
                    log.info("Image stats: %s", stats)

            elif mode == "video":
                items = query_unembedded_videos(conn, limit=args.limit)
                _end_read_txn()
                log.info("Found %d unembedded videos.", len(items))
                if items:
                    stats = process_video_batch(conn, items, batch_size=args.batch_size)
                    log.info("Video stats: %s", stats)

            elif mode == "tags":
                rows = query_unembedded_tags(conn, limit=args.limit)
                items = _group_tag_rows(rows)
                _end_read_txn()
                log.info("Found %d posts with unembedded tags.", len(items))
                if items:
                    stats = process_tag_batch(conn, items, batch_size=args.batch_size)
                    log.info("Tag stats: %s", stats)

            elif mode == "aggregate-posts":
                if args.dry_run:
                    log.info("[DRY RUN] aggregate-posts has no preview; run without --dry-run.")
                    continue
                stats = process_aggregate_posts(conn, limit=args.limit)
                log.info("Aggregate stats: %s", stats)

    finally:
        conn.close()
        log.info("Done.")


if __name__ == "__main__":
    main()