# app.py

import os
import gradio as gr
import vertexai
from vertexai.preview.vision_models import ImageGenerationModel

# --- 1. Authentication and Initialization ---
# This section handles secure authentication with Google Cloud.
# It expects the GCP credentials to be stored in a Hugging Face Secret
# named "GOOGLE_APPLICATION_CREDENTIALS_JSON".

# Check if the secret is available
if "GOOGLE_APPLICATION_CREDENTIALS_JSON" in os.environ:
    # Get the JSON content from the environment variable
    gcp_creds_json_str = os.environ["GOOGLE_APPLICATION_CREDENTIALS_JSON"]
    
    # Write the JSON content to a temporary file
    with open("gcp_creds.json", "w") as f:
        f.write(gcp_creds_json_str)
        
    # Set the environment variable to point to the temporary file
    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "gcp_creds.json"
    print("GCP credentials loaded successfully from Hugging Face secrets.")
else:
    # This block will run if the secret is not set.
    # The app will likely fail to authenticate with Vertex AI.
    print("WARNING: GCP credentials secret not found. The app may not work.")
    print("Please set the 'GOOGLE_APPLICATION_CREDENTIALS_JSON' secret in your Space settings.")

# Initialize Vertex AI
# IMPORTANT: Replace with your Google Cloud Project ID and desired location.
GCP_PROJECT_ID = "gen-lang-client-0193353123" 
GCP_LOCATION = "us-central1"

try:
    vertexai.init(project=GCP_PROJECT_ID, location=GCP_LOCATION)
    # Load the pre-trained model (this is done only once when the app starts)
    generation_model = ImageGenerationModel.from_pretrained("imagen-4.0-generate-preview-06-06")
    print("Vertex AI and Imagen Model initialized successfully.")
    INITIALIZATION_SUCCESS = True
except Exception as e:
    print(f"ERROR: Failed to initialize Vertex AI or the model. {e}")
    INITIALIZATION_SUCCESS = False


# --- 2. The Core Image Generation Function ---
# This function takes user inputs and returns a generated image.

def generate_image(prompt: str, negative_prompt: str, seed: int):
    """
    Generates an image using the Vertex AI Imagen model.
    """
    if not INITIALIZATION_SUCCESS:
        raise gr.Error("Model not initialized. Check server logs for authentication issues.")

    print(f"Received request with prompt: '{prompt}'")
    try:
        # The API expects a seed value between 0 and 2147483647
        # We ensure the Gradio input (0-99999) fits this range.
        api_seed = seed
        
        images = generation_model.generate_images(
            prompt=prompt,
            number_of_images=1,  # We generate one image for the web UI
            aspect_ratio="1:1",
            negative_prompt=negative_prompt,
            add_watermark=False,
            safety_filter_level="block_few",
            person_generation="allow_adult",
            seed=api_seed,
        )
        
        print("Image generated successfully.")
        # The SDK returns a list of objects; we get the PIL image from the first one.
        return images[0]._pil_image
        
    except Exception as e:
        print(f"An error occurred during image generation: {e}")
        # Provide a user-friendly error message in the UI
        raise gr.Error(f"Failed to generate image. Error: {e}")


# --- 3. Gradio User Interface Definition ---
# This section defines the web interface using Gradio.

with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container {max-width: 960px !important}") as demo:
    gr.Markdown("# 🎨 Vertex AI Imagen 4.0 Image Generator")
    gr.Markdown(
        "Generate high-quality images with Google's latest Imagen 4.0 model. "
        "Describe what you want to see in the prompt box below."
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(
                label="Prompt", 
                placeholder="A photorealistic image of a futuristic city at sunset, with flying cars and neon lights.",
                lines=3
            )
            negative_prompt_input = gr.Textbox(
                label="Negative Prompt (what to avoid)", 
                placeholder="text, watermark, blurry, cartoon, ugly",
                lines=2
            )
            seed_input = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=99999,
                step=1,
                randomize=True,
                info="A seed controls randomness. Different seeds produce different images for the same prompt."
            )
            submit_button = gr.Button("Generate Image", variant="primary")
            
        with gr.Column(scale=1):
            image_output = gr.Image(label="Generated Image", type="pil", interactive=False)

    gr.Examples(
        examples=[
            ["A high-detail, epic oil painting of a lone astronaut discovering a lush, alien jungle.", "blurry, low quality"],
            ["A cute, fluffy corgi wearing a tiny chef's hat, sitting in a kitchen, cinematic lighting.", "human, text, watermark"],
            ["A minimalist black and white logo for a coffee shop named 'The Grind', vector art.", "photograph, complex"],
        ],
        inputs=[prompt_input, negative_prompt_input]
    )

    # Connect the button click to the generation function
    submit_button.click(
        fn=generate_image,
        inputs=[prompt_input, negative_prompt_input, seed_input],
        outputs=image_output
    )

# Launch the Gradio app
demo.launch(debug=True)