"""
AIClient: Pluggable AI provider for embeddings and LLM operations.

Supports: OpenAI, Azure OpenAI, OpenAI-compatible APIs.
Provider selection via EMBED_PROVIDER / LLM_PROVIDER env vars.
"""
from __future__ import annotations

import functools
import hashlib
import logging
import os
import time
from typing import Protocol

logger = logging.getLogger(__name__)


# ==============================================================
# Retry utilities
# ==============================================================


def _is_transient_error(exc: Exception) -> bool:
    """Check if an error is transient (retryable).

    Detects OpenAI SDK transient exceptions via class name and status_code
    attribute (duck typing — no import required).

    Transient errors:
        - RateLimitError (429)
        - InternalServerError (500)
        - APIConnectionError (connection failure)
        - Any APIStatusError with status_code in {429, 500, 502, 503, 504}

    Non-transient (permanent) errors:
        - AuthenticationError (401)
        - BadRequestError (400)
        - NotFoundError (404)
        - PermissionDeniedError (403)
    """
    exc_type = type(exc).__name__
    if exc_type in ("RateLimitError", "InternalServerError", "APIConnectionError"):
        return True
    status = getattr(exc, "status_code", None)
    if status and status in (429, 500, 502, 503, 504):
        return True
    return False


def retry_on_transient(
    max_retries: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 30.0,
):
    """Decorator: exponential backoff retry for transient API errors.

    Non-transient errors (401, 400) are raised immediately without retry.

    Args:
        max_retries: Maximum number of retry attempts.
        base_delay: Initial delay in seconds (doubled each retry).
        max_delay: Maximum delay cap in seconds.
    """
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            last_exc = None
            for attempt in range(max_retries + 1):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    last_exc = e
                    if not _is_transient_error(e) or attempt == max_retries:
                        raise
                    delay = min(base_delay * (2 ** attempt), max_delay)
                    logger.warning(
                        "Transient error (attempt %d/%d), retrying in %.1fs: %s",
                        attempt + 1, max_retries, delay, e,
                    )
                    time.sleep(delay)
            raise last_exc  # pragma: no cover — unreachable, satisfies type checker
        return wrapper
    return decorator


# ==============================================================
# AIClient Protocol
# ==============================================================


class AIClient(Protocol):
    """Interface for AI operations (embedding + LLM judge)."""

    def embed(self, text: str) -> list[float]:
        """Generate embedding vector for text."""
        ...

    def llm_judge(self, prompt: str) -> str:
        """Run LLM judgment (contradiction detection, screening, etc.)."""
        ...



class OpenAIClient:
    """AI client using OpenAI or OpenAI-compatible API.

    Requires: pip install trustmem[openai]
    Includes automatic retry for transient errors (429, 500, 503).
    """

    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
        embed_model: str = "text-embedding-3-small",
        llm_model: str = "gpt-4o-mini",
    ) -> None:
        try:
            from openai import OpenAI
        except ImportError:
            raise ImportError(
                "openai is required. Install with: pip install trustmem[openai]"
            )
        self._client = OpenAI(
            api_key=api_key or os.getenv("OPENAI_API_KEY"),
            base_url=base_url,
        )
        self._embed_model = embed_model
        self._llm_model = llm_model

    @retry_on_transient(max_retries=3, base_delay=1.0)
    def embed(self, text: str) -> list[float]:
        response = self._client.embeddings.create(
            model=self._embed_model,
            input=text,
        )
        return response.data[0].embedding

    @retry_on_transient(max_retries=3, base_delay=1.0)
    def batch_embed(self, texts: list[str]) -> list[list[float]]:
        """Batch embed multiple texts in a single API call."""
        if not texts:
            return []
        response = self._client.embeddings.create(
            model=self._embed_model,
            input=texts,
        )
        # Sort by index to ensure correct ordering
        sorted_data = sorted(response.data, key=lambda d: d.index)
        return [d.embedding for d in sorted_data]

    @retry_on_transient(max_retries=2, base_delay=1.5)
    def llm_judge(self, prompt: str) -> str:
        response = self._client.chat.completions.create(
            model=self._llm_model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.0,
            max_tokens=int(os.getenv("TRUSTMEM_LLM_MAX_TOKENS", "1000")),
        )
        return response.choices[0].message.content or ""


class AzureOpenAIClient:
    """AI client using Azure OpenAI.

    Requires: pip install trustmem[openai]
    Includes automatic retry for transient errors (429, 500, 503).
    """

    def __init__(
        self,
        api_key: str | None = None,
        azure_endpoint: str | None = None,
        api_version: str = "2024-02-01",
        embed_deployment: str = "text-embedding-3-small",
        llm_deployment: str = "gpt-4o-mini",
    ) -> None:
        try:
            from openai import AzureOpenAI
        except ImportError:
            raise ImportError(
                "openai is required. Install with: pip install trustmem[openai]"
            )
        self._client = AzureOpenAI(
            api_key=api_key or os.getenv("AZURE_OPENAI_API_KEY"),
            azure_endpoint=azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT", ""),
            api_version=api_version,
        )
        self._embed_deployment = embed_deployment
        self._llm_deployment = llm_deployment

    @retry_on_transient(max_retries=3, base_delay=1.0)
    def embed(self, text: str) -> list[float]:
        response = self._client.embeddings.create(
            model=self._embed_deployment,
            input=text,
        )
        return response.data[0].embedding

    @retry_on_transient(max_retries=3, base_delay=1.0)
    def batch_embed(self, texts: list[str]) -> list[list[float]]:
        """Batch embed multiple texts in a single API call."""
        if not texts:
            return []
        response = self._client.embeddings.create(
            model=self._embed_deployment,
            input=texts,
        )
        sorted_data = sorted(response.data, key=lambda d: d.index)
        return [d.embedding for d in sorted_data]

    @retry_on_transient(max_retries=2, base_delay=1.5)
    def llm_judge(self, prompt: str) -> str:
        response = self._client.chat.completions.create(
            model=self._llm_deployment,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.0,
            max_tokens=int(os.getenv("TRUSTMEM_LLM_MAX_TOKENS", "1000")),
        )
        return response.choices[0].message.content or ""


def create_ai_client(provider: str = "openai", **kwargs) -> AIClient:
    """Factory function to create an AI client.

    Args:
        provider: One of "openai", "azure", "openai_compat"
        **kwargs: Provider-specific configuration

    Returns:
        An AIClient instance
    """
    if provider == "openai":
        return OpenAIClient(**kwargs)
    elif provider == "azure":
        return AzureOpenAIClient(**kwargs)
    elif provider == "openai_compat":
        return OpenAIClient(**kwargs)
    else:
        raise ValueError(f"Unknown AI provider: {provider}")
