from google import genai
from google.genai import types
import time
import logging
from typing import Optional
from dotenv import load_dotenv
import os
load_dotenv()

client = genai.Client(api_key=os.getenv("GOOGLE_API_KEY"))

def get_response(prompt: str) -> str:
    response = client.models.generate_content(
        model="gemini-2.5-flash-preview-05-20", contents=prompt
    )
    return(response.text)

def generate_with_gemini(
    prompt: str,
    purpose: str,
    max_retries: int = 5,
    retry_delay: float = 1.0
) -> Optional[str]:
    """
    Generate text using Gemini API with fallback and retry logic.
    
    Args:
        prompt (str): The input prompt for text generation
        api_key (str): Your Gemini API key
        max_retries (int): Maximum number of retry attempts (default: 5)
        retry_delay (float): Initial delay between retries in seconds (default: 1.0)
    
    Returns:
        Optional[str]: Generated text response or None if all attempts fail
    """
    
    print(f"Purpose: {purpose}")
    # Model configurations with their respective max output tokens
    models_config = [
        {
            "name": "gemini-2.5-flash-preview-05-20",
            "max_output_tokens": 65536,
            "description": "Gemini 2.5 Flash Preview"
        },
        {
            "name": "gemini-2.0-flash",
            "max_output_tokens": 8192,
            "description": "Gemini 2.0 Flash"
        }
    ]
    
    # Try each model with retry logic
    for model_config in models_config:
        model_name = model_config["name"]
        max_tokens = model_config["max_output_tokens"]
        model_desc = model_config["description"]
        
        print(f"Attempting to use {model_desc} ({model_name})...")
        
        for attempt in range(max_retries):
            try:
                # Create generation config
                config = types.GenerateContentConfig(
                    max_output_tokens=max_tokens,
                    temperature=0.0  
                )
                
                # Make the API call
                response = client.models.generate_content(
                    model=model_name,
                    contents=[prompt],
                    config=config
                )
                
                # Check if response has content
                if response and response.text:
                    print(f"✅ Success with {model_desc} on attempt {attempt + 1}")
                    return response.text
                else:
                    print(f"⚠️  Empty response from {model_desc} on attempt {attempt + 1}")
                    
            except Exception as e:
                print(f"❌ Error with {model_desc} on attempt {attempt + 1}: {str(e)}")
                
                # If this is not the last attempt, wait before retrying
                if attempt < max_retries - 1:
                    wait_time = retry_delay * (2 ** attempt)  # Exponential backoff
                    print(f"⏳ Waiting {wait_time:.1f} seconds before retry...")
                    time.sleep(wait_time)
                else:
                    print(f"💥 All {max_retries} attempts failed for {model_desc}")
                    break
    
    print("❌ All models and retry attempts exhausted. Unable to generate response.")
    return None