src/services.py

"""Backend services for AION Search."""

import time
import logging
from pathlib import Path
import torch
import torch.nn.functional as F
import numpy as np
import pandas as pd
import requests
from typing import List, Optional
from openai import OpenAI

from src.config import (
    AION_PCA96_MODEL_PATH,
    ZILLIZ_BEARER,
    ZILLIZ_ENDPOINT,
    ZILLIZ_COLLECTION_NAME,
    ZILLIZ_IMAGE_SEARCH_COLLECTION_NAME,
    ZILLIZ_ANNS_FIELD,
    ZILLIZ_PRIMARY_KEY,
    ZILLIZ_OUTPUT_FIELDS,
    ZILLIZ_VECTOR_DIM,
    ZILLIZ_QUERY_PROJECTION,
    COLLECTION_CONFIGS,
    OPENAI_API_KEY,
    OPENAI_EMBEDDING_MODEL,
    CLIP_NORMALIZE_EPS,
    DEFAULT_TOP_K,
)
from src.utils import cutout_url, log_zilliz_query

logger = logging.getLogger(__name__)


class CLIPModelService:
    """Service for managing CLIP model loading and inference."""

    def __init__(self):
        self.model = None
        self.device = None
        self.loaded = False

    def load_model(self, checkpoint_path: str) -> None:
        """Load the CLIP model from checkpoint.

        Args:
            checkpoint_path: Path to the CLIP model checkpoint file
        """
        logger.info(f"Loading CLIP model from {checkpoint_path}...")

        from clip.models.clip_model import GalaxyClipModel

        # Set device
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # Load checkpoint
        checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
        model_config = checkpoint['model_config']

        # Initialize model with saved configuration
        self.model = GalaxyClipModel(
            image_input_dim=model_config['image_input_dim'],
            text_input_dim=model_config['text_input_dim'],
            embedding_dim=model_config['embedding_dim'],
            use_mean_embeddings=model_config.get('use_mean_embeddings', True)
        )

        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.model.to(self.device)
        self.model.eval()
        self.loaded = True

        logger.info("CLIP model loaded successfully")

    def encode_text(self, text_embedding: np.ndarray) -> np.ndarray:
        """Project text embedding through CLIP text projector.

        Args:
            text_embedding: OpenAI text embedding (1536-dim)

        Returns:
            CLIP-projected embedding (1024-dim)
        """
        if not self.loaded:
            raise RuntimeError("CLIP model not loaded. Call load_model() first.")

        with torch.no_grad():
            text_tensor = torch.from_numpy(text_embedding).float().unsqueeze(0).to(self.device)
            clip_features = self.model.text_projector(text_tensor)
            # Normalize as per CLIP
            clip_features = F.normalize(clip_features, dim=-1, eps=CLIP_NORMALIZE_EPS)
            query_embedding = clip_features.cpu().numpy().squeeze(0)

        return query_embedding


class AIONPCA96Projection:
    """Project 1024-d AION-Search vectors into the 96-d PCA search space."""

    def __init__(self, model_path: str):
        self.model_path = Path(model_path)
        if not self.model_path.exists():
            raise FileNotFoundError(f"PCA model not found: {self.model_path}")

        model = np.load(self.model_path)
        self.mean = np.asarray(model["mean"], dtype=np.float64)
        self.components = np.asarray(model["components"], dtype=np.float64)

        if self.mean.shape != (1024,):
            raise ValueError(f"Expected PCA mean shape (1024,), got {self.mean.shape}")
        if self.components.shape != (96, 1024):
            raise ValueError(f"Expected PCA components shape (96, 1024), got {self.components.shape}")

        logger.info(f"Loaded AION PCA96 projection from {self.model_path}")

    def project(self, embedding: np.ndarray) -> np.ndarray:
        """Apply the trained PCA projection and L2-normalize the output."""
        x = np.asarray(embedding, dtype=np.float64)
        squeeze = False

        if x.ndim == 1:
            x = x[None, :]
            squeeze = True

        if x.ndim != 2 or x.shape[-1] != self.mean.shape[0]:
            raise ValueError(f"Expected embedding shape (..., 1024), got {x.shape}")

        y = (x - self.mean) @ self.components.T
        y = y / np.maximum(np.linalg.norm(y, axis=-1, keepdims=True), 1e-12)
        y = y.astype(np.float32)

        if squeeze:
            return y[0]
        return y


