utils/embedding_generator.py

"""
Spot Embedding Generator

Gemini text-embedding-004 모델을 사용하여 story_spots의 벡터 임베딩을 생성하고
Supabase spot_embeddings 테이블에 저장합니다.

Usage:
    from utils.embedding_generator import generate_all_embeddings, update_changed_embeddings

    # 전체 생성 (초기 1회)
    stats = await generate_all_embeddings()

    # 변경분만 업데이트 (주기적)
    stats = await update_changed_embeddings()

    # CLI 직접 실행
    python -m utils.embedding_generator [--force]
"""

import asyncio
import hashlib
import json
import logging
import os
import time
from typing import Any, Optional

from google import genai
from google.genai import types

logger = logging.getLogger(__name__)

# ============ Constants ============

EMBEDDING_MODEL = "text-embedding-004"
EMBEDDING_DIMENSIONS = 768
BATCH_SIZE = 100  # Gemini batch embed limit
RATE_LIMIT_DELAY = 0.5  # seconds between batch calls

# Gemini task types for embedding
TASK_TYPE_RETRIEVAL_DOCUMENT = "RETRIEVAL_DOCUMENT"
TASK_TYPE_RETRIEVAL_QUERY = "RETRIEVAL_QUERY"

# story_spots columns needed for embedding text construction
SPOT_SELECT_COLUMNS = (
    "id, name, name_en, name_zh, category, lat, lng, address, "
    "story_title, story_content, story_source, tips, "
    "tags_tier1, tags_tier2, meta, "
    "main_image_url, thumbnail_url, generated_image_url, "
    "priority_score, status, zone, cluster_id, "
    "village, source_book, historical_period"
)


# ============ Gemini Client ============

def _get_gemini_client() -> genai.Client:
    """Get Gemini API client (lazy init)."""
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("GEMINI_API_KEY environment variable is required")
    return genai.Client(api_key=api_key)


# ============ Content Hash ============

def compute_content_hash(spot: dict[str, Any]) -> str:
    """
    스팟 데이터의 콘텐츠 해시 생성.
    임베딩에 사용되는 필드만 포함하여,
    관련 없는 필드 변경으로 불필요한 재생성 방지.
    """
    hash_fields = {
        "name": spot.get("name", ""),
        "category": spot.get("category", ""),
        "story_title": spot.get("story_title", ""),
        "story_content": spot.get("story_content", ""),
        "tips": spot.get("tips", ""),
        "tags_tier1": json.dumps(spot.get("tags_tier1") or {}, sort_keys=True, ensure_ascii=False),
        "tags_tier2": json.dumps(spot.get("tags_tier2") or [], sort_keys=True, ensure_ascii=False),
        "village": spot.get("village", ""),
        "historical_period": spot.get("historical_period", ""),
    }
    content_str = json.dumps(hash_fields, sort_keys=True, ensure_ascii=False)
    return hashlib.sha256(content_str.encode("utf-8")).hexdigest()[:16]


# ============ Embedding Text Construction ============

def build_embedding_text(spot: dict[str, Any]) -> str:
    """
    스팟 데이터를 임베딩용 텍스트로 변환.

    구성:
    - 이름 (한국어 + 영어/중국어)
    - 카테고리
    - 스토리 제목 + 내용 (핵심 시맨틱 정보)
    - 태그 (테마, 분위기, 활동성)
    - 팁
    - 마을, 시대 정보
    """
    parts: list[str] = []

    # 이름 (다국어)
    name = spot.get("name", "")
    parts.append(f"장소: {name}")
    if spot.get("name_en"):
        parts.append(f"({spot['name_en']})")
    if spot.get("name_zh"):
        parts.append(f"({spot['name_zh']})")

    # 카테고리
    category = spot.get("category", "")
    if category:
        parts.append(f"분류: {category}")

    # 마을/시대 (향토지 메타데이터)
    if spot.get("village"):
        parts.append(f"마을: {spot['village']}")
    if spot.get("historical_period"):
        parts.append(f"시대: {spot['historical_period']}")

    # 스토리 (핵심 시맨틱)
    story_title = spot.get("story_title", "")
    story_content = spot.get("story_content", "")
    if story_title:
        parts.append(f"이야기: {story_title}")
    if story_content and "카테고리에 속합니다" not in story_content:
        # 자동 생성된 의미없는 콘텐츠 제외
        # 임베딩 토큰 예산: ~500자로 제한
        content_trimmed = story_content[:500]
        parts.append(content_trimmed)

    # 태그
    tags_tier1 = spot.get("tags_tier1") or {}
    themes = tags_tier1.get("theme", [])
    moods = tags_tier1.get("mood", [])
    activity = tags_tier1.get("activity_level", "")
    time_of_day = tags_tier1.get("time_of_day", [])

    tag_parts: list[str] = []
    if themes:
        tag_parts.append(f"테마: {', '.join(themes)}")
    if moods:
        tag_parts.append(f"분위기: {', '.join(moods)}")
    if activity:
        tag_parts.append(f"활동: {activity}")
    if time_of_day:
        tag_parts.append(f"시간대: {', '.join(time_of_day)}")
    if tag_parts:
        parts.append(" | ".join(tag_parts))

    # tier2 태그
    tags_tier2 = spot.get("tags_tier2") or []
    if tags_tier2:
        parts.append(f"키워드: {', '.join(tags_tier2)}")

    # 팁
    tips = spot.get("tips", "")
    if tips:
        parts.append(f"팁: {tips[:200]}")

    return "\n".join(parts)


