"""LLM client wrapper for GuardRail."""

import os
import time
from typing import Optional

from openai import OpenAI


class LLMClient:
    """Wrapper for LLM API calls."""

    def __init__(
        self,
        api_key: Optional[str] = None,
        model: str = "gpt-4o-mini",
        base_url: Optional[str] = None,
    ):
        """
        Initialize LLM client.

        Args:
            api_key: OpenAI API key. If None, reads from OPENAI_API_KEY env var or config.
            model: Model name to use (default: gpt-4o-mini).
            base_url: Optional base URL for API (for custom endpoints).
        """
        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
        if not self.api_key:
            raise ValueError(
                "OpenAI API key not provided. Set OPENAI_API_KEY environment variable, "
                "configure it via 'guardrail init', or pass api_key parameter."
            )

        self.model = model
        self.client = OpenAI(api_key=self.api_key, base_url=base_url)

    def generate(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        temperature: float = 0.3,
        max_tokens: int = 2000,
        max_retries: int = 3,
        retry_delay: float = 1.0,
        timeout: int = 60,
    ) -> str:
        """
        Generate text using the LLM with retry logic and timeout.

        Args:
            prompt: User prompt.
            system_prompt: Optional system prompt.
            temperature: Temperature for generation (0.0-2.0).
            max_tokens: Maximum tokens to generate.
            max_retries: Maximum number of retry attempts.
            retry_delay: Initial delay between retries (exponential backoff).
            timeout: Request timeout in seconds.

        Returns:
            Generated text.

        Raises:
            RuntimeError: If all retry attempts fail.
        """
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        last_error = None
        for attempt in range(max_retries):
            try:
                response = self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    temperature=temperature,
                    max_tokens=max_tokens,
                    timeout=timeout,
                )
                return response.choices[0].message.content.strip()
            except Exception as e:
                last_error = e
                error_str = str(e).lower()

                # Check if it's a rate limit error
                if "rate limit" in error_str or "429" in error_str:
                    if attempt < max_retries - 1:
                        wait_time = retry_delay * (2**attempt)  # Exponential backoff
                        time.sleep(wait_time)
                        continue
                # Check if it's a timeout error
                elif "timeout" in error_str or "timed out" in error_str:
                    if attempt < max_retries - 1:
                        time.sleep(retry_delay)
                        continue
                # For other errors, retry once
                elif attempt < max_retries - 1:
                    time.sleep(retry_delay)
                    continue

                # If we've exhausted retries, raise the error
                if attempt == max_retries - 1:
                    raise RuntimeError(
                        f"LLM API call failed after {max_retries} attempts: {e}"
                    ) from e

        # Should never reach here, but just in case
        raise RuntimeError(f"LLM API call failed: {last_error}") from last_error
