app.py

import os
import io
import time
import base64
import uuid
import PIL.Image
from flask import Flask, render_template, request, jsonify
from dotenv import load_dotenv

# Google Cloud & GenAI specific imports
from google.cloud import storage
from google.api_core import exceptions as google_exceptions
from google import genai
from google.genai import types

# --- Configuration & Initialization ---
# load_dotenv('.env')

app = Flask(__name__)

LOCAL_IMAGE_DIR = os.path.join('static', 'generated_images')
os.makedirs(LOCAL_IMAGE_DIR, exist_ok=True)

# Gemini Image Generation Client (using your existing setup)
API_KEY = os.environ.get("GOOGLE_API_KEY")
MODEL_ID_IMAGE = 'gemini-2.0-flash-exp-image-generation'

# Veo Video Generation Client (NEW)
PROJECT_ID = os.environ.get("PROJECT_ID")
LOCATION = os.environ.get("GOOGLE_CLOUD_REGION", "us-central1")
GCS_BUCKET_NAME = os.environ.get("GCS_BUCKET_NAME")
MODEL_ID_VIDEO = "veo-3.0-generate-preview" # Your Veo model ID

if not all([API_KEY, PROJECT_ID, GCS_BUCKET_NAME, LOCATION]):
    raise RuntimeError("Missing required environment variables. Check your .env file.")

# Initialize clients
try:
    # Client for Gemini Image Generation
    gemini_image_client = genai.Client(api_key=API_KEY)
    print(f"Gemini Image Client initialized successfully for model: {MODEL_ID_IMAGE}")

    # Client for Veo Video Generation (Vertex AI)
    veo_video_client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)
    print(f"Veo Video Client (Vertex AI) initialized successfully for project: {PROJECT_ID}")

    # Client for Google Cloud Storage
    gcs_client = storage.Client(project=PROJECT_ID)
    print("Google Cloud Storage Client initialized successfully.")

except Exception as e:
    print(f"Error during client initialization: {e}")
    gemini_image_client = veo_video_client = gcs_client = None


# --- Helper Function to Upload to GCS (NEW) ---
def upload_bytes_to_gcs(image_bytes: bytes, bucket_name: str, destination_blob_name: str) -> str:
    """Uploads image bytes to GCS and returns the GCS URI."""
    if not gcs_client:
        raise ConnectionError("GCS client is not initialized.")
    
    bucket = gcs_client.bucket(bucket_name)
    blob = bucket.blob(destination_blob_name)
    blob.upload_from_string(image_bytes, content_type='image/png')
    
    gcs_uri = f"gs://{bucket_name}/{destination_blob_name}"
    print(f"Image successfully uploaded to {gcs_uri}")
    return gcs_uri


# --- Main Routes ---
@app.route('/')
def index():
    """Renders the main HTML page."""
    return render_template('index.html')

