"""
cortex/llm/venice_provider.py — Venice.ai Implementation of LLMProvider.

Venice.ai offers an OpenAI-compatible chat completions API with access to
powerful models (grok-41-fast, gemini-3-flash-preview) at zero cost for
self-hosted / API-key access. This makes Venice an excellent middle-tier
provider: more capable than Ollama (no local GPU required), cheaper than
Anthropic/OpenAI.

Venice is fully OpenAI-compatible, so we reuse the openai SDK with a custom
base URL. Two Venice-specific parameters are injected into every completion
call to prevent interference and reduce overhead:

    venice_parameters.include_venice_system_prompt = False
        Prevents Venice's default system prompt from overriding ours.

    venice_parameters.disable_thinking = True
        Skips the reasoning step on models that support it (saves tokens
        and reduces latency for extraction tasks).

API keys / endpoints (from environment):
    VENICE_API_KEY   — Required for all completion/extraction operations.
                       If absent, the provider silently fails to init so
                       the FallbackLLMProvider skips it gracefully.
    VENICE_BASE_URL  — Optional override. Default: https://api.venice.ai/api/v1

Supported models (as of 2026):
    Completion: grok-41-fast (default), gemini-3-flash-preview
    Embedding: NOT SUPPORTED — raises NotImplementedError.
               Use Voyage or OpenAI embeddings instead.

Key implementation details:
    - Uses openai SDK with custom base_url pointing at Venice
    - Injects venice_parameters extra_body on every completion call
    - response_format=json_object is supported but skipped for custom
      base URLs that may not handle it (matches OpenAI provider pattern)
    - Cost estimation uses a conservative pricing table (Venice is free-tier
      for most models — cost_estimate will be 0.0 unless pricing is added)
    - Logs Venice-specific response headers for observability:
        x-venice-balance-usd    — remaining account balance
        x-ratelimit-remaining-requests  — requests left this minute
    - Retry on 429/503 via retry_with_backoff (inherited behavior)
    - Circuit breaker support via FallbackLLMProvider (no changes needed here)

References:
    https://docs.venice.ai/api-reference

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 Any

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

logger = logging.getLogger(__name__)

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

_DEFAULT_BASE_URL = "https://api.venice.ai/api/v1"
_DEFAULT_COMPLETION_MODEL = "grok-41-fast"

# Venice-specific extra parameters injected into every completion request.
# include_venice_system_prompt=False: prevent Venice default prompt override.
# disable_thinking=True: skip reasoning overhead for extraction tasks.
_VENICE_PARAMETERS: dict[str, Any] = {
    "include_venice_system_prompt": False,
    "disable_thinking": True,
}

# ---------------------------------------------------------------------------
# Pricing (Venice free-tier models → 0.0 cost)
# Update if Venice introduces paid tiers.
# ---------------------------------------------------------------------------
_VENICE_COMPLETION_PRICING: dict[str, tuple[float, float]] = {
    # model_key: (input $/1M, output $/1M)
    # Currently Venice is free — all prices are 0.0.
    # Add paid model pricing here if Venice introduces billing.
    "grok-41-fast": (0.0, 0.0),
    "gemini-3-flash-preview": (0.0, 0.0),
}


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


# ---------------------------------------------------------------------------
# VeniceProvider
# ---------------------------------------------------------------------------

class VeniceProvider(LLMProvider):
    """
    LLMProvider backed by Venice.ai Chat Completions API.

    Venice is OpenAI-compatible. Uses the openai SDK pointed at Venice's
    base URL with Venice-specific extra_body parameters on every call.

    This provider is **completely optional** — if VENICE_API_KEY is not set,
    provider init raises RuntimeError and FallbackLLMProvider silently skips
    it (no crash, no noise in normal operation).

    Environment variables:
        VENICE_API_KEY   — required
        VENICE_BASE_URL  — optional (default: https://api.venice.ai/api/v1)

    Embeddings:
        Venice does not offer high-quality embedding models. embed() raises
        NotImplementedError. The FallbackLLMProvider's embed cascade will
        skip Venice and fall through to OpenAI/Voyage/Ollama for embeddings.
    """

    DEFAULT_COMPLETION_MODEL = _DEFAULT_COMPLETION_MODEL
    DEFAULT_EMBED_MODEL = ""  # not supported
    PROVIDER_NAME = "venice"

    def __init__(
        self,
        completion_model: str | None = None,
        base_url: str | None = None,
    ) -> None:
        super().__init__()
        self._completion_model = completion_model or self.DEFAULT_COMPLETION_MODEL
        self._base_url = (
            base_url
            or os.environ.get("VENICE_BASE_URL", _DEFAULT_BASE_URL)
        ).rstrip("/")
        self._openai_client: Any = None
        # Validate API key at init time so FallbackLLMProvider skips us cleanly.
        api_key = os.environ.get("VENICE_API_KEY")
        if not api_key:
            raise RuntimeError(
                "VENICE_API_KEY environment variable is not set. "
                "Venice provider will be skipped in the cascade."
            )
        self._api_key = api_key

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

    def _get_client(self) -> Any:
        """Return a lazily-initialised AsyncOpenAI client pointed at Venice."""
        if self._openai_client is None:
            try:
                import openai  # type: ignore
            except ImportError:
                raise RuntimeError(
                    "openai package not installed. Run: pip install openai"
                )
            self._openai_client = openai.AsyncOpenAI(
                api_key=self._api_key,
                base_url=self._base_url,
            )
        return self._openai_client

    def _log_venice_headers(self, response: Any) -> None:
        """Log Venice-specific observability headers if present."""
        # The openai SDK exposes raw response headers via response.headers on
        # the HTTP response object. Headers may be absent if the SDK version
        # or call path doesn't expose them — handle gracefully.
        try:
            headers = getattr(response, "headers", None) or {}
            balance = headers.get("x-venice-balance-usd")
            remaining = headers.get("x-ratelimit-remaining-requests")
            if balance is not None:
                logger.debug("venice balance_usd=%s", balance)
            if remaining is not None:
                logger.debug("venice ratelimit_remaining_requests=%s", remaining)
        except Exception:
            pass  # Never let header logging crash a completion

    # --- 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:
        """
        Chat completion via Venice OpenAI-compatible API.

        Venice-specific parameters are injected via extra_body on every call:
            - include_venice_system_prompt: False (suppress Venice defaults)
            - disable_thinking: True (skip reasoning overhead)

        response_format='json' is supported but skipped here (we use a custom
        base URL). The system prompt should instruct JSON output when needed.
        """
        resolved_model = self._resolve_model(model, 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},
            ],
            # Venice-specific parameters injected on every request.
            # These must be passed as extra_body since they are not part of
            # the standard OpenAI API schema.
            "extra_body": {
                "venice_parameters": _VENICE_PARAMETERS,
            },
        }
        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens
        if response_format == "json":
            # Venice is a custom base URL — skip response_format=json_object
            # to avoid compatibility issues (matches OpenAI provider pattern).
            # The system prompt handles JSON enforcement instead.
            logger.debug(
                "VeniceProvider: skipping response_format=json_object "
                "(custom base URL — rely on system prompt for JSON enforcement)"
            )

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

        self._log_venice_headers(response)

        text = response.choices[0].message.content or ""
        tokens_in = response.usage.prompt_tokens
        tokens_out = response.usage.completion_tokens
        cost = _venice_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 JSON extraction via Venice API.

        Augments the system prompt with the JSON schema and relies on
        Venice's grok-41-fast instruction-following to produce valid JSON.
        Venice parameters (no system prompt override, no thinking) are
        injected automatically via complete().
        """
        resolved_model = self._resolve_model(model, self._completion_model)

        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}"
        )

        t0 = time.perf_counter()
        cr = await self.complete(
            system=augmented_system,
            user=user,
            model=resolved_model,
            response_format="json",
            max_tokens=max_tokens,
            temperature=temperature,
        )
        latency_ms = (time.perf_counter() - t0) * 1000

        try:
            data = json.loads(cr.text)
        except json.JSONDecodeError as exc:
            logger.warning(
                "VeniceProvider.extract: JSON parse failed — %s. Raw: %.200s",
                exc,
                cr.text,
            )
            data = {"_raw": cr.text, "_parse_error": str(exc)}

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

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

    async def embed(
        self,
        texts: list[str],
        model: str | None = None,
        input_type: str = "query",
    ) -> EmbedResponse:
        """
        Embeddings are NOT supported by VeniceProvider.

        Venice does not offer high-quality embedding models. The
        FallbackLLMProvider will skip this provider and fall through to
        OpenAI or Voyage for embedding generation.

        Raises:
            NotImplementedError: Always. Use OpenAI or Voyage for embeddings.
        """
        raise NotImplementedError(
            "VeniceProvider does not support embeddings. "
            "Configure an OpenAI or Voyage provider for embedding generation."
        )