def create_query_projection() -> Optional[AIONPCA96Projection]:
    """Create the configured query projection, if the active collection needs one."""
    if ZILLIZ_QUERY_PROJECTION is None:
        return None

    if ZILLIZ_QUERY_PROJECTION == "aion_pca96":
        return AIONPCA96Projection(AION_PCA96_MODEL_PATH)

    raise ValueError(f"Unknown query projection: {ZILLIZ_QUERY_PROJECTION}")


class ImageProcessingService:
    """Service for retrieving pre-existing image embeddings from Zilliz."""

    def __init__(self):
        pass

    def encode_image(self, ra: float, dec: float, fov: float = 0.025, size: int = 256) -> np.ndarray:
        """Query Zilliz for pre-existing embedding at the given coordinates.

        Args:
            ra: Right ascension in degrees
            dec: Declination in degrees
            fov: Field of view in degrees (used to define search box)
            size: Image size in pixels (unused, kept for API compatibility)

        Returns:
            Pre-existing AION-Search embedding vector from Zilliz
        """
        logger.info(f"Querying Zilliz for pre-existing embedding at RA={ra}, Dec={dec}, FoV={fov}")

        # Calculate bounding box based on field of view
        ra_min = ra - fov/2
        ra_max = ra + fov/2
        dec_min = dec - fov/2
        dec_max = dec + fov/2

        # Build filter expression for coordinate range
        filter_expr = f"ra > {ra_min} AND ra < {ra_max} AND dec > {dec_min} AND dec < {dec_max}"

        # Get the ANNS field for the image search collection
        image_search_config = COLLECTION_CONFIGS.get(ZILLIZ_IMAGE_SEARCH_COLLECTION_NAME)
        image_anns_field = image_search_config["anns_field"]

        # Prepare query payload from the active image-search collection.
        payload = {
            "collectionName": ZILLIZ_IMAGE_SEARCH_COLLECTION_NAME,
            "filter": filter_expr,
            "outputFields": [image_anns_field],
            "limit": 1
        }

        headers = {
            "Authorization": f"Bearer {ZILLIZ_BEARER}",
            "Accept": "application/json",
            "Content-Type": "application/json"
        }

        try:
            # Use query endpoint (replace /search with /query)
            query_endpoint = ZILLIZ_ENDPOINT.replace("/search", "/query")
            response = requests.post(query_endpoint, json=payload, headers=headers)
            response.raise_for_status()

            result = response.json()

            if result.get("code") == 0 and "data" in result:
                data = result["data"]
                if data and len(data) > 0:
                    # Extract the embedding from the first result using the image search ANNS field
                    embedding = data[0].get(image_anns_field)
                    if embedding:
                        embedding_array = np.array(embedding, dtype=np.float32)
                        logger.info(f"Retrieved pre-existing embedding with shape: {embedding_array.shape}")
                        return embedding_array
                    else:
                        logger.error(f"No embedding field found in result: {data[0].keys()}")
                        raise RuntimeError(f"No embedding found at coordinates RA={ra}, Dec={dec}")
                else:
                    logger.error(f"No galaxies found at coordinates RA={ra}, Dec={dec} with FoV={fov}")
                    raise RuntimeError(f"No galaxies found at coordinates RA={ra}, Dec={dec}")
            else:
                logger.error(f"Zilliz query failed: {result}")
                raise RuntimeError(f"Failed to query Zilliz: {result}")

        except Exception as e:
            logger.error(f"Error querying Zilliz for embedding: {e}")
            raise