# ============ Embedding API ============

async def generate_embeddings_batch(
    texts: list[str],
    client: Optional[genai.Client] = None,
    task_type: str = TASK_TYPE_RETRIEVAL_DOCUMENT,
) -> list[list[float]]:
    """
    Gemini text-embedding-004로 텍스트 배치 임베딩 생성.

    Args:
        texts: 임베딩할 텍스트 리스트 (최대 BATCH_SIZE)
        client: Gemini 클라이언트 (None이면 자동 생성)
        task_type: RETRIEVAL_DOCUMENT (저장용) 또는 RETRIEVAL_QUERY (검색용)

    Returns:
        임베딩 벡터 리스트 (각 768차원)
    """
    if not texts:
        return []

    if client is None:
        client = _get_gemini_client()

    if len(texts) > BATCH_SIZE:
        raise ValueError(f"Batch size {len(texts)} exceeds limit {BATCH_SIZE}")

    try:
        response = await asyncio.to_thread(
            client.models.embed_content,
            model=EMBEDDING_MODEL,
            contents=texts,
            config=types.EmbedContentConfig(
                task_type=task_type,
                output_dimensionality=EMBEDDING_DIMENSIONS,
            ),
        )

        embeddings = []
        for emb in response.embeddings:
            embeddings.append(list(emb.values))

        logger.info(
            f"[embedding] Generated {len(embeddings)} embeddings "
            f"(dim={len(embeddings[0]) if embeddings else 0})"
        )
        return embeddings

    except Exception as e:
        logger.error(f"[embedding] Gemini embed_content failed: {e}")
        raise


async def generate_query_embedding(
    query_text: str,
    client: Optional[genai.Client] = None,
) -> list[float]:
    """
    검색 쿼리를 임베딩으로 변환.
    RETRIEVAL_QUERY task type 사용 (문서 임베딩과 쌍을 이룸).

    Args:
        query_text: 검색 쿼리 텍스트

    Returns:
        768차원 임베딩 벡터
    """
    embeddings = await generate_embeddings_batch(
        [query_text],
        client=client,
        task_type=TASK_TYPE_RETRIEVAL_QUERY,
    )
    return embeddings[0]


# ============ DB Operations ============

def _load_active_spots() -> list[dict[str, Any]]:
    """Supabase에서 active 스팟 로드 (임베딩 생성에 필요한 컬럼만)."""
    from db import get_supabase
    supabase = get_supabase()

    result = supabase.table("story_spots") \
        .select(SPOT_SELECT_COLUMNS) \
        .eq("status", "active") \
        .execute()

    return result.data or []


def _load_existing_embeddings() -> dict[str, str]:
    """기존 임베딩의 spot_id -> content_hash 매핑 로드."""
    from db import get_supabase
    supabase = get_supabase()

    result = supabase.table("spot_embeddings") \
        .select("spot_id, content_hash") \
        .execute()

    return {row["spot_id"]: row["content_hash"] for row in (result.data or [])}


def _upsert_embeddings(records: list[dict[str, Any]]) -> int:
    """임베딩 레코드 배치 upsert. 성공 건수 반환."""
    if not records:
        return 0

    from db import get_supabase
    supabase = get_supabase()

    # Supabase upsert: spot_id UNIQUE constraint로 ON CONFLICT
    result = supabase.table("spot_embeddings") \
        .upsert(records, on_conflict="spot_id") \
        .execute()

    return len(result.data) if result.data else 0


def _delete_orphan_embeddings(active_spot_ids: set[str]) -> int:
    """더 이상 active가 아닌 스팟의 임베딩 삭제."""
    from db import get_supabase
    supabase = get_supabase()

    # 현재 임베딩이 있는 spot_id 조회
    result = supabase.table("spot_embeddings") \
        .select("spot_id") \
        .execute()

    orphan_ids = [
        row["spot_id"] for row in (result.data or [])
        if row["spot_id"] not in active_spot_ids
    ]

    if not orphan_ids:
        return 0

    for orphan_id in orphan_ids:
        supabase.table("spot_embeddings") \
            .delete() \
            .eq("spot_id", orphan_id) \
            .execute()

    logger.info(f"[embedding] Deleted {len(orphan_ids)} orphan embeddings")
    return len(orphan_ids)


# ============ Main Functions ============

