"""Instructor client for structured LLM outputs.

Provides type-safe, validated responses from LLMs using Instructor
with automatic retry on validation failures and transient error handling.
"""

from __future__ import annotations

import logging
import os
import time
from typing import Callable, Optional, Type, TypeVar

from pydantic import BaseModel

T = TypeVar("T", bound=BaseModel)

logger = logging.getLogger(__name__)

# Transient HTTP status codes that warrant a retry
_RETRYABLE_STATUS_CODES = {429, 500, 502, 503, 529}

# Default retry configuration
_LLM_MAX_ATTEMPTS = 3
_LLM_BASE_DELAY = 1.0  # seconds
_LLM_MAX_DELAY = 60.0  # seconds


def _is_transient_error(exc: Exception) -> bool:
    """Check if an exception represents a transient LLM API error.

    Detects rate limits (429), server errors (500/502/503/529), and
    timeout/connection errors from Anthropic and OpenAI SDKs.

    Args:
        exc: The exception to check.

    Returns:
        True if the error is transient and should be retried.
    """
    # Check for status_code attribute (Anthropic/OpenAI APIStatusError)
    status_code = getattr(exc, "status_code", None)
    if status_code and status_code in _RETRYABLE_STATUS_CODES:
        return True

    # Check for httpx-based errors (both SDKs use httpx internally)
    try:
        import httpx
        if isinstance(exc, (httpx.TimeoutException, httpx.ConnectError)):
            return True
    except ImportError:
        pass

    # Check exception class names for SDK-specific timeout/connection errors
    exc_type = type(exc).__name__
    if exc_type in (
        "APITimeoutError",
        "APIConnectionError",
        "InternalServerError",
        "RateLimitError",
        "OverloadedError",
    ):
        return True

    # Check for generic timeout/connection errors
    if isinstance(exc, (TimeoutError, ConnectionError, OSError)):
        return True

    return False


def _get_retry_delay(attempt: int, exc: Exception) -> float:
    """Calculate retry delay with exponential backoff and jitter.

    Respects Retry-After headers when present (e.g. from 429 responses).

    Args:
        attempt: The current attempt number (0-indexed).
        exc: The exception that triggered the retry.

    Returns:
        Delay in seconds before the next retry.
    """
    # Respect Retry-After header if present
    retry_after = getattr(exc, "retry_after", None)
    if retry_after is not None:
        try:
            return min(float(retry_after), _LLM_MAX_DELAY)
        except (ValueError, TypeError):
            pass

    # Check response headers for Retry-After
    response = getattr(exc, "response", None)
    if response is not None:
        headers = getattr(response, "headers", {})
        retry_after_header = headers.get("retry-after") or headers.get("Retry-After")
        if retry_after_header:
            try:
                return min(float(retry_after_header), _LLM_MAX_DELAY)
            except (ValueError, TypeError):
                pass

    # Exponential backoff: base_delay * 2^attempt, capped at max_delay
    import random
    delay = _LLM_BASE_DELAY * (2 ** attempt)
    # Add jitter (0.5x to 1.5x)
    delay *= 0.5 + random.random()
    return min(delay, _LLM_MAX_DELAY)


def _call_with_llm_retry(fn: Callable[[], T], max_attempts: int = _LLM_MAX_ATTEMPTS) -> T:
    """Call a function with retry logic for transient LLM API errors.

    Retries on rate limits (429), server errors (500/502/503/529),
    timeouts, and connection errors with exponential backoff.

    Args:
        fn: Zero-argument callable that makes the LLM API call.
        max_attempts: Maximum number of attempts (including the first).

    Returns:
        The result of the function call.

    Raises:
        The last exception if all retries are exhausted.
    """
    last_exc: Optional[Exception] = None

    for attempt in range(max_attempts):
        try:
            return fn()
        except Exception as exc:
            if not _is_transient_error(exc) or attempt >= max_attempts - 1:
                raise
            last_exc = exc
            delay = _get_retry_delay(attempt, exc)
            logger.warning(
                "LLM API call failed (attempt %d/%d), retrying in %.1fs: %s",
                attempt + 1,
                max_attempts,
                delay,
                str(exc),
            )
            time.sleep(delay)

    # Should not reach here, but satisfy type checker
    raise last_exc  # type: ignore[misc]