class EmbeddingService:
    """Service for encoding text queries into embeddings."""

    def __init__(
        self,
        clip_service: CLIPModelService,
        query_projection: Optional[AIONPCA96Projection] = None,
    ):
        self.clip_service = clip_service
        self.query_projection = query_projection
        self.openai_client = None

    def _project_for_search(self, embedding: np.ndarray) -> np.ndarray:
        """Project CLIP-space embeddings into the active Zilliz collection space."""
        if self.query_projection is None:
            return np.asarray(embedding, dtype=np.float32)
        return self.query_projection.project(embedding)

    def _encode_text_clip_query(self, query: str) -> np.ndarray:
        """Encode text into the original 1024-d CLIP search space."""
        if not self._moderate_content(query):
            raise ValueError("Content moderation filter triggered")

        client = self._get_openai_client()
        response = client.embeddings.create(
            input=query,
            model=OPENAI_EMBEDDING_MODEL
        )
        text_embedding = np.array(response.data[0].embedding)

        return self.clip_service.encode_text(text_embedding)

    def _get_openai_client(self) -> OpenAI:
        """Get or create OpenAI client."""
        if self.openai_client is None:
            if not OPENAI_API_KEY:
                raise ValueError("OPENAI_API_KEY environment variable not set")
            self.openai_client = OpenAI(api_key=OPENAI_API_KEY)
        return self.openai_client

    def _moderate_content(self, text: str) -> bool:
        """Check if text content is appropriate using OpenAI Moderation API.

        Args:
            text: Text to moderate

        Returns:
            True if content is safe, False if flagged
        """
        try:
            client = self._get_openai_client()
            response = client.moderations.create(input=text)

            # If any category is flagged, reject the content
            if response.results[0].flagged:
                logger.warning(f"Content moderation flagged input")
                return False
            return True
        except Exception as e:
            logger.error(f"Moderation API error: {e}")
            # On error, allow the content through (fail open)
            return True

    def encode_text_query(self, query: str) -> np.ndarray:
        """Encode text query using OpenAI embeddings + CLIP text projector.

        Args:
            query: Text search query

        Returns:
            CLIP embedding vector
        """
        clip_embedding = self._encode_text_clip_query(query)
        return self._project_for_search(clip_embedding)

    def encode_vector_queries(
        self,
        queries: List[str],
        operations: List[str]
    ) -> np.ndarray:
        """Encode multiple text queries and combine them using vector addition/subtraction.

        Args:
            queries: List of text queries
            operations: List of operations ('+' or '-') for each query

        Returns:
            Combined normalized embedding vector
        """
        weights = [1.0 if operation == "+" else -1.0 for operation in operations]
        return self.encode_weighted_text_queries(queries, weights)

    def encode_weighted_text_queries(
        self,
        queries: List[str],
        weights: List[float]
    ) -> np.ndarray:
        """Encode, combine, and project text queries.

        Text-only vector arithmetic stays in the original 1024-d CLIP space until
        the final combined vector is normalized, then projects once into the
        active Zilliz search space.
        """
        for query in queries:
            if not self._moderate_content(query):
                raise ValueError("Content moderation filter triggered")

        client = self._get_openai_client()
        response = client.embeddings.create(
            input=queries,
            model=OPENAI_EMBEDDING_MODEL
        )

        combined_embedding = None
        for embedding_data, weight in zip(response.data, weights):
            text_embedding = np.array(embedding_data.embedding)
            query_embedding = self.clip_service.encode_text(text_embedding)
            weighted_embedding = query_embedding * weight

            if combined_embedding is None:
                combined_embedding = weighted_embedding
            else:
                combined_embedding += weighted_embedding

        norm = np.linalg.norm(combined_embedding)
        if norm > 0:
            combined_embedding = combined_embedding / norm

        return self._project_for_search(combined_embedding)


