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

import os
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,
    ) -> str:
        """
        Generate text using the LLM.

        Args:
            prompt: User prompt.
            system_prompt: Optional system prompt.
            temperature: Temperature for generation (0.0-2.0).
            max_tokens: Maximum tokens to generate.

        Returns:
            Generated text.
        """
        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        messages.append({"role": "user", "content": prompt})

        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
            )
            return response.choices[0].message.content.strip()
        except Exception as e:
            raise RuntimeError(f"LLM API call failed: {e}") from e

