"""
cortex/llm/host_provider.py — HostProvider: LLM/embedding provider using host-supplied credentials.

In "host" mode, the OpenClaw plugin passes its already-configured provider
credentials to Cortex at startup via POST /api/v1/config/providers. This
provider wraps those credentials and routes LLM calls to the appropriate
vendor SDK (OpenAI, Anthropic, Voyage).

Key design decisions:
    - OpenAI SDK (AsyncOpenAI) for OpenAI, Venice, and any OpenAI-compatible provider
    - Anthropic SDK (AsyncAnthropic) for Claude models (different API format)
    - Voyage SDK (AsyncClient) for Voyage embeddings
    - OpenAI SDK for OpenAI embeddings
    - Auto-detects provider from model name, base_url, or explicit provider field
    - No fallback chains — if the configured provider fails, the operation fails cleanly

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

from __future__ import annotations

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

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

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Pricing tables (copied from existing providers for cost estimation)
# ---------------------------------------------------------------------------

_COMPLETION_PRICING: dict[str, tuple[float, float]] = {
    # model_prefix: (input $/1M tokens, output $/1M tokens)
    # OpenAI
    "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),
    # Anthropic
    "claude-opus-4-5": (15.0, 75.0),
    "claude-sonnet-4-5": (3.0, 15.0),
    "claude-sonnet-4-6": (3.0, 15.0),
    "claude-haiku-4-5": (0.80, 4.0),
    "claude-3-opus": (15.0, 75.0),
    "claude-3-5-sonnet": (3.0, 15.0),
    "claude-3-haiku": (0.25, 1.25),
    # Venice (OpenAI-compatible, mostly free-tier)
    "grok-41-fast": (0.0, 0.0),
}

_EMBED_PRICING: dict[str, float] = {
    # model: $/1M tokens
    "text-embedding-3-small": 0.02,
    "text-embedding-3-large": 0.13,
    "text-embedding-ada-002": 0.10,
    "voyage-code-3": 0.18,
    "voyage-3": 0.06,
    "voyage-3-lite": 0.02,
    "voyage-3-large": 0.18,
}

_VOYAGE_MODELS = {"voyage-code-3", "voyage-3", "voyage-3-lite", "voyage-3-large"}


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


def _estimate_embed_cost(model: str, tokens: int) -> float:
    """Estimate cost in USD for an embedding call."""
    price = _EMBED_PRICING.get(model, 0.02)
    return tokens * price / 1_000_000


def _is_anthropic_model(model: str) -> bool:
    """Return True if the model should use the Anthropic SDK."""
    return model.startswith("claude-")


def _is_voyage_model(model: str) -> bool:
    """Return True if the model should use the Voyage SDK."""
    return model in _VOYAGE_MODELS or model.startswith("voyage-")


class HostProvider(LLMProvider):
    """LLM/embedding provider that uses credentials supplied by the host application.

    In host mode, the OpenClaw plugin passes its provider credentials at startup.
    HostProvider wraps those credentials and routes calls to the appropriate SDK:
    - OpenAI SDK for GPT models and OpenAI-compatible APIs (Venice, etc.)
    - Anthropic SDK for Claude models
    - Voyage SDK for Voyage embedding models
    - OpenAI SDK for OpenAI embedding models

    Args:
        llm_config: Dict with keys: api_key, base_url, model, provider (optional).
        embedding_config: Dict with keys: api_key, base_url, model, dim, provider (optional).
    """

    PROVIDER_NAME = "host"

    def __init__(
        self,
        llm_config: dict[str, Any],
        embedding_config: dict[str, Any],
    ) -> None:
        super().__init__()
        self._llm_config = llm_config
        self._embedding_config = embedding_config

        # Default models from config
        self.DEFAULT_COMPLETION_MODEL = llm_config.get("model", "gpt-4.1-nano")
        self.DEFAULT_EMBED_MODEL = embedding_config.get("model", "voyage-code-3")

        # Lazy-initialized SDK clients
        self._openai_client: Any = None
        self._anthropic_client: Any = None
        self._voyage_client: Any = None
        self._openai_embed_client: Any = None

    # --- Client factories (lazy init) ---------------------------------------

    def _get_openai(self) -> Any:
        """Get or create AsyncOpenAI client for LLM completions."""
        if self._openai_client is None:
            try:
                import openai
            except ImportError:
                raise RuntimeError("openai package not installed. Run: pip install openai")
            api_key = self._llm_config.get("api_key", "")
            if not api_key:
                raise RuntimeError("No LLM API key provided in host config.")
            kwargs: dict[str, Any] = {"api_key": api_key}
            base_url = self._llm_config.get("base_url", "")
            if base_url:
                kwargs["base_url"] = base_url
            self._openai_client = openai.AsyncOpenAI(**kwargs)
        return self._openai_client

    def _get_anthropic(self) -> Any:
        """Get or create AsyncAnthropic client."""
        if self._anthropic_client is None:
            try:
                import anthropic
            except ImportError:
                raise RuntimeError("anthropic package not installed. Run: pip install anthropic")
            api_key = self._llm_config.get("api_key", "")
            if not api_key:
                raise RuntimeError("No LLM API key provided in host config.")
            kwargs: dict[str, Any] = {"api_key": api_key}
            base_url = self._llm_config.get("base_url", "")
            if base_url:
                kwargs["base_url"] = base_url
            self._anthropic_client = anthropic.AsyncAnthropic(**kwargs)
        return self._anthropic_client

    def _get_voyage(self) -> Any:
        """Get or create Voyage AsyncClient for embeddings."""
        if self._voyage_client is None:
            try:
                import voyageai
            except ImportError:
                raise RuntimeError("voyageai package not installed. Run: pip install voyageai")
            api_key = self._embedding_config.get("api_key", "")
            if not api_key:
                raise RuntimeError("No embedding API key provided in host config.")
            self._voyage_client = voyageai.AsyncClient(api_key=api_key)
        return self._voyage_client

    def _get_openai_embed(self) -> Any:
        """Get or create AsyncOpenAI client for OpenAI embeddings."""
        if self._openai_embed_client is None:
            try:
                import openai
            except ImportError:
                raise RuntimeError("openai package not installed. Run: pip install openai")
            api_key = self._embedding_config.get("api_key", "")
            if not api_key:
                raise RuntimeError("No embedding API key provided in host config.")
            kwargs: dict[str, Any] = {"api_key": api_key}
            base_url = self._embedding_config.get("base_url", "")
            if base_url:
                kwargs["base_url"] = base_url
            self._openai_embed_client = openai.AsyncOpenAI(**kwargs)
        return self._openai_embed_client

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

    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:
        """Send a completion request, routing to OpenAI or Anthropic SDK."""
        resolved_model = self._resolve_model(model, self.DEFAULT_COMPLETION_MODEL)

        if _is_anthropic_model(resolved_model):
            return await self._complete_anthropic(
                system, user, resolved_model, response_format, max_tokens, temperature,
            )
        return await self._complete_openai(
            system, user, resolved_model, response_format, max_tokens, temperature,
        )

    async def _complete_openai(
        self,
        system: str,
        user: str,
        model: str,
        response_format: str,
        max_tokens: int | None,
        temperature: float,
    ) -> CompletionResponse:
        """Complete via OpenAI-compatible API (OpenAI, Venice, etc.)."""
        client = self._get_openai()

        kwargs: dict[str, Any] = {
            "model": 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":
            # Venice and other OpenAI-compatible APIs support json_object
            base_url = self._llm_config.get("base_url", "")
            if base_url and "venice" in base_url.lower():
                # Venice is OpenAI-compatible, json_object works
                kwargs["response_format"] = {"type": "json_object"}
            elif base_url:
                # Unknown custom endpoint — skip response_format to be safe
                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"host.openai.complete({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 = _estimate_completion_cost(model, tokens_in, tokens_out)

        self._log_usage(UsageRecord(
            provider="host/openai",
            model=model,
            call_type="complete",
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            cost_estimate=cost,
            latency_ms=latency_ms,
        ))

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

    async def _complete_anthropic(
        self,
        system: str,
        user: str,
        model: str,
        response_format: str,
        max_tokens: int | None,
        temperature: float,
    ) -> CompletionResponse:
        """Complete via Anthropic Messages API."""
        client = self._get_anthropic()

        kwargs: dict[str, Any] = {
            "model": model,
            "max_tokens": max_tokens or 4096,
            "temperature": temperature,
            "system": system,
            "messages": [{"role": "user", "content": user}],
        }

        # JSON mode: use prefill trick (Anthropic doesn't have json_object mode)
        if response_format == "json":
            kwargs["messages"].append({
                "role": "assistant",
                "content": "{",
            })

        t0 = time.perf_counter()
        response = await retry_with_backoff(
            lambda: client.messages.create(**kwargs),
            operation_name=f"host.anthropic.complete({model})",
        )
        latency_ms = (time.perf_counter() - t0) * 1000

        text = response.content[0].text if response.content else ""
        # Anthropic prefill: prepend the '{' we sent
        if response_format == "json" and text and not text.startswith("{"):
            text = "{" + text

        tokens_in = response.usage.input_tokens
        tokens_out = response.usage.output_tokens
        cost = _estimate_completion_cost(model, tokens_in, tokens_out)

        self._log_usage(UsageRecord(
            provider="host/anthropic",
            model=model,
            call_type="complete",
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            cost_estimate=cost,
            latency_ms=latency_ms,
        ))

        return CompletionResponse(
            text=text, model=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,
        schema: dict[str, Any],
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float = 0.0,
    ) -> ExtractionResponse:
        """Structured JSON extraction — forces output to match schema."""
        resolved_model = self._resolve_model(model, self.DEFAULT_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}"
        )

        # Use complete() with json format, then parse
        completion = await self.complete(
            system=augmented_system,
            user=user,
            model=resolved_model,
            response_format="json",
            max_tokens=max_tokens,
            temperature=temperature,
        )

        try:
            data = json.loads(completion.text)
        except json.JSONDecodeError:
            logger.warning(
                "host_provider.extract: failed to parse JSON from model response "
                "(model=%s, text_len=%d)", resolved_model, len(completion.text),
            )
            data = {}

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

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

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

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

    async def _embed_voyage(
        self, texts: list[str], model: str, input_type: str,
    ) -> EmbedResponse:
        """Embed via Voyage SDK."""
        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"host.voyage.embed({model})",
        )
        latency_ms = (time.perf_counter() - t0) * 1000

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

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

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

    async def _embed_openai(
        self, texts: list[str], model: str,
    ) -> EmbedResponse:
        """Embed via OpenAI Embeddings API."""
        client = self._get_openai_embed()

        t0 = time.perf_counter()
        response = await retry_with_backoff(
            lambda: client.embeddings.create(
                model=model, input=texts, encoding_format="float",
            ),
            operation_name=f"host.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 = _estimate_embed_cost(model, tokens_in)

        self._log_usage(UsageRecord(
            provider="host/openai",
            model=model,
            call_type="embed",
            tokens_in=tokens_in,
            tokens_out=0,
            cost_estimate=cost,
            latency_ms=latency_ms,
        ))

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

    # --- Utility methods ----------------------------------------------------

    def reconfigure(
        self,
        llm_config: Optional[dict[str, Any]] = None,
        embedding_config: Optional[dict[str, Any]] = None,
    ) -> None:
        """Update credentials at runtime (e.g. from POST /api/v1/config/providers).

        Resets cached SDK clients so they re-initialize with new credentials.
        """
        if llm_config is not None:
            self._llm_config = llm_config
            self.DEFAULT_COMPLETION_MODEL = llm_config.get("model", self.DEFAULT_COMPLETION_MODEL)
            self._openai_client = None
            self._anthropic_client = None

        if embedding_config is not None:
            self._embedding_config = embedding_config
            self.DEFAULT_EMBED_MODEL = embedding_config.get("model", self.DEFAULT_EMBED_MODEL)
            self._voyage_client = None
            self._openai_embed_client = None

        logger.info(
            "HostProvider reconfigured: llm_model=%s, embed_model=%s",
            self.DEFAULT_COMPLETION_MODEL, self.DEFAULT_EMBED_MODEL,
        )
