import os
import json
import gradio as gr
import vertexai
from vertexai.preview.vision_models import ImageGenerationModel
from huggingface_hub import login
from google.oauth2 import service_account # Added for secure in-memory auth

# --- 1. Configuration and Authentication ---

# Load configuration from Hugging Face Secrets instead of hardcoding.
GCP_PROJECT_ID = os.environ.get("GCP_PROJECT_ID")
GCP_LOCATION = os.environ.get("GCP_LOCATION")

# --- Authentication and Sanity Checks Block ---

# Part A: Hugging Face Hub Authentication
hf_token = os.environ.get("IMAGEN_TOKEN")
if hf_token:
    print("Hugging Face token found. Logging in.")
    login(token=hf_token)
else:
    print("WARNING: Hugging Face token ('IMAGEN_TOKEN') not found. Hub-related features may be disabled.")

# Part B: Google Cloud Credentials and Initialization
creds_json_str = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON")

# This global variable will hold a generic error message for the UI if initialization fails.
generic_init_error_message = "FATAL: The image generation service is not configured correctly. Please contact the administrator."

# Check if all necessary secrets are loaded
if not all([GCP_PROJECT_ID, GCP_LOCATION, creds_json_str]):
    missing_secrets = []
    if not GCP_PROJECT_ID: missing_secrets.append("GCP_PROJECT_ID")
    if not GCP_LOCATION: missing_secrets.append("GCP_LOCATION")
    if not creds_json_str: missing_secrets.append("GOOGLE_APPLICATION_CREDENTIALS_JSON")
    
    # Log the detailed error for the administrator
    error_message_for_logs = f"FATAL: The following required secrets are missing: {', '.join(missing_secrets)}. Please set them in the Space settings."
    print(error_message_for_logs)
    
    # Set the global error message to the generic one for the UI
    error_message = generic_init_error_message
    INITIALIZATION_SUCCESS = False
else:
    # This block runs only if all secrets are present.
    print("All required GCP secrets (Project, Location, Credentials) are loaded.")
    
    try:
        # <<< START: SECURITY FIX >>>
        # Securely initialize the Vertex AI client by passing credentials in-memory.
        # This avoids writing sensitive credential files to the container's filesystem.
        
        # 1. Parse the JSON string from the environment variable into a Python dict.
        creds_info = json.loads(creds_json_str)
        
        # 2. Create a credentials object from the dictionary.
        credentials = service_account.Credentials.from_service_account_info(creds_info)
        
        # 3. Initialize Vertex AI with the in-memory credentials object.
        vertexai.init(project=GCP_PROJECT_ID, location=GCP_LOCATION, credentials=credentials)
        # <<< END: SECURITY FIX >>>

        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:
        # Log the detailed, sensitive error for the administrator ONLY.
        print(f"ERROR: Failed to initialize Vertex AI or the model. Exception: {e}")
        # Set the generic error message for the UI. DO NOT expose the raw exception 'e'.
        error_message = generic_init_error_message
        INITIALIZATION_SUCCESS = False

# --- 2. The Core Image Generation Function ---
def generate_image(prompt: str, negative_prompt: str, seed: int):
    """Generates an image using the Vertex AI Imagen model."""
    if not INITIALIZATION_SUCCESS:
        # Provide the clear, generic error message in the UI if initialization failed.
        raise gr.Error(error_message)

    # <<< START: PRIVACY UPDATE >>>
    # Do not log the user's prompt content.
    print("Received new image generation request.")
    # <<< END: PRIVACY UPDATE >>>
    
    try:
        images = generation_model.generate_images(
            prompt=prompt, number_of_images=1, aspect_ratio="1:1",
            negative_prompt=negative_prompt, add_watermark=False,
            safety_filter_level="block_few", person_generation="allow_adult",
            seed=seed
        )
        print("Image generation successful for the request.")
        return images[0]._pil_image
    except Exception as e:
        # <<< START: SECURITY UPDATE >>>
        # Log the detailed error for the administrator's review.
        print(f"An error occurred during image generation: {e}")
        # Raise a generic, user-friendly error in the UI instead of the raw exception.
        raise gr.Error("An unexpected error occurred while generating the image. Please try again or contact support if the issue persists.")
        # <<< END: SECURITY UPDATE >>>

# --- 3. Gradio User Interface Definition ---
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.")
    
    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(label="Prompt", placeholder="A photorealistic image of a futuristic city...", lines=3)
            negative_prompt_input = gr.Textbox(label="Negative Prompt", placeholder="text, watermark, blurry...", lines=2)
            seed_input = gr.Slider(label="Seed", minimum=0, maximum=99999, step=1, randomize=True, info="Controls randomness.")
            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(
        [
            ["An epic oil painting of a lone astronaut discovering a lush, alien jungle.", "blurry"],
            ["A cute, fluffy corgi wearing a tiny chef's hat, in a kitchen, cinematic lighting.", "text"],
        ],
        inputs=[prompt_input, negative_prompt_input]
    )

    submit_button.click(
        fn=generate_image,
        inputs=[prompt_input, negative_prompt_input, seed_input],
        outputs=image_output
    )

# Launch the Gradio app
# <<< START: SECURITY UPDATE >>>
# Disabled debug mode for production to prevent leaking sensitive information through stack traces.
demo.launch(debug=False)
# <<< END: SECURITY UPDATE >>>