class ZillizService:
    """Service for interacting with Zilliz vector database."""

    def get_collection_count(self) -> int:
        """Get the total number of entities in the collection.

        Returns:
            Total count of entities in the collection
        """
        logger.info("Getting collection count from Zilliz...")

        # Use query endpoint with count to get total entities
        payload = {
            "collectionName": ZILLIZ_COLLECTION_NAME,
            "filter": "",  # Empty filter to count all entities
            "outputFields": ["count(*)"]
        }

        headers = {
            "Authorization": f"Bearer {ZILLIZ_BEARER}",
            "Accept": "application/json",
            "Content-Type": "application/json"
        }

        try:
            # Use the query endpoint (replace /search with /query in the endpoint)
            query_endpoint = ZILLIZ_ENDPOINT.replace("/search", "/query")
            response = requests.post(query_endpoint, json=payload, headers=headers)
            response.raise_for_status()

            result = response.json()

            if result.get("code") == 0 and "data" in result:
                # The count should be in the response data
                data = result["data"]
                if data and len(data) > 0:
                    count = data[0].get("count(*)", 0)
                    logger.info(f"Collection count: {count:,}")
                    return count
            else:
                logger.error(f"Failed to get collection count: {result}")
                return 0

        except Exception as e:
            logger.error(f"Error getting collection count: {e}")
            return 0

    def search(self, query_embedding: np.ndarray, top_k: int = DEFAULT_TOP_K, filter_expr: str = None) -> pd.DataFrame:
        """Search Zilliz for top-k most similar galaxies.

        Args:
            query_embedding: Query embedding vector
            top_k: Number of results to return
            filter_expr: Optional filter expression for filtering results

        Returns:
            DataFrame with search results
        """
        query_embedding = np.asarray(query_embedding, dtype=np.float32)
        if query_embedding.ndim != 1:
            raise ValueError(f"Expected one query vector, got shape {query_embedding.shape}")
        if len(query_embedding) != ZILLIZ_VECTOR_DIM:
            raise ValueError(
                f"Query vector dimension {len(query_embedding)} does not match "
                f"{ZILLIZ_COLLECTION_NAME}.{ZILLIZ_ANNS_FIELD} dimension {ZILLIZ_VECTOR_DIM}"
            )

        logger.info(
            f"Querying Zilliz collection {ZILLIZ_COLLECTION_NAME} "
            f"({ZILLIZ_ANNS_FIELD}, {len(query_embedding)}-dim)..."
        )
        start_time = time.time()

        # Prepare the search payload
        payload = {
            "collectionName": ZILLIZ_COLLECTION_NAME,
            "data": [query_embedding.tolist()],
            "annsField": ZILLIZ_ANNS_FIELD,
            "limit": top_k,
            "outputFields": ZILLIZ_OUTPUT_FIELDS
        }

        # Add filter if provided
        if filter_expr:
            payload["filter"] = filter_expr
            logger.info(f"Applying filter: {filter_expr}")

        headers = {
            "Authorization": f"Bearer {ZILLIZ_BEARER}",
            "Accept": "application/json",
            "Content-Type": "application/json"
        }

        try:
            response = requests.post(ZILLIZ_ENDPOINT, json=payload, headers=headers)
            response.raise_for_status()

            result = response.json()

            if result.get("code") == 0 and "data" in result:
                # Extract cost from response
                cost_vcu = result.get("cost", 0)

                # Convert to DataFrame
                data_list = result["data"]
                df = pd.DataFrame(data_list)

                # Add cutout URLs
                if not df.empty:
                    df["cutout_url"] = [cutout_url(ra, dec) for ra, dec in zip(df["ra"], df["dec"])]

                query_time = time.time() - start_time

                # Log the query
                log_zilliz_query(
                    query_type="vector_search",
                    query_info={
                        "top_k": top_k,
                        "embedding_dim": len(query_embedding)
                    },
                    result_count=len(df),
                    query_time=query_time,
                    cost_vcu=cost_vcu
                )

                return df
            else:
                logger.error(f"Zilliz search failed: {result}")
                return pd.DataFrame()

        except Exception as e:
            logger.error(f"Zilliz search error: {e}")
            return pd.DataFrame()


