import requests

class AIClient:
    """
    Client for interacting with Google Gemini AI models
    using an OpenAI-compatible API endpoint.

    This client sends prompts to the AI model and
    returns generated responses.
    """

    def __init__(self, api_key: str, model: str = "gemini-2.0-flash"):
        """
        Initialize the AIClient.

        Parameters:
            api_key (str): Google Gemini API key.
            model (str): Model name to use (default: gemini-2.0-flash).

        Raises:
            ValueError: If API key is not provided.
        """
        if not api_key:
            raise ValueError("API key required. Pass api_key='YOUR_KEY'.")

        self.api_key = api_key
        self.model = model
        self.url = "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions"

    def get_response(self, prompt: str, max_tokens: int = 300) -> str:
        """
        Send a prompt to the AI model and return the generated response.

        Parameters:
            prompt (str): The input message to send to the AI.
            max_tokens (int): Maximum number of tokens in the response.

        Returns:
            str: AI-generated response text.

        Raises:
            Exception: If the API request fails.
        """

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        body = {
            "model": self.model,
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens
        }

        response = requests.post(self.url, json=body, headers=headers)

        if response.status_code != 200:
            raise Exception(f"API request failed: {response.text}")

        data = response.json()
        return data["choices"][0]["message"]["content"].strip()     