from __future__ import annotations

import contextlib
import time
from typing import Any

from awaf.providers.base import (
    LLMProvider,
    ProviderAuthError,
    ProviderConfig,
    ProviderConfigError,
    ProviderError,
    ProviderRateLimitError,
    ProviderResponse,
    ProviderTimeoutError,
)


class AnthropicProvider(LLMProvider):
    """Adapter for the Anthropic Messages API."""

    def __init__(self, config: ProviderConfig) -> None:
        super().__init__(config)
        self._client: object = None  # lazy-initialised on first use

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

    @property
    def default_model(self) -> str:
        return "claude-haiku-4-5-20251001"

    @property
    def supports_system_prompt(self) -> bool:
        return True

    def validate_config(self) -> None:
        try:
            import anthropic  # noqa: F401
        except ImportError as exc:
            raise ProviderConfigError(
                "Provider 'anthropic' requires additional dependencies. Run: pip install awaf[anthropic]",
                provider=self.config.provider_name,
                model=self.config.model,
            ) from exc

        if not self.config.api_key:
            raise ProviderConfigError(
                "Anthropic provider requires an API key. Set ANTHROPIC_API_KEY or api_key in awaf.toml.",
                provider=self.config.provider_name,
                model=self.config.model,
            )

    def complete(
        self,
        system_prompt: str,
        user_prompt: str,
        artifact_content: str | None = None,
    ) -> ProviderResponse:
        import anthropic

        if self._client is None:
            self._client = anthropic.Anthropic(api_key=self.config.api_key)
        client: anthropic.Anthropic = self._client  # type: ignore[assignment]
        model = self.config.model or self.default_model

        t0 = time.monotonic()
        try:
            # Artifact goes first in system as a cached block — cache key is just the
            # artifact content, shared across all 10 pillar calls. Pillar criteria come
            # second (uncached) so each pillar's different prompt doesn't break the key.
            # Without artifact: fall back to caching the system prompt directly.
            if artifact_content:
                system_blocks: list[Any] = [
                    {
                        "type": "text",
                        "text": artifact_content,
                        "cache_control": {"type": "ephemeral"},
                    },
                    {
                        "type": "text",
                        "text": system_prompt,
                    },
                ]
            else:
                system_blocks = [
                    {
                        "type": "text",
                        "text": system_prompt,
                        "cache_control": {"type": "ephemeral"},
                    }
                ]
            response = client.messages.create(
                model=model,
                max_tokens=self.config.max_tokens,
                temperature=self.config.temperature,
                system=system_blocks,
                messages=[{"role": "user", "content": user_prompt}],
                extra_headers={"anthropic-beta": "prompt-caching-2024-07-31"},
            )
        except anthropic.APIStatusError as exc:
            if exc.status_code == 429:
                retry_after: int | None = None
                raw_after = exc.response.headers.get("retry-after")
                if raw_after is not None:
                    with contextlib.suppress(ValueError):
                        retry_after = int(raw_after)
                raise ProviderRateLimitError(
                    str(exc),
                    provider=self.config.provider_name,
                    model=model,
                    retry_after_seconds=retry_after,
                ) from exc
            if exc.status_code in (401, 403):
                raise ProviderAuthError(
                    str(exc),
                    provider=self.config.provider_name,
                    model=model,
                ) from exc
            raise ProviderError(
                str(exc),
                provider=self.config.provider_name,
                model=model,
            ) from exc
        except anthropic.APITimeoutError as exc:
            raise ProviderTimeoutError(
                str(exc),
                provider=self.config.provider_name,
                model=model,
            ) from exc

        latency_ms = int((time.monotonic() - t0) * 1000)

        # Duck typing: only TextBlock has a str .text; other block types are skipped
        text_blocks = [
            b for b in response.content if hasattr(b, "text") and isinstance(b.text, str)
        ]
        content = text_blocks[0].text if text_blocks else ""

        # Preserve cache token counts separately for accurate cost estimation.
        # cache_creation/read not present on older SDK versions → default 0.
        cache_create = getattr(response.usage, "cache_creation_input_tokens", 0) or 0
        cache_read = getattr(response.usage, "cache_read_input_tokens", 0) or 0
        total_input = response.usage.input_tokens + cache_create + cache_read

        return ProviderResponse(
            content=content,
            input_tokens=total_input,
            output_tokens=response.usage.output_tokens,
            model=response.model,
            provider=self.config.provider_name,
            latency_ms=latency_ms,
            cache_creation_input_tokens=cache_create,
            cache_read_input_tokens=cache_read,
            raw=response.model_dump(),
        )

    def count_tokens(self, text: str) -> int:
        # Heuristic (~4 chars/token) is accurate enough for budget decisions and
        # avoids a live API round-trip per file during ingest (100+ calls otherwise).
        # Set AWAF_EXACT_TOKEN_COUNT=1 to use the Anthropic count_tokens API instead.
        import os

        if not os.environ.get("AWAF_EXACT_TOKEN_COUNT"):
            return max(1, len(text) // 4)
        try:
            import anthropic

            if self._client is None:
                self._client = anthropic.Anthropic(api_key=self.config.api_key)
            client: anthropic.Anthropic = self._client  # type: ignore[assignment]
            model = self.config.model or self.default_model
            result = client.beta.messages.count_tokens(
                model=model,
                messages=[{"role": "user", "content": text}],
            )
            return result.input_tokens
        except Exception:
            return max(1, len(text) // 4)
