"""
cortex/llm/google_embeddings.py — Google/Gemini embedding provider.

Provides embedding generation using Google's Gemini embedding models via
the OpenAI-compatible endpoint. Supports gemini-embedding-001 (3072-dim)
and gemini-embedding-2-preview (768/1536/3072-dim configurable).

Auth: GOOGLE_API_KEY or GEMINI_API_KEY environment variable, or explicit
api_key in embedding_config.

Endpoint: https://generativelanguage.googleapis.com/v1beta/openai/embeddings

Design decisions:
    - Uses the OpenAI SDK against Google's OpenAI-compatible endpoint
      (simpler than using the google-genai SDK, same wire format)
    - Batch limit: 100 texts per request (Google's limit)
    - Rate limiting: 0.5s delay between batches to stay under quota
    - Configurable output_dimensionality for gemini-embedding-2-preview

Dependencies: openai SDK.
Implements: standalone embed() that matches HostProvider._embed_google() interface.
"""

from __future__ import annotations

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

from cortex.llm.provider import (
    EmbedResponse,
    UsageRecord,
    retry_with_backoff,
)

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

GOOGLE_EMBED_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"

GOOGLE_EMBED_MODELS = {
    "gemini-embedding-001",
    "gemini-embedding-2-preview",
}

# Default dimensions per model
GOOGLE_MODEL_DEFAULT_DIM: dict[str, int] = {
    "gemini-embedding-001": 3072,
    "gemini-embedding-2-preview": 3072,
}

# Valid dimensions per model
GOOGLE_MODEL_VALID_DIMS: dict[str, set[int]] = {
    "gemini-embedding-001": {3072},
    "gemini-embedding-2-preview": {768, 1536, 3072},
}

# Pricing: $/1M tokens (free tier currently, but track for cost estimation)
GOOGLE_EMBED_PRICING: dict[str, float] = {
    "gemini-embedding-001": 0.0,
    "gemini-embedding-2-preview": 0.0,
}

# Max texts per batch request
GOOGLE_BATCH_LIMIT = 100

# Delay between batches (seconds) to respect rate limits
GOOGLE_BATCH_DELAY = 0.5


def _estimate_google_embed_cost(model: str, tokens: int) -> float:
    """Estimate cost in USD for a Google embedding call."""
    price = GOOGLE_EMBED_PRICING.get(model, 0.0)
    return tokens * price / 1_000_000


def is_google_embed_model(model: str) -> bool:
    """Return True if the model is a Google/Gemini embedding model."""
    return model in GOOGLE_EMBED_MODELS or model.startswith("gemini-embedding")


def get_google_api_key(embedding_config: Optional[dict[str, Any]] = None) -> str:
    """Resolve Google API key from config or environment.

    Priority: embedding_config["api_key"] > GOOGLE_API_KEY > GEMINI_API_KEY.
    """
    if embedding_config:
        key = embedding_config.get("api_key", "")
        if key:
            return key
    return os.environ.get("GOOGLE_API_KEY", "") or os.environ.get("GEMINI_API_KEY", "")


def _get_google_embed_client(api_key: str) -> Any:
    """Create an AsyncOpenAI client pointing at Google's OpenAI-compatible endpoint."""
    if not api_key:
        raise RuntimeError(
            "No Google API key provided. Set GOOGLE_API_KEY or GEMINI_API_KEY "
            "environment variable, or pass api_key in embedding config."
        )
    try:
        import openai
    except ImportError:
        raise RuntimeError("openai package not installed. Run: pip install openai")
    return openai.AsyncOpenAI(
        api_key=api_key,
        base_url=GOOGLE_EMBED_BASE_URL,
    )


async def embed_google(
    texts: list[str],
    model: str,
    api_key: str,
    dim: Optional[int] = None,
    _client: Any = None,
) -> EmbedResponse:
    """Generate embeddings using Google's Gemini embedding models.

    Handles batching (up to 100 texts per request) and rate limiting
    (0.5s delay between batches).

    Args:
        texts: List of strings to embed.
        model: Gemini embedding model ID.
        api_key: Google API key.
        dim: Output dimensionality (optional, uses model default if not set).
        _client: Pre-created AsyncOpenAI client (for testing/reuse).

    Returns:
        EmbedResponse with embeddings + cost metrics.
    """
    client = _client or _get_google_embed_client(api_key)

    all_embeddings: list[list[float]] = []
    total_tokens = 0
    t0 = time.perf_counter()

    # Process in batches of GOOGLE_BATCH_LIMIT
    for i in range(0, len(texts), GOOGLE_BATCH_LIMIT):
        if i > 0:
            await asyncio.sleep(GOOGLE_BATCH_DELAY)

        batch = texts[i : i + GOOGLE_BATCH_LIMIT]

        kwargs: dict[str, Any] = {
            "model": model,
            "input": batch,
            "encoding_format": "float",
        }
        # Add dimensions parameter if specified
        if dim is not None:
            kwargs["dimensions"] = dim

        response = await retry_with_backoff(
            lambda kw=kwargs: client.embeddings.create(**kw),
            operation_name=f"google.embed({model})",
        )

        batch_embeddings = [item.embedding for item in response.data]
        all_embeddings.extend(batch_embeddings)
        if response.usage and hasattr(response.usage, "prompt_tokens"):
            total_tokens += response.usage.prompt_tokens or 0
        else:
            # Gemini embed API may not return usage; estimate from input
            total_tokens += sum(len(t.split()) for t in batch)

    latency_ms = (time.perf_counter() - t0) * 1000
    cost = _estimate_google_embed_cost(model, total_tokens)

    return EmbedResponse(
        embeddings=all_embeddings,
        model=model,
        tokens_in=total_tokens,
        tokens_out=0,
        cost_estimate=cost,
        latency_ms=latency_ms,
    )
