"""
cortex/llm/openai_provider.py — OpenAI Implementation of LLMProvider.

Wraps OpenAI's Chat Completions API for inference and their Embeddings API
for vector generation. This is the default/recommended provider for Cortex
due to the cost-effectiveness of gpt-4.1-nano for extraction tasks.

API keys / endpoints (from environment):
    OPENAI_API_KEY   — Required for all operations.
    OPENAI_BASE_URL  — Optional. Custom OpenAI-compatible base URL
                       (e.g., DashScope/Bailian OpenAI endpoint).
    VOYAGE_API_KEY   — Optional. If set, embed() can use Voyage models.

Supported models:
    Completion: gpt-4.1, gpt-4.1-mini, gpt-4.1-nano, gpt-4o, gpt-4o-mini, etc.
    Embedding: text-embedding-3-small (default, 1536d), text-embedding-3-large (3072d),
               text-embedding-ada-002 (legacy), voyage-code-3 (via voyageai SDK)

Key implementation details:
    - Uses OpenAI's native response_format={"type": "json_object"} for JSON mode
    - extract() uses response_format + schema validation for structured output
    - embed() routes to OpenAI's embeddings API by default, or Voyage if
      embed_model starts with "voyage-"
    - Cost estimation uses a hardcoded pricing table

References:
    https://platform.openai.com/docs/api-reference
    https://platform.openai.com/docs/guides/embeddings

Dependencies: openai SDK, voyageai SDK (optional for Voyage embeddings).
Implements: cortex.llm.provider.LLMProvider ABC.
"""

from __future__ import annotations

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

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

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Per-model pricing (USD / 1M tokens)
# ---------------------------------------------------------------------------
_OPENAI_COMPLETION_PRICING: dict[str, tuple[float, float]] = {
    # model_key: (input $/1M, output $/1M)
    "gpt-4.1": (2.0, 8.0),
    "gpt-4.1-mini": (0.40, 1.60),
    "gpt-4.1-nano": (0.10, 0.40),
    "gpt-4o": (2.50, 10.0),
    "gpt-4o-mini": (0.15, 0.60),
    "gpt-4-turbo": (10.0, 30.0),
    "gpt-3.5-turbo": (0.50, 1.50),
}

_OPENAI_EMBED_PRICING: dict[str, float] = {
    # model_key: $/1M tokens
    "text-embedding-3-small": 0.02,
    "text-embedding-3-large": 0.13,
    "text-embedding-ada-002": 0.10,
}

_VOYAGE_PRICING: dict[str, float] = {
    "voyage-code-3": 0.18,
    "voyage-3": 0.06,
    "voyage-3-lite": 0.02,
    "voyage-3-large": 0.18,
}

_VOYAGE_MODELS = set(_VOYAGE_PRICING.keys())


def _openai_completion_cost(model: str, tokens_in: int, tokens_out: int) -> float:
    # Sort by key length descending so more-specific keys (e.g. "gpt-4.1-mini")
    # are matched before prefix keys (e.g. "gpt-4.1").
    key = next(
        (k for k in sorted(_OPENAI_COMPLETION_PRICING, key=len, reverse=True)
         if model.startswith(k)),
        None,
    )
    if key is None:
        return 0.0
    in_p, out_p = _OPENAI_COMPLETION_PRICING[key]
    return (tokens_in * in_p + tokens_out * out_p) / 1_000_000


def _openai_embed_cost(model: str, tokens: int) -> float:
    key = next(
        (k for k in sorted(_OPENAI_EMBED_PRICING, key=len, reverse=True)
         if model.startswith(k)),
        None,
    )
    price = _OPENAI_EMBED_PRICING.get(key, 0.02) if key else 0.02
    return tokens * price / 1_000_000


def _voyage_cost(model: str, tokens: int) -> float:
    price = _VOYAGE_PRICING.get(model, 0.18)
    return tokens * price / 1_000_000