class SearchService:
    """High-level search orchestration service."""

    def __init__(
        self,
        embedding_service: EmbeddingService,
        zilliz_service: ZillizService,
        image_service: 'ImageProcessingService' = None
    ):
        self.embedding_service = embedding_service
        self.zilliz_service = zilliz_service
        self.image_service = image_service

    def _build_rmag_filter(self, rmag_min=None, rmag_max=None) -> str:
        """Build r_mag filter expression.

        Args:
            rmag_min: Minimum r_mag value (inclusive)
            rmag_max: Maximum r_mag value (inclusive)

        Returns:
            Filter expression string, or None if no filter
        """
        filter_parts = []

        if rmag_min is not None:
            filter_parts.append(f"r_mag >= {rmag_min}")

        if rmag_max is not None:
            filter_parts.append(f"r_mag <= {rmag_max}")

        if filter_parts:
            return " AND ".join(filter_parts)

        return None

    def search_text(self, query: str, top_k: int = DEFAULT_TOP_K, rmag_min=None, rmag_max=None) -> pd.DataFrame:
        """Search galaxies using text query.

        Args:
            query: Text search query
            top_k: Number of results to return
            rmag_min: Minimum r_mag value (inclusive)
            rmag_max: Maximum r_mag value (inclusive)

        Returns:
            DataFrame with search results
        """
        try:
            # Encode query
            query_embedding = self.embedding_service.encode_text_query(query)

            # Build filter
            filter_expr = self._build_rmag_filter(rmag_min, rmag_max)

            # Search Zilliz
            return self.zilliz_service.search(query_embedding, top_k, filter_expr)
        except ValueError as e:
            # Content moderation triggered - return empty results silently
            if "moderation" in str(e).lower():
                logger.info("Search blocked by content moderation")
                return pd.DataFrame()
            raise

    def search_vector(
        self,
        queries: List[str],
        operations: List[str],
        top_k: int = DEFAULT_TOP_K,
        rmag_min=None,
        rmag_max=None
    ) -> pd.DataFrame:
        """Search galaxies using vector addition/subtraction.

        Args:
            queries: List of text queries
            operations: List of operations ('+' or '-') for each query
            top_k: Number of results to return
            rmag_min: Minimum r_mag value (inclusive)
            rmag_max: Maximum r_mag value (inclusive)

        Returns:
            DataFrame with search results
        """
        try:
            # Encode and combine vectors
            combined_embedding = self.embedding_service.encode_vector_queries(queries, operations)

            # Build filter
            filter_expr = self._build_rmag_filter(rmag_min, rmag_max)

            # Search Zilliz
            return self.zilliz_service.search(combined_embedding, top_k, filter_expr)
        except ValueError as e:
            # Content moderation triggered - return empty results silently
            if "moderation" in str(e).lower():
                logger.info("Search blocked by content moderation")
                return pd.DataFrame()
            raise

    def search_advanced(
        self,
        text_queries: List[str] = None,
        text_weights: List[float] = None,
        image_queries: List[dict] = None,
        image_weights: List[float] = None,
        top_k: int = DEFAULT_TOP_K,
        rmag_min=None,
        rmag_max=None
    ) -> pd.DataFrame:
        """Search galaxies using advanced vector addition/subtraction with text and/or images.

        Args:
            text_queries: List of text query strings
            text_weights: List of weight magnitudes for text queries (e.g., 1.0, -1.0, 2.0, -5.0)
            image_queries: List of dicts with 'ra', 'dec', 'fov' keys
            image_weights: List of weight magnitudes for image queries (e.g., 1.0, -1.0, 2.0, -5.0)
            top_k: Number of results to return
            rmag_min: Minimum r_mag value (inclusive)
            rmag_max: Maximum r_mag value (inclusive)

        Returns:
            DataFrame with search results
        """
        try:
            combined_embedding = None
            has_image_queries = bool(image_queries and len(image_queries) > 0)

            if text_queries and len(text_queries) > 0 and not has_image_queries:
                combined_embedding = self.embedding_service.encode_weighted_text_queries(
                    text_queries,
                    text_weights
                )
            elif text_queries and len(text_queries) > 0:
                for query, weight in zip(text_queries, text_weights):
                    query_embedding = self.embedding_service.encode_text_query(query)

                    # Apply weight
                    weighted_embedding = query_embedding * weight

                    if combined_embedding is None:
                        combined_embedding = weighted_embedding
                    else:
                        combined_embedding += weighted_embedding

            # Process image queries. Once legacy 1024-d vectors are removed, mixed
            # text+image arithmetic must happen in the active 96-d search space.
            if has_image_queries:
                if self.image_service is None:
                    raise RuntimeError("Image service not initialized")

                for img_query, weight in zip(image_queries, image_weights):
                    # Encode image
                    image_embedding = self.image_service.encode_image(
                        ra=img_query['ra'],
                        dec=img_query['dec'],
                        fov=img_query.get('fov', 0.025),
                        size=256
                    )

                    # Apply weight
                    weighted_embedding = image_embedding * weight

                    if combined_embedding is None:
                        combined_embedding = weighted_embedding
                    else:
                        combined_embedding += weighted_embedding

            # Normalize the final combined embedding
            if combined_embedding is not None:
                norm = np.linalg.norm(combined_embedding)
                if norm > 0:
                    combined_embedding = combined_embedding / norm

            # Build filter
            filter_expr = self._build_rmag_filter(rmag_min, rmag_max)

            # Search Zilliz
            return self.zilliz_service.search(combined_embedding, top_k, filter_expr)
        except ValueError as e:
            # Content moderation triggered - return empty results silently
            if "moderation" in str(e).lower():
                logger.info("Search blocked by content moderation")
                return pd.DataFrame()
            raise