class InstructorClient:
    """LLM client with Instructor for structured outputs.

    Supports Anthropic, OpenAI, and DeepSeek providers with automatic
    retry on validation failures.
    """

    def __init__(
        self,
        provider: str = "anthropic",
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
    ):
        """Initialize the Instructor client.

        Args:
            provider: LLM provider ("anthropic", "openai", or "deepseek")
            api_key: Optional API key (defaults to environment variable)
            base_url: Optional base URL for OpenAI-compatible endpoints
        """
        self.provider = provider
        self.api_key = api_key
        self.base_url = base_url
        self._client = None
        self._instructor_available: Optional[bool] = None

    @property
    def instructor_available(self) -> bool:
        """Check if Instructor is available."""
        if self._instructor_available is None:
            try:
                import instructor
                self._instructor_available = True
            except ImportError:
                self._instructor_available = False
        return self._instructor_available

    @property
    def client(self):
        """Get or create the instructor-patched client."""
        if self._client is None:
            self._client = self._create_client()
        return self._client

    def _create_client(self):
        """Create instructor-patched client for the provider."""
        if not self.instructor_available:
            raise ImportError(
                "Instructor is not installed. Install with: pip install instructor"
            )

        import instructor

        if self.provider == "anthropic":
            try:
                import anthropic
            except ImportError:
                raise ImportError(
                    "anthropic package not installed. Install with: pip install anthropic"
                )

            api_key = self.api_key or os.environ.get("ANTHROPIC_API_KEY")
            if not api_key:
                raise ValueError(
                    "ANTHROPIC_API_KEY environment variable not set"
                )

            return instructor.from_anthropic(anthropic.Anthropic(api_key=api_key))

        elif self.provider in ("openai", "deepseek"):
            try:
                import openai
            except ImportError:
                raise ImportError(
                    "openai package not installed. Install with: pip install openai"
                )

            if self.provider == "deepseek":
                api_key = self.api_key or os.environ.get("DEEPSEEK_API_KEY") or os.environ.get("HIKAFLOW_LLM_API_KEY")
                base_url = self.base_url or os.environ.get("HIKAFLOW_LLM_BASE_URL") or "https://api.deepseek.com"
            else:
                api_key = self.api_key or os.environ.get("OPENAI_API_KEY")
                base_url = self.base_url or os.environ.get("HIKAFLOW_LLM_BASE_URL")

            if not api_key:
                raise ValueError("API key not set for provider: " + self.provider)

            client_kwargs = {"api_key": api_key}
            if base_url:
                client_kwargs["base_url"] = base_url

            return instructor.from_openai(openai.OpenAI(**client_kwargs))

        else:
            raise ValueError(f"Unknown provider: {self.provider}")

    def generate(
        self,
        prompt: str,
        response_model: Type[T],
        model: Optional[str] = None,
        max_retries: int = 2,
        temperature: float = 0.3,
        max_tokens: int = 4096,
        system_prompt: Optional[str] = None,
    ) -> T:
        """Generate structured response with automatic retry.

        Retries on transient errors (rate limits, server errors, timeouts)
        with exponential backoff. Validation retries are handled separately
        by the Instructor library via max_retries.

        Args:
            prompt: The prompt to send to the LLM
            response_model: Pydantic model for the response
            model: Model name (defaults based on provider)
            max_retries: Number of retries on validation failure
            temperature: Sampling temperature (0-1)
            max_tokens: Maximum tokens in response
            system_prompt: Optional system-level instructions that set the
                LLM's role and constraints.  For Anthropic this is passed
                as the ``system`` parameter; for OpenAI it is prepended as
                a ``{"role": "system", ...}`` message (C3 -- no system prompt).

        Returns:
            Validated Pydantic model instance
        """
        if model is None:
            model = self._default_model()

        if self.provider == "anthropic":
            def _call_anthropic():
                kwargs: dict = dict(
                    model=model,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    messages=[{"role": "user", "content": prompt}],
                    response_model=response_model,
                    max_retries=max_retries,
                )
                if system_prompt:
                    kwargs["system"] = system_prompt
                return self.client.messages.create(**kwargs)
            return _call_with_llm_retry(_call_anthropic)
        elif self.provider in ("openai", "deepseek"):
            def _call_openai():
                messages = []
                if system_prompt:
                    messages.append({"role": "system", "content": system_prompt})
                messages.append({"role": "user", "content": prompt})
                return self.client.chat.completions.create(
                    model=model,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    messages=messages,
                    response_model=response_model,
                    max_retries=max_retries,
                )
            return _call_with_llm_retry(_call_openai)
        else:
            raise ValueError(f"Unknown provider: {self.provider}")

    def _default_model(self) -> str:
        """Get default model for the provider."""
        defaults = {
            "anthropic": "claude-sonnet-4-20250514",
            "openai": "gpt-5",
            "deepseek": "deepseek-chat",
        }
        return defaults.get(self.provider, "deepseek-chat")

    def generate_hypothesis(
        self,
        prompt: str,
        model: Optional[str] = None,
        max_retries: int = 2,
    ) -> HypothesisResponse:
        """Generate a single hypothesis with structured output.

        Args:
            prompt: The prompt describing the error and context
            model: Optional model override
            max_retries: Number of retries on validation failure

        Returns:
            HypothesisResponse with analysis, fix, and verification
        """
        from autodebug.prompts.structured_output import HypothesisResponse

        return self.generate(
            prompt=prompt,
            response_model=HypothesisResponse,
            model=model,
            max_retries=max_retries,
            temperature=0.3,
        )

    def generate_hypotheses(
        self,
        prompt: str,
        num_hypotheses: int = 3,
        model: Optional[str] = None,
        max_retries: int = 2,
    ) -> MultiHypothesisResponse:
        """Generate multiple diverse hypotheses.

        Args:
            prompt: The prompt describing the error and context
            num_hypotheses: Number of hypotheses to request
            model: Optional model override
            max_retries: Number of retries on validation failure

        Returns:
            MultiHypothesisResponse with list of hypotheses
        """
        from autodebug.prompts.structured_output import MultiHypothesisResponse

        # Modify prompt to request multiple hypotheses
        multi_prompt = f"""{prompt}

Generate exactly {num_hypotheses} different fix hypotheses, ordered by confidence.
Each hypothesis should use a different approach or address a different possible root cause."""

        return self.generate(
            prompt=multi_prompt,
            response_model=MultiHypothesisResponse,
            model=model,
            max_retries=max_retries,
            temperature=0.5,  # Slightly higher for diversity
            max_tokens=8192,  # More tokens for multiple hypotheses
        )


def create_instructor_client(
    provider: Optional[str] = None,
) -> InstructorClient:
    """Create an InstructorClient with auto-detected provider.

    Checks for available API keys and creates client for the
    first available provider. Prefers DeepSeek when available.

    Args:
        provider: Optional provider override

    Returns:
        InstructorClient instance
    """
    if provider:
        return InstructorClient(provider=provider)

    # Auto-detect: prefer DeepSeek (cheapest), then Anthropic, then OpenAI
    if os.environ.get("DEEPSEEK_API_KEY") or os.environ.get("HIKAFLOW_LLM_PROVIDER", "").lower() == "deepseek":
        return InstructorClient(provider="deepseek")
    elif os.environ.get("ANTHROPIC_API_KEY"):
        return InstructorClient(provider="anthropic")
    elif os.environ.get("OPENAI_API_KEY"):
        return InstructorClient(provider="openai")
    else:
        raise ValueError(
            "No API key found. Set DEEPSEEK_API_KEY, ANTHROPIC_API_KEY, or OPENAI_API_KEY"
        )
