"""
cortex/llm/google_provider.py — Google/Gemini Implementation of LLMProvider.

Google's Gemini API offers an OpenAI-compatible chat completions endpoint,
allowing us to reuse the openai SDK with a custom base URL. This provider
handles LLM completions (extraction, reconciliation, synthesis) — embedding
is handled separately by google_embeddings.py.

API keys / endpoints (from environment):
    GOOGLE_API_KEY   — Required. Also accepts GEMINI_API_KEY as fallback.
    GOOGLE_BASE_URL  — Optional override. Default:
                       https://generativelanguage.googleapis.com/v1beta/openai/

Supported models (as of 2026):
    Completion: gemini-2.5-flash (default), gemini-3-flash-preview,
                gemini-2.5-flash-lite, gemini-3-pro-preview, gemini-3.1-pro-preview
    Embedding: NOT SUPPORTED here — use google_embeddings.py instead.

Key implementation details:
    - Uses openai SDK with custom base_url pointing at Google's OpenAI-compat endpoint
    - response_format=json_object IS supported by Gemini models
    - Retry on 429/503 via retry_with_backoff (inherited behavior)
    - Cost estimation uses Google's published pricing

References:
    https://ai.google.dev/gemini-api/docs/openai

Dependencies: openai SDK (same as OpenAIProvider — no new requirements).
Implements: cortex.llm.provider.LLMProvider ABC.
"""

from __future__ import annotations

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

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

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Default configuration
# ---------------------------------------------------------------------------

_DEFAULT_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
_DEFAULT_COMPLETION_MODEL = "gemini-2.5-flash"

# ---------------------------------------------------------------------------
# Pricing ($/1M tokens — Google AI Studio pricing as of 2026)
# ---------------------------------------------------------------------------
_GOOGLE_COMPLETION_PRICING: dict[str, tuple[float, float]] = {
    # model_key: (input $/1M, output $/1M)
    "gemini-2.5-flash-lite": (0.0, 0.0),       # free tier
    "gemini-2.5-flash": (0.15, 0.60),           # pay-as-you-go
    "gemini-3-flash-preview": (0.15, 0.60),     # preview pricing
    "gemini-3-pro-preview": (1.25, 10.0),       # pro tier
    "gemini-3.1-pro-preview": (1.25, 10.0),     # pro tier
}


def _google_completion_cost(model: str, tokens_in: int, tokens_out: int) -> float:
    """Estimate cost for a Google completion call."""
    key = next(
        (k for k in sorted(_GOOGLE_COMPLETION_PRICING, key=len, reverse=True)
         if model.startswith(k)),
        None,
    )
    if key is None:
        return 0.0
    in_p, out_p = _GOOGLE_COMPLETION_PRICING[key]
    return (tokens_in * in_p + tokens_out * out_p) / 1_000_000


def is_google_llm_model(model: str) -> bool:
    """Return True if the model is a Google/Gemini LLM (non-embedding) model."""
    return model.startswith("gemini-") and "embedding" not in model


# ---------------------------------------------------------------------------
# GoogleProvider
# ---------------------------------------------------------------------------