async def generate_all_embeddings(force: bool = False) -> dict[str, Any]:
    """
    전체 스팟 임베딩 생성.

    Args:
        force: True면 content_hash 무시하고 전부 재생성

    Returns:
        {
            "total_spots": int,
            "generated": int,
            "skipped": int,
            "deleted_orphans": int,
            "errors": int,
            "duration_seconds": float
        }
    """
    start = time.monotonic()
    stats = {
        "total_spots": 0,
        "generated": 0,
        "skipped": 0,
        "deleted_orphans": 0,
        "errors": 0,
        "duration_seconds": 0.0,
    }

    try:
        # 1. Active 스팟 로드
        spots = await asyncio.to_thread(_load_active_spots)
        stats["total_spots"] = len(spots)
        logger.info(f"[embedding] Loaded {len(spots)} active spots")

        if not spots:
            return stats

        # 2. 기존 임베딩 해시 로드
        existing_hashes = await asyncio.to_thread(_load_existing_embeddings)
        logger.info(f"[embedding] Found {len(existing_hashes)} existing embeddings")

        # 3. 변경 감지: 새로운/변경된 스팟만 필터링
        spots_to_embed: list[dict[str, Any]] = []
        for spot in spots:
            content_hash = compute_content_hash(spot)
            spot["_content_hash"] = content_hash

            if force or spot["id"] not in existing_hashes:
                spots_to_embed.append(spot)
            elif existing_hashes[spot["id"]] != content_hash:
                spots_to_embed.append(spot)
            else:
                stats["skipped"] += 1

        logger.info(
            f"[embedding] To embed: {len(spots_to_embed)}, "
            f"skipped (unchanged): {stats['skipped']}"
        )

        # 4. 배치 임베딩 생성
        client = _get_gemini_client()

        for batch_start in range(0, len(spots_to_embed), BATCH_SIZE):
            batch = spots_to_embed[batch_start:batch_start + BATCH_SIZE]
            batch_texts = [build_embedding_text(s) for s in batch]

            try:
                embeddings = await generate_embeddings_batch(
                    batch_texts, client=client
                )

                # DB upsert 레코드 준비
                records = []
                for spot, embedding in zip(batch, embeddings):
                    records.append({
                        "spot_id": spot["id"],
                        "embedding": embedding,
                        "embedding_model": EMBEDDING_MODEL,
                        "content_hash": spot["_content_hash"],
                        "updated_at": "now()",
                    })

                upserted = await asyncio.to_thread(_upsert_embeddings, records)
                stats["generated"] += upserted
                logger.info(
                    f"[embedding] Batch {batch_start // BATCH_SIZE + 1}: "
                    f"embedded {upserted}/{len(batch)} spots"
                )

            except Exception as e:
                stats["errors"] += len(batch)
                logger.error(
                    f"[embedding] Batch {batch_start // BATCH_SIZE + 1} failed: {e}"
                )

            # Rate limiting
            if batch_start + BATCH_SIZE < len(spots_to_embed):
                await asyncio.sleep(RATE_LIMIT_DELAY)

        # 5. 고아 임베딩 정리
        active_ids = {s["id"] for s in spots}
        stats["deleted_orphans"] = await asyncio.to_thread(
            _delete_orphan_embeddings, active_ids
        )

    except Exception as e:
        logger.error(f"[embedding] generate_all_embeddings failed: {e}")
        stats["errors"] += 1

    stats["duration_seconds"] = round(time.monotonic() - start, 2)
    logger.info(f"[embedding] Complete: {stats}")
    return stats


async def update_changed_embeddings() -> dict[str, Any]:
    """
    변경된 스팟만 임베딩 업데이트.
    content_hash 비교로 실제 변경분만 처리.

    generate_all_embeddings(force=False)의 편의 래퍼.
    """
    return await generate_all_embeddings(force=False)


# ============ CLI Entry Point ============

async def _main():
    """CLI 실행용 메인 함수."""
    import argparse
    import sys

    # backend/ 디렉토리를 Python path에 추가
    from pathlib import Path
    backend_dir = str(Path(__file__).parent.parent)
    if backend_dir not in sys.path:
        sys.path.insert(0, backend_dir)

    from dotenv import load_dotenv
    load_dotenv()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )

    parser = argparse.ArgumentParser(description="Generate spot embeddings")
    parser.add_argument(
        "--force", action="store_true",
        help="Force regenerate all embeddings (ignore content_hash)"
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help="Show what would be embedded without actually calling API"
    )
    args = parser.parse_args()

    if args.dry_run:
        spots = _load_active_spots()
        existing = _load_existing_embeddings()
        changed = 0
        new = 0
        for spot in spots:
            h = compute_content_hash(spot)
            if spot["id"] not in existing:
                new += 1
                print(f"  NEW: {spot['id']} ({spot['name']})")
            elif existing[spot["id"]] != h:
                changed += 1
                print(f"  CHANGED: {spot['id']} ({spot['name']})")

        print(f"\nTotal: {len(spots)} active spots")
        print(f"  New: {new}")
        print(f"  Changed: {changed}")
        print(f"  Unchanged: {len(spots) - new - changed}")
        return

    stats = await generate_all_embeddings(force=args.force)
    print(f"\nEmbedding generation complete:")
    for k, v in stats.items():
        print(f"  {k}: {v}")


if __name__ == "__main__":
    asyncio.run(_main())