embeddings.py

"""
Embedding Service for RAG vector operations.

This module handles:
- Generating embeddings using OpenAI's text-embedding-3-small
- Indexing generated assets and events in Supabase (pgvector columns)
- Retrieving similar content via cosine similarity (computed in Python)
"""

import json
import logging
import time
from typing import List, Dict, Any, Optional, Union
import numpy as np

from openai import OpenAI
from app.core.config import settings
from app.db.supabase_client import supabase

logger = logging.getLogger(__name__)

# Embedding model configuration
EMBEDDING_MODEL = "text-embedding-3-small"
EMBEDDING_DIMENSIONS = 1536


def parse_embedding(embedding: Union[str, List[float], None]) -> Optional[List[float]]:
    """
    Parse embedding from database which may be stored as string or list.

    Supabase/pgvector can return embeddings in different formats:
    - List[float]: Already parsed
    - str: JSON string representation of the array
    - str: PostgreSQL array format like "[0.1,0.2,...]"

    Returns None if parsing fails.
    """
    if embedding is None:
        return None

    if isinstance(embedding, list):
        # Already a list, ensure all elements are floats
        try:
            return [float(x) for x in embedding]
        except (ValueError, TypeError):
            logger.warning("Failed to convert list elements to floats")
            return None

    if isinstance(embedding, str):
        # Try to parse as JSON
        try:
            parsed = json.loads(embedding)
            if isinstance(parsed, list):
                return [float(x) for x in parsed]
        except (json.JSONDecodeError, ValueError, TypeError):
            pass

        # Try PostgreSQL array format (starts with [ or {)
        try:
            cleaned = embedding.strip()
            if cleaned.startswith('{') and cleaned.endswith('}'):
                # PostgreSQL array format: {0.1,0.2,...}
                cleaned = cleaned[1:-1]  # Remove braces
                return [float(x) for x in cleaned.split(',')]
            elif cleaned.startswith('[') and cleaned.endswith(']'):
                # JSON-like format without being valid JSON
                cleaned = cleaned[1:-1]
                return [float(x) for x in cleaned.split(',')]
        except (ValueError, TypeError) as e:
            logger.warning(f"Failed to parse embedding string: {e}")

    return None


def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
    """Compute cosine similarity between two vectors."""
    a = np.array(vec1, dtype=np.float64)
    b = np.array(vec2, dtype=np.float64)

    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)

    if norm_a == 0 or norm_b == 0:
        return 0.0

    return float(np.dot(a, b) / (norm_a * norm_b))