@app.route('/generate', methods=['POST'])
def generate_video_from_sketch():
    """Full pipeline: sketch -> image -> video."""
    if not all([gemini_image_client]):
    # if not all([gemini_image_client, veo_video_client, gcs_client]):
        return jsonify({"error": "A server-side client is not initialized. Check server logs."}), 500

    if not request.json or 'image_data' not in request.json:
        return jsonify({"error": "Missing image_data in request"}), 400

    base64_image_data = request.json['image_data']
    user_prompt = request.json.get('prompt', '').strip()

    # --- Step 1: Generate Image with Gemini ---
    try:
        print("--- Step 1: Generating image from sketch with Gemini ---")
        if ',' in base64_image_data:
            base64_data = base64_image_data.split(',', 1)[1]
        else:
            base64_data = base64_image_data
        
        image_bytes = base64.b64decode(base64_data)
        sketch_pil_image = PIL.Image.open(io.BytesIO(image_bytes))

        # default_prompt = "Create a photorealistic image based on this sketch. Focus on realistic lighting, textures, and shadows to make it look like a photograph taken with a professional DSLR camera."
        default_prompt = "Convert this sketch into a photorealistic image as if it were taken from a real DSLR camera. The elements and objects should look real."
        #prompt_text = f"{default_prompt} {user_prompt}" if user_prompt else default_prompt

        response = gemini_image_client.models.generate_content(
            model=MODEL_ID_IMAGE,
            contents=[default_prompt, sketch_pil_image],
            config=types.GenerateContentConfig(response_modalities=['TEXT', 'IMAGE'])
        )

        if not response.candidates:
            raise ValueError("Gemini image generation returned no candidates.")

        generated_image_bytes = None
        for part in response.candidates[0].content.parts:
            if part.inline_data and part.inline_data.mime_type.startswith('image/'):
                generated_image_bytes = part.inline_data.data
                break
        
        if not generated_image_bytes:
            raise ValueError("Gemini did not return an image in the response.")
        
        print("Image generated successfully.")

        try:
            # Use a unique filename to prevent overwrites
            local_filename = f"generated-image-{uuid.uuid4()}.png"
            local_image_path = os.path.join(LOCAL_IMAGE_DIR, local_filename)
            # Write the bytes to a file in binary mode ('wb')
            with open(local_image_path, "wb") as f:
                f.write(generated_image_bytes)
            print(f"Image also saved locally to: {local_image_path}")
        except Exception as e:
            # This is not a critical error, so we just print a warning and continue.
            print(f"[Warning] Could not save image locally: {e}")

    except Exception as e:
        print(f"Error during Gemini image generation: {e}")
        return jsonify({"error": f"Failed to generate image: {e}"}), 500

    # --- Step 2 & 3: Upload Image to GCS and Generate Video with Veo ---
    try:
        print("\n--- Step 2: Uploading generated image to GCS ---")
        unique_id = uuid.uuid4()
        image_blob_name = f"images/generated-image-{unique_id}.png"
        output_gcs_prefix = f"gs://{GCS_BUCKET_NAME}/videos/" # Folder for video outputs

        image_gcs_uri = upload_bytes_to_gcs(generated_image_bytes, GCS_BUCKET_NAME, image_blob_name)
        
        print("\n--- Step 3: Calling Veo to generate video ---")
        default_video_prompt = "Animate this image. Add subtle, cinematic motion."
        video_prompt = f"{user_prompt}" if user_prompt else default_video_prompt
        print(video_prompt)
        
        operation = veo_video_client.models.generate_videos(
            model=MODEL_ID_VIDEO,
            prompt=video_prompt,
            image=types.Image(gcs_uri=image_gcs_uri, mime_type="image/png"),
            config=types.GenerateVideosConfig(
                aspect_ratio="16:9", 
                output_gcs_uri=output_gcs_prefix,
                duration_seconds=8,
                person_generation="allow_adult",
                enhance_prompt=True,
                generate_audio=True, # Keep it simple for now
            ),
        )

       
        # WARNING: This is a synchronous poll, which will block the server thread.
        # For production, consider an asynchronous pattern (e.g., websockets or long polling).
        timeout_seconds = 300 # 5 minutes
        start_time = time.time()
        while not operation.done:
            if time.time() - start_time > timeout_seconds:
                raise TimeoutError("Video generation timed out.")
            time.sleep(15)
            # You must get the operation object again to refresh its status
            operation = veo_video_client.operations.get(operation)
            print(operation)

        print("Video generation operation complete.")
        
        if not operation.response or not operation.result.generated_videos:
            raise ValueError("Veo operation completed but returned no video.")

        video_gcs_uri = operation.result.generated_videos[0].video.uri
        print(f"Video saved to GCS at: {video_gcs_uri}")
        
        # Convert gs:// URI to public https:// URL
        video_blob_name = video_gcs_uri.replace(f"gs://{GCS_BUCKET_NAME}/", "")

        public_video_url = f"https://storage.googleapis.com/{GCS_BUCKET_NAME}/{video_blob_name}"
        print(f"Video generated successfully. Public URL: {public_video_url}")

        return jsonify({"generated_video_url": public_video_url})

    except Exception as e:
        print(f"An error occurred during video generation: {e}")
        return jsonify({"error": f"Failed to generate video: {e}"}), 500


if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)