"""
Anthropic Claude LLM Provider

Primary LLM provider for SAGE, using Claude for NCS Q&A.
"""

import logging
import time

from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from sage.llm.base import (
    BaseLLMProvider,
    LLMAuthenticationError,
    LLMContextLengthError,
    LLMError,
    LLMProvider,
    LLMProviderError,
    LLMRateLimitError,
    LLMResponse,
    LLMUsage,
)
from sage.llm.config import AnthropicConfig, ModelFamily

logger = logging.getLogger(__name__)

try:
    import anthropic
    from anthropic import (
        APIConnectionError,
        APIStatusError,
        AuthenticationError,
        RateLimitError,
    )

    ANTHROPIC_AVAILABLE = True
except ImportError:
    ANTHROPIC_AVAILABLE = False
    anthropic = None  # type: ignore[misc, assignment]


class ClaudeProvider(BaseLLMProvider):
    """
    Claude LLM provider using the Anthropic API.

    Primary provider for SAGE NCS, optimized for safety-critical
    technical Q&A with proper citation handling.

    Example:
        config = AnthropicConfig(api_key="sk-ant-...")
        async with ClaudeProvider(config) as provider:
            response = await provider.generate(
                prompt="What is the subcritical mass limit for U-235?",
                system="You are a nuclear safety expert. Always cite sources.",
                temperature=0.3
            )
            print(response.content)
    """

    def __init__(self, config: AnthropicConfig):
        """Initialize Claude provider."""
        if not ANTHROPIC_AVAILABLE:
            raise ImportError(
                "anthropic package not installed. Install with: pip install anthropic"
            )

        self.config = config
        self._client: anthropic.AsyncAnthropic | None = None
        self._initialized = False

    async def initialize(self) -> None:
        """Initialize the Anthropic client."""
        if self._initialized:
            return

        logger.info(f"Initializing Claude provider with model {self.config.model}")

        self._client = anthropic.AsyncAnthropic(
            api_key=self.config.api_key,
            base_url=self.config.api_base_url,
            timeout=self.config.timeout,
        )

        # Validate API key with a minimal request
        try:
            # Just verify connection - actual validation happens on first request
            self._initialized = True
            logger.info("Claude provider initialized")
        except AuthenticationError as e:
            raise LLMAuthenticationError(f"Invalid Anthropic API key: {e}") from e

    async def close(self) -> None:
        """Close the client."""
        if self._client:
            await self._client.close()
            self._client = None
        self._initialized = False

    @property
    def provider_name(self) -> LLMProvider:
        """Get provider identifier."""
        return LLMProvider.ANTHROPIC

    @property
    def model_name(self) -> str:
        """Get model name."""
        return self.config.model

    @retry(
        retry=retry_if_exception_type(LLMRateLimitError),
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=60),
    )
    async def generate(
        self,
        prompt: str,
        system: str | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        stop_sequences: list[str] | None = None,
        **kwargs,
    ) -> LLMResponse:
        """
        Generate a response from Claude.

        Args:
            prompt: User prompt
            system: System prompt
            temperature: Sampling temperature (default: config value)
            max_tokens: Max response tokens (default: config value)
            stop_sequences: Stop sequences
            **kwargs: Additional Anthropic-specific parameters

        Returns:
            LLMResponse with generated content
        """
        if not self._initialized:
            await self.initialize()

        messages = [{"role": "user", "content": prompt}]

        return await self.generate_with_messages(
            messages=messages,
            system=system,
            temperature=temperature,
            max_tokens=max_tokens,
            stop_sequences=stop_sequences,
            **kwargs,
        )

    @retry(
        retry=retry_if_exception_type(LLMRateLimitError),
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=1, max=60),
    )
    async def generate_with_messages(
        self,
        messages: list[dict[str, str]],
        system: str | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        stop_sequences: list[str] | None = None,
        **kwargs,
    ) -> LLMResponse:
        """
        Generate a response from a message history.

        Args:
            messages: List of messages
            system: System prompt
            temperature: Sampling temperature
            max_tokens: Max response tokens
            stop_sequences: Stop sequences
            **kwargs: Additional parameters

        Returns:
            LLMResponse with generated content
        """
        if not self._initialized:
            await self.initialize()

        temperature = temperature if temperature is not None else self.config.temperature
        max_tokens = max_tokens or self.config.max_tokens

        start_time = time.time()

        try:
            # Build request parameters
            request_params = {
                "model": self.config.model,
                "messages": messages,
                "max_tokens": max_tokens,
                "temperature": temperature,
            }

            if system:
                request_params["system"] = system

            if stop_sequences:
                request_params["stop_sequences"] = stop_sequences

            # Add any extra kwargs
            request_params.update(kwargs)

            # Make API call
            response = await self._client.messages.create(**request_params)

            latency_ms = (time.time() - start_time) * 1000

            # Extract content
            content = ""
            if response.content:
                content = response.content[0].text

            # Build usage info
            usage = LLMUsage(
                prompt_tokens=response.usage.input_tokens,
                completion_tokens=response.usage.output_tokens,
                total_tokens=response.usage.input_tokens + response.usage.output_tokens,
                cached_tokens=getattr(response.usage, "cache_read_input_tokens", 0),
            )

            return LLMResponse(
                content=content,
                model=response.model,
                provider=LLMProvider.ANTHROPIC,
                usage=usage,
                latency_ms=latency_ms,
                finish_reason=response.stop_reason or "",
                raw_response=response,
            )

        except RateLimitError as e:
            logger.warning(f"Rate limit hit: {e}")
            raise LLMRateLimitError(str(e), retry_after=60.0) from e

        except AuthenticationError as e:
            raise LLMAuthenticationError(f"Authentication failed: {e}") from e

        except APIStatusError as e:
            if "context_length" in str(e).lower() or "too long" in str(e).lower():
                raise LLMContextLengthError(f"Prompt too long: {e}") from e
            raise LLMProviderError(f"Anthropic API error: {e}") from e

        except APIConnectionError as e:
            raise LLMProviderError(f"Connection error: {e}") from e

        except Exception as e:
            logger.error(f"Unexpected error calling Claude: {e}")
            raise LLMError(f"Unexpected error: {e}") from e


async def create_claude_provider(
    api_key: str,
    model: str = ModelFamily.CLAUDE_3_5_SONNET.value,
    **kwargs,
) -> ClaudeProvider:
    """
    Factory function to create and initialize a Claude provider.

    Args:
        api_key: Anthropic API key
        model: Model to use
        **kwargs: Additional config parameters

    Returns:
        Initialized ClaudeProvider
    """
    config = AnthropicConfig(api_key=api_key, model=model, **kwargs)
    provider = ClaudeProvider(config)
    await provider.initialize()
    return provider