class GoogleProvider(LLMProvider):
    """
    LLMProvider backed by Google's Gemini Chat Completions API.

    Google's Gemini API is OpenAI-compatible. Uses the openai SDK pointed
    at Google's ``/v1beta/openai/`` endpoint.

    The provider reads ``GOOGLE_API_KEY`` (or ``GEMINI_API_KEY``) from the
    environment. If neither is set, ``__init__`` raises ``RuntimeError``
    so the FallbackLLMProvider cascade can skip it gracefully.

    Args:
        completion_model: Override the default completion model.
        embed_model: Ignored (embeddings handled by google_embeddings.py).
    """

    PROVIDER_NAME = "google"

    def __init__(
        self,
        completion_model: Optional[str] = None,
        embed_model: Optional[str] = None,
    ) -> None:
        super().__init__()
        self._api_key = os.environ.get("GOOGLE_API_KEY", "") or os.environ.get(
            "GEMINI_API_KEY", ""
        )
        if not self._api_key:
            raise RuntimeError(
                "GoogleProvider requires GOOGLE_API_KEY or GEMINI_API_KEY "
                "environment variable. Set one and retry."
            )
        self._base_url = os.environ.get("GOOGLE_BASE_URL", _DEFAULT_BASE_URL)
        self._completion_model = completion_model or _DEFAULT_COMPLETION_MODEL
        self._embed_model = embed_model or "gemini-embedding-2-preview"
        self._client: Any = None

    def _get_client(self) -> Any:
        """Get or create AsyncOpenAI client for Google completions."""
        if self._client is None:
            try:
                import openai
            except ImportError:
                raise RuntimeError(
                    "openai package not installed. Run: pip install openai"
                )
            self._client = openai.AsyncOpenAI(
                api_key=self._api_key, base_url=self._base_url
            )
        return self._client

    # --- LLMProvider interface: complete() ----------------------------------

    async def complete(
        self,
        system: str,
        user: str,
        model: Optional[str] = None,
        response_format: str = "text",
        max_tokens: Optional[int] = None,
        temperature: float = 0.0,
    ) -> CompletionResponse:
        """Send a completion request to Gemini via OpenAI-compatible API."""
        resolved_model = model or self._completion_model
        client = self._get_client()

        kwargs: dict[str, Any] = {
            "model": resolved_model,
            "temperature": temperature,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
        }
        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens
        if response_format == "json":
            kwargs["response_format"] = {"type": "json_object"}

        t0 = time.perf_counter()
        response = await retry_with_backoff(
            lambda: client.chat.completions.create(**kwargs),
            operation_name=f"google.complete({resolved_model})",
        )
        latency_ms = (time.perf_counter() - t0) * 1000

        text = response.choices[0].message.content or ""
        tokens_in = getattr(response.usage, "prompt_tokens", 0) or 0
        tokens_out = getattr(response.usage, "completion_tokens", 0) or 0
        cost = _google_completion_cost(resolved_model, tokens_in, tokens_out)

        self._log_usage(
            UsageRecord(
                provider="google",
                model=resolved_model,
                call_type="complete",
                tokens_in=tokens_in,
                tokens_out=tokens_out,
                cost_estimate=cost,
                latency_ms=latency_ms,
            )
        )

        return CompletionResponse(
            text=text,
            model=resolved_model,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            cost_estimate=cost,
            latency_ms=latency_ms,
        )

    # --- LLMProvider interface: extract() -----------------------------------

    async def extract(
        self,
        system: str,
        user: str,
        model: Optional[str] = None,
        max_tokens: Optional[int] = None,
        temperature: float = 0.0,
    ) -> ExtractionResponse:
        """Extract structured data as JSON via Gemini."""
        resolved_model = model or self._completion_model

        # Use JSON mode for extraction
        resp = await self.complete(
            system=system,
            user=user,
            model=resolved_model,
            response_format="json",
            max_tokens=max_tokens,
            temperature=temperature,
        )

        raw_text = resp.text.strip()
        # Attempt JSON parse
        try:
            parsed = json.loads(raw_text)
        except json.JSONDecodeError:
            # Try extracting JSON from markdown code block
            if "```" in raw_text:
                start = raw_text.find("```")
                end = raw_text.rfind("```")
                if start != end:
                    block = raw_text[start:end].split("\n", 1)[-1]
                    try:
                        parsed = json.loads(block)
                    except json.JSONDecodeError:
                        parsed = {"raw": raw_text}
                else:
                    parsed = {"raw": raw_text}
            else:
                parsed = {"raw": raw_text}

        return ExtractionResponse(
            data=parsed,
            raw_text=raw_text,
            model=resp.model,
            tokens_in=resp.tokens_in,
            tokens_out=resp.tokens_out,
            cost_estimate=resp.cost_estimate,
            latency_ms=resp.latency_ms,
        )

    # --- LLMProvider interface: embed() ------------------------------------

    SUPPORTS_EMBED = True

    async def embed(
        self,
        texts: list[str],
        model: Optional[str] = None,
    ) -> EmbedResponse:
        """Generate embeddings via Gemini Embedding models (google_embeddings.py)."""
        from cortex.llm.google_embeddings import embed_google

        embed_model = model or self._embed_model
        return await embed_google(
            texts=texts,
            model=embed_model,
            api_key=self._api_key,
        )
