import os
from typing import Optional
from google import genai
from google.genai import types

API_KEY = "[Your API Key]"
Model = "gemini-2.5-pro-preview-05-06"

class AIClient:
    
    def __init__(self, api_key=API_KEY, model=Model):
        self.api_key = api_key or os.environ.get("GOOGLE_API_KEY")
        if not self.api_key:
            raise ValueError("API key not provided!")
        
        # Configure client
        self.client = genai.Client(api_key=self.api_key)
        self.model = model
        
    def generate(self, user_prompt: str, system_prompt: str = "", model: Optional[str] = None) -> str:
        """
        Generate text response
        
        Parameters:
            user_prompt: User prompt
            system_prompt: System prompt
            model: Model to use, if not specified the default model will be used
            
        Returns:
            Generated text
        """
        model_to_use = model or self.model
        
        # Build request content
        contents = []
        
        # Add system prompt
        if system_prompt:
            contents.append({
                "role": "system",
                "parts": [{"text": system_prompt}]
            })
        
        # Add user prompt
        contents.append({
            "role": "user", 
            "parts": [{"text": user_prompt}]
        })
        
        try:
            # Call Google API to generate content
            response = self.client.models.generate_content(
                model=model_to_use,
                contents=contents,
                config=types.GenerateContentConfig(
                    thinking_config=types.ThinkingConfig(thinking_budget=16384)
                )
            )
            
            # Get generated text
            return response.text
            
        except Exception as e:
            raise Exception(f"Failed to generate content: {str(e)}")