"""
cortex/llm/anthropic_provider.py — Anthropic Claude Implementation of LLMProvider.

Wraps Anthropic's Messages API for chat completion and structured extraction.
For embeddings, delegates to Voyage AI (since Anthropic has no native embedding
endpoint).

API keys / endpoints (from environment — never hardcoded):
    ANTHROPIC_API_KEY   — Required for complete() and extract().
    ANTHROPIC_BASE_URL  — Optional. Custom Anthropic-compatible base URL
                          (e.g., DashScope/Bailian Anthropic endpoint).
    VOYAGE_API_KEY      — Required for embed(). If missing, embed() raises RuntimeError.

Supported models:
    Completion: claude-opus-4-5, claude-sonnet-4-5, claude-haiku-4-5 (and variants)
    Embedding: voyage-code-3 (via voyageai SDK)

Key implementation details:
    - Uses the Anthropic Python SDK's messages.create() for completions
    - Maps Cortex's 'json' response_format to Anthropic's prefill trick
      (starts the response with '{' to encourage JSON output)
    - Cost estimation uses a hardcoded pricing table (kept up-to-date manually)
    - extract() reuses complete() with response_format='json' — Anthropic doesn't
      have a native structured extraction API like OpenAI's

References:
    https://docs.anthropic.com/en/api/messages
    https://docs.voyageai.com/reference/embeddings-api

Dependencies: anthropic SDK, voyageai SDK (optional for 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 per 1M tokens) — update as pricing changes
# ---------------------------------------------------------------------------
_ANTHROPIC_PRICING: dict[str, tuple[float, float]] = {
    # model_id: (input $/1M, output $/1M)
    "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-20240229": (15.0, 75.0),
    "claude-3-5-sonnet-20241022": (3.0, 15.0),
    "claude-3-haiku-20240307": (0.25, 1.25),
}

_VOYAGE_PRICING: dict[str, float] = {
    # model_id: $/1M tokens
    "voyage-code-3": 0.18,
    "voyage-3": 0.06,
    "voyage-3-lite": 0.02,
    "voyage-3-large": 0.18,
}


def _anthropic_cost(model: str, tokens_in: int, tokens_out: int) -> float:
    key = next((k for k in _ANTHROPIC_PRICING if model.startswith(k)), None)
    if key is None:
        return 0.0
    in_price, out_price = _ANTHROPIC_PRICING[key]
    return (tokens_in * in_price + tokens_out * out_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 AnthropicProvider(LLMProvider):
    """
    LLMProvider backed by Anthropic Messages API + Voyage embeddings.

    Environment variables:
        ANTHROPIC_API_KEY   — required for complete() / extract()
        ANTHROPIC_BASE_URL  — optional custom Anthropic-compatible endpoint
        VOYAGE_API_KEY      — required for embed()
    """

    DEFAULT_COMPLETION_MODEL = "claude-haiku-4-5"
    DEFAULT_EMBED_MODEL = "voyage-code-3"
    PROVIDER_NAME = "anthropic"

    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

        # Lazy-init clients so we don't import at module load if SDK missing
        self._anthropic_client: Any = None
        self._voyage_client: Any = None

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

    def _get_anthropic(self) -> Any:
        if self._anthropic_client is None:
            try:
                import anthropic  # type: ignore
            except ImportError:
                raise RuntimeError(
                    "anthropic package not installed. Run: pip install anthropic"
                )
            api_key = os.environ.get("ANTHROPIC_API_KEY")
            if not api_key:
                raise RuntimeError(
                    "ANTHROPIC_API_KEY environment variable is not set."
                )
            base_url = os.environ.get("ANTHROPIC_BASE_URL")
            kwargs: dict[str, Any] = {"api_key": api_key}
            if base_url:
                kwargs["base_url"] = base_url
            self._anthropic_client = anthropic.AsyncAnthropic(**kwargs)
        return self._anthropic_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_anthropic()

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

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

        text = response.content[0].text
        tokens_in = response.usage.input_tokens
        tokens_out = response.usage.output_tokens
        cost = _anthropic_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 Anthropic tool-use (guaranteed JSON).
        We expose the schema as a single tool called 'extract_result'.
        """
        resolved_model = self._resolve_model(model, self._completion_model)
        client = self._get_anthropic()

        tool_def = {
            "name": "extract_result",
            "description": "Return the extracted structured data.",
            "input_schema": schema,
        }

        t0 = time.perf_counter()
        response = await retry_with_backoff(
            lambda: client.messages.create(
                model=resolved_model,
                max_tokens=max_tokens or 4096,
                temperature=temperature,
                system=system,
                messages=[{"role": "user", "content": user}],
                tools=[tool_def],
                tool_choice={"type": "tool", "name": "extract_result"},
            ),
            operation_name=f"anthropic.extract({resolved_model})",
        )
        latency_ms = (time.perf_counter() - t0) * 1000

        # Extract the tool_use block
        tool_block = next(
            (b for b in response.content if b.type == "tool_use"),
            None,
        )
        if tool_block is None:
            # Fallback: parse text as JSON
            raw_text = response.content[0].text if response.content else "{}"
            data = json.loads(raw_text)
        else:
            data = tool_block.input  # already a dict

        tokens_in = response.usage.input_tokens
        tokens_out = response.usage.output_tokens
        cost = _anthropic_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 Voyage API (Anthropic has no native embed endpoint).
        input_type: "query" or "document" — Voyage uses this for asymmetric search.
        """
        resolved_model = model or self._embed_model
        client = self._get_voyage()

        t0 = time.perf_counter()
        result = await retry_with_backoff(
            lambda: client.embed(texts, model=resolved_model, input_type=input_type),
            operation_name=f"anthropic.embed_voyage({resolved_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(resolved_model, tokens_in)

        record = UsageRecord(
            provider="voyage",
            model=resolved_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=resolved_model,
            tokens_in=tokens_in,
            tokens_out=0,
            cost_estimate=cost,
            latency_ms=latency_ms,
        )