class OpenAIProvider(LLMProvider):
    """
    LLMProvider backed by OpenAI Chat Completions + Embeddings APIs.

    Also supports Voyage embeddings when embed_model is a Voyage model ID.

    Environment variables:
        OPENAI_API_KEY   — required for complete() / extract() / OpenAI embed()
        OPENAI_BASE_URL  — optional custom OpenAI-compatible endpoint
        VOYAGE_API_KEY   — required when embed_model is a Voyage model
    """

    DEFAULT_COMPLETION_MODEL = "gpt-4.1-mini"
    DEFAULT_EMBED_MODEL = "text-embedding-3-small"
    PROVIDER_NAME = "openai"

    def __init__(
        self,
        completion_model: str | None = None,
        embed_model: str | None = None,
    ) -> None:
        super().__init__()
        self._completion_model = completion_model or self.DEFAULT_COMPLETION_MODEL
        self._embed_model = embed_model or self.DEFAULT_EMBED_MODEL
        self._openai_client: Any = None
        self._voyage_client: Any = None
        self._custom_base_url: bool = bool(os.environ.get("OPENAI_BASE_URL"))

    # --- internal helpers ---------------------------------------------------

    def _get_openai(self) -> Any:
        if self._openai_client is None:
            try:
                import openai  # type: ignore
            except ImportError:
                raise RuntimeError(
                    "openai package not installed. Run: pip install openai"
                )
            api_key = os.environ.get("OPENAI_API_KEY")
            if not api_key:
                raise RuntimeError(
                    "OPENAI_API_KEY environment variable is not set."
                )
            base_url = os.environ.get("OPENAI_BASE_URL")
            kwargs: dict[str, Any] = {"api_key": api_key}
            if base_url:
                kwargs["base_url"] = base_url
            self._openai_client = openai.AsyncOpenAI(**kwargs)
        return self._openai_client

    def _get_voyage(self) -> Any:
        if self._voyage_client is None:
            try:
                import voyageai  # type: ignore
            except ImportError:
                raise RuntimeError(
                    "voyageai package not installed. Run: pip install voyageai"
                )
            api_key = os.environ.get("VOYAGE_API_KEY")
            if not api_key:
                raise RuntimeError(
                    "VOYAGE_API_KEY environment variable is not set."
                )
            self._voyage_client = voyageai.AsyncClient(api_key=api_key)
        return self._voyage_client

    # --- LLMProvider interface ----------------------------------------------

    async def complete(
        self,
        system: str,
        user: str,
        model: str | None = None,
        response_format: str = "text",
        max_tokens: int | None = None,
        temperature: float = 0.0,
    ) -> CompletionResponse:
        resolved_model = self._resolve_model(model, self._completion_model)
        client = self._get_openai()

        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":
            if self._custom_base_url:
                # Custom base URLs (e.g. Bailian/DashScope) may not support
                # response_format=json_object reliably. Skip the parameter and
                # rely on the system prompt to enforce JSON output.
                logger.debug(
                    "Skipping response_format=json_object for custom base URL"
                )
            else:
                kwargs["response_format"] = {"type": "json_object"}

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

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

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

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

    async def extract(
        self,
        system: str,
        user: str,
        schema: dict[str, Any],
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float = 0.0,
    ) -> ExtractionResponse:
        """
        Structured extraction via OpenAI JSON Schema response format.
        Uses response_format={"type":"json_schema"} when supported,
        falls back to response_format={"type":"json_object"} + prompt injection.
        """
        resolved_model = self._resolve_model(model, self._completion_model)
        client = self._get_openai()

        schema_hint = json.dumps(schema, indent=2)
        augmented_system = (
            f"{system}\n\n"
            f"Respond with valid JSON that matches this schema:\n{schema_hint}"
        )

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

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

        raw = response.choices[0].message.content or "{}"
        data = json.loads(raw)

        tokens_in = response.usage.prompt_tokens
        tokens_out = response.usage.completion_tokens
        cost = _openai_completion_cost(resolved_model, tokens_in, tokens_out)

        record = UsageRecord(
            provider=self.PROVIDER_NAME,
            model=resolved_model,
            call_type="extract",
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            cost_estimate=cost,
            latency_ms=latency_ms,
        )
        self._log_usage(record)

        return ExtractionResponse(
            data=data,
            model=resolved_model,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            cost_estimate=cost,
            latency_ms=latency_ms,
        )

    async def embed(
        self,
        texts: list[str],
        model: str | None = None,
        input_type: str = "query",
    ) -> EmbedResponse:
        """
        Embeddings via OpenAI or Voyage, depending on model.

        If resolved model is a Voyage model (voyage-*), routes to Voyage API.
        Otherwise uses OpenAI Embeddings API.

        input_type is passed to Voyage (ignored by OpenAI).
        """
        resolved_model = model or self._embed_model

        if resolved_model in _VOYAGE_MODELS:
            return await self._embed_voyage(texts, resolved_model, input_type)
        return await self._embed_openai(texts, resolved_model)

    async def _embed_openai(
        self, texts: list[str], model: str
    ) -> EmbedResponse:
        client = self._get_openai()

        t0 = time.perf_counter()
        response = await retry_with_backoff(
            lambda: client.embeddings.create(
                model=model, input=texts, encoding_format="float",
            ),
            operation_name=f"openai.embed({model})",
        )
        latency_ms = (time.perf_counter() - t0) * 1000

        embeddings = [item.embedding for item in response.data]
        tokens_in = response.usage.prompt_tokens
        cost = _openai_embed_cost(model, tokens_in)

        record = UsageRecord(
            provider=self.PROVIDER_NAME,
            model=model,
            call_type="embed",
            tokens_in=tokens_in,
            tokens_out=0,
            cost_estimate=cost,
            latency_ms=latency_ms,
        )
        self._log_usage(record)

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

    async def _embed_voyage(
        self, texts: list[str], model: str, input_type: str
    ) -> EmbedResponse:
        client = self._get_voyage()

        t0 = time.perf_counter()
        result = await retry_with_backoff(
            lambda: client.embed(texts, model=model, input_type=input_type),
            operation_name=f"openai.embed_voyage({model})",
        )
        latency_ms = (time.perf_counter() - t0) * 1000

        embeddings = result.embeddings
        tokens_in = result.total_tokens if hasattr(result, "total_tokens") else 0
        cost = _voyage_cost(model, tokens_in)

        record = UsageRecord(
            provider="voyage",
            model=model,
            call_type="embed",
            tokens_in=tokens_in,
            tokens_out=0,
            cost_estimate=cost,
            latency_ms=latency_ms,
        )
        self._log_usage(record)

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