class EmbeddingService:
    """Service for generating and managing vector embeddings."""

    def __init__(self):
        # Always use OpenAI for embeddings (most reliable, cost-effective)
        self.client = OpenAI(api_key=settings.OPENAI_API_KEY)
        self.model = EMBEDDING_MODEL

    def embed_text(self, text: str, retries: int = 2) -> List[float]:
        """
        Generate embedding vector for text.

        Args:
            text: The text to embed
            retries: Number of retry attempts for API errors

        Returns:
            List of floats representing the embedding vector (1536 dimensions)
        """
        last_error = None
        for attempt in range(retries + 1):
            try:
                response = self.client.embeddings.create(
                    model=self.model,
                    input=text,
                )
                return response.data[0].embedding
            except Exception as e:
                last_error = e
                error_str = str(e)

                # Check if it's a retryable error (500, 503, rate limit)
                if any(code in error_str for code in ['500', '503', '529', 'rate_limit']):
                    if attempt < retries:
                        wait_time = (attempt + 1) * 2  # Exponential backoff
                        logger.warning(f"OpenAI API error (attempt {attempt + 1}/{retries + 1}), retrying in {wait_time}s: {e}")
                        time.sleep(wait_time)
                        continue

                logger.error(f"Error generating embedding: {e}")
                raise

        logger.error(f"Error generating embedding after {retries + 1} attempts: {last_error}")
        raise last_error

    def embed_texts(self, texts: List[str], retries: int = 2) -> List[List[float]]:
        """
        Generate embeddings for multiple texts in a single API call.

        Args:
            texts: List of texts to embed
            retries: Number of retry attempts for API errors

        Returns:
            List of embedding vectors
        """
        last_error = None
        for attempt in range(retries + 1):
            try:
                response = self.client.embeddings.create(
                    model=self.model,
                    input=texts,
                )
                return [item.embedding for item in response.data]
            except Exception as e:
                last_error = e
                error_str = str(e)

                # Check if it's a retryable error (500, 503, rate limit)
                if any(code in error_str for code in ['500', '503', '529', 'rate_limit']):
                    if attempt < retries:
                        wait_time = (attempt + 1) * 2  # Exponential backoff
                        logger.warning(f"OpenAI API error (attempt {attempt + 1}/{retries + 1}), retrying in {wait_time}s: {e}")
                        time.sleep(wait_time)
                        continue

                logger.error(f"Error generating batch embeddings: {e}")
                raise

        logger.error(f"Error generating batch embeddings after {retries + 1} attempts: {last_error}")
        raise last_error

    def index_generated_asset(self, asset_id: str, content: str) -> bool:
        """
        Generate embedding for a generated asset and store in database.

        Args:
            asset_id: UUID of the generated_assets row
            content: The content text to embed

        Returns:
            True if successful, False otherwise
        """
        try:
            embedding = self.embed_text(content)

            # Update the generated_assets row with embedding
            result = supabase.table("generated_assets").update({
                "content_embedding": embedding
            }).eq("id", asset_id).execute()

            logger.info(f"Indexed generated asset {asset_id}")
            return True
        except Exception as e:
            logger.error(f"Error indexing asset {asset_id}: {e}")
            return False

    def index_event(self, event_id: str, title: str, description: str, audience: str = "") -> bool:
        """
        Generate embedding for an event and store in database.

        Combines title, description, and audience into a single text for embedding.

        Args:
            event_id: UUID of the events row
            title: Event title
            description: Event description
            audience: Target audience (optional)

        Returns:
            True if successful, False otherwise
        """
        try:
            # Combine event fields for embedding
            event_text = f"{title}\n{description}"
            if audience:
                event_text += f"\nAudience: {audience}"

            embedding = self.embed_text(event_text)

            # Update the events row with embedding
            result = supabase.table("events").update({
                "event_embedding": embedding
            }).eq("id", event_id).execute()

            logger.info(f"Indexed event {event_id}")
            return True
        except Exception as e:
            logger.error(f"Error indexing event {event_id}: {e}")
            return False

    def get_similar_assets(
        self,
        org_id: str,
        channel: str,
        query_text: str,
        limit: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Find similar generated assets using cosine similarity.

        Fetches assets with embeddings, computes similarity in Python,
        and returns the top matches.

        Args:
            org_id: Organization UUID to scope the search
            channel: Content channel (instagram, whatsapp, etc.)
            query_text: Text to find similar content for
            limit: Maximum number of results

        Returns:
            List of similar assets with similarity scores, sorted by similarity desc
        """
        try:
            # Generate embedding for query
            query_embedding = self.embed_text(query_text)

            # Get all events for this org
            events_result = supabase.table("events").select(
                "id, title"
            ).eq("org_id", org_id).execute()

            if not events_result.data:
                return []

            org_event_ids = [e["id"] for e in events_result.data]
            event_titles = {e["id"]: e["title"] for e in events_result.data}

            # Fetch assets for these events with embeddings
            # Supabase doesn't support IN with large lists well, so we filter in Python
            assets_result = supabase.table("generated_assets").select(
                "id, event_id, type, content, feedback, iteration, content_embedding"
            ).eq("type", channel).not_.is_("content_embedding", "null").execute()

            if not assets_result.data:
                return []

            # Filter to org's events and compute similarity
            scored_assets = []
            for asset in assets_result.data:
                if asset["event_id"] not in org_event_ids:
                    continue

                raw_embedding = asset.get("content_embedding")
                embedding = parse_embedding(raw_embedding)
                if not embedding:
                    logger.debug(f"Skipping asset {asset['id']} - could not parse embedding")
                    continue

                try:
                    similarity = cosine_similarity(query_embedding, embedding)
                except Exception as e:
                    logger.warning(f"Error computing similarity for asset {asset['id']}: {e}")
                    continue

                scored_assets.append({
                    "id": asset["id"],
                    "event_id": asset["event_id"],
                    "event_title": event_titles.get(asset["event_id"], ""),
                    "type": asset["type"],
                    "content": asset["content"],
                    "feedback": asset.get("feedback"),
                    "iteration": asset["iteration"],
                    "similarity": similarity
                })

            # Sort by similarity descending and return top matches
            scored_assets.sort(key=lambda x: x["similarity"], reverse=True)
            return scored_assets[:limit]

        except Exception as e:
            logger.error(f"Error finding similar assets: {e}")
            return []

    def get_similar_events(
        self,
        org_id: str,
        query_text: str,
        limit: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Find similar events using cosine similarity.

        Args:
            org_id: Organization UUID to scope the search
            query_text: Text to find similar events for
            limit: Maximum number of results

        Returns:
            List of similar events with similarity scores
        """
        try:
            query_embedding = self.embed_text(query_text)

            # Fetch events for this org with embeddings
            events_result = supabase.table("events").select(
                "id, title, description, audience, tone, event_embedding"
            ).eq("org_id", org_id).not_.is_("event_embedding", "null").execute()

            if not events_result.data:
                return []

            # Compute similarity for each event
            scored_events = []
            for event in events_result.data:
                raw_embedding = event.get("event_embedding")
                embedding = parse_embedding(raw_embedding)
                if not embedding:
                    logger.debug(f"Skipping event {event['id']} - could not parse embedding")
                    continue

                try:
                    similarity = cosine_similarity(query_embedding, embedding)
                except Exception as e:
                    logger.warning(f"Error computing similarity for event {event['id']}: {e}")
                    continue

                scored_events.append({
                    "id": event["id"],
                    "title": event["title"],
                    "description": event["description"],
                    "audience": event.get("audience"),
                    "tone": event.get("tone"),
                    "similarity": similarity
                })

            # Sort by similarity descending
            scored_events.sort(key=lambda x: x["similarity"], reverse=True)
            return scored_events[:limit]

        except Exception as e:
            logger.error(f"Error finding similar events: {e}")
            return []


# Global instance
embedding_service = EmbeddingService()