import gradio as gr
from google import genai
from google.genai import types
from PIL import Image
from io import BytesIO
import os
from dotenv import load_dotenv
import warnings

# Suppress specific warnings if needed (optional)
warnings.filterwarnings("ignore", message="IMAGE_SAFETY is not a valid FinishReason")

# Load environment variables
load_dotenv()
api_key = os.environ.get('GEMINI_API_KEY')

# Initialize the client
client = genai.Client(api_key=api_key)

# Constants for metrics
MODEL_NAME = "gemini-2.0-flash-exp-image-generation"
PRICE_PER_INPUT_TOKEN = 0.0000  # $ per token (example)
PRICE_PER_OUTPUT_TOKEN = 0.0000  # $ per token (example)

def estimate_tokens(text):
    """Rough estimation of tokens (words + some overhead)"""
    return len(str(text).split()) + 10

def generate_image(prompt):
    """Generate an image based on text prompt with metrics"""
    try:
        response = client.models.generate_content(
            model=MODEL_NAME,
            contents=prompt,
            config=types.GenerateContentConfig(
                response_modalities=['Text', 'Image']
            )
        )
        
        # Check if response was blocked or incomplete
        if not response.candidates:
            return "Content generation blocked or failed", f"Model: {MODEL_NAME}\nReason: No candidates returned"
        
        # Estimate tokens
        input_tokens = estimate_tokens(prompt)
        output_tokens = 100  # Fixed estimate for image output
        
        # Calculate price
        total_price = (input_tokens * PRICE_PER_INPUT_TOKEN) + (output_tokens * PRICE_PER_OUTPUT_TOKEN)
        
        # Get image
        image = None
        for part in response.candidates[0].content.parts:
            if part.inline_data is not None:
                image = Image.open(BytesIO(part.inline_data.data))
                break
        
        if image:
            metrics = f"Model: {MODEL_NAME}\nInput Tokens: {input_tokens}\nOutput Tokens: {output_tokens}\nTotal Price: ${total_price:.4f}"
            return image, metrics
        return "No image generated", f"Model: {MODEL_NAME}\nReason: No image data in response"
    
    except Exception as e:
        return f"Error generating image: {str(e)}", f"Model: {MODEL_NAME}\nError: {str(e)}"

def edit_uploaded_image(image, edit_prompt):
    """Edit the uploaded image based on the prompt with metrics"""
    try:
        if image is None:
            return "No image uploaded", "No metrics available"
        
        # Convert Gradio image input to PIL Image
        pil_image = Image.fromarray(image)
        
        # Prepare content
        text_input = f"Hi, this is an image. {edit_prompt}"
        contents = [text_input, pil_image]
        
        # Generate edited image
        response = client.models.generate_content(
            model=MODEL_NAME,
            contents=contents,
            config=types.GenerateContentConfig(
                response_modalities=['Text', 'Image']
            )
        )
        
        # Check if response was blocked or incomplete
        if not response.candidates:
            return "Content generation blocked or failed", f"Model: {MODEL_NAME}\nReason: No candidates returned"
        
        # Estimate tokens
        input_tokens = estimate_tokens(text_input) + 100  # Add approximate tokens for image
        output_tokens = 100  # Fixed estimate for image output
        
        # Calculate price
        total_price = (input_tokens * PRICE_PER_INPUT_TOKEN) + (output_tokens * PRICE_PER_OUTPUT_TOKEN)
        
        # Get edited image
        edited_image = None
        for part in response.candidates[0].content.parts:
            if part.inline_data is not None:
                edited_image = Image.open(BytesIO(part.inline_data.data))
                break
        
        if edited_image:
            metrics = f"Model: {MODEL_NAME}\nInput Tokens: {input_tokens}\nOutput Tokens: {output_tokens}\nTotal Price: ${total_price:.4f}"
            return edited_image, metrics
        return "No edited image generated", f"Model: {MODEL_NAME}\nReason: No image data in response"
    
    except Exception as e:
        return f"Error editing image: {str(e)}", f"Model: {MODEL_NAME}\nError: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="Image Generation and Editing App") as app:
    gr.Markdown("# Image Generation and Editing App")
    gr.Markdown("Generate or edit images with usage metrics")
    
    with gr.Tab("Generate Image"):
        prompt_input = gr.Textbox(
            label="Enter your prompt",
            placeholder="e.g., 'A pig with wings and a top hat flying over a futuristic city'"
        )
        generate_button = gr.Button("Generate")
        with gr.Row():
            generated_output = gr.Image(label="Generated Image")
            generate_metrics = gr.Textbox(label="Metrics", interactive=False)
        
    with gr.Tab("Edit Uploaded Image"):
        image_upload = gr.Image(label="Upload your image")
        edit_prompt = gr.Textbox(
            label="Edit prompt",
            placeholder="e.g., 'Add a llama next to the subject'"
        )
        edit_button = gr.Button("Edit Image")
        with gr.Row():
            edited_output = gr.Image(label="Edited Image")
            edit_metrics = gr.Textbox(label="Metrics", interactive=False)
    
    # Connect the components with functions
    generate_button.click(
        fn=generate_image,
        inputs=prompt_input,
        outputs=[generated_output, generate_metrics]
    )
    
    edit_button.click(
        fn=edit_uploaded_image,
        inputs=[image_upload, edit_prompt],
        outputs=[edited_output, edit_metrics]
    )

# Launch the app
# app.launch(share=True)
app.launch()