app.py

"""
Marketing Image Generator with Gradio MCP Server
Professional AI image generation using Google Imagen3 with marketing review
Deployed on HuggingFace Spaces with built-in MCP server support
"""

import gradio as gr
import os
import logging
import json
import base64
import asyncio
from typing import Dict, Any, Tuple
from PIL import Image
import io

# Google Service Account Authentication Setup
def setup_google_credentials():
    """Setup Google credentials from service account JSON"""
    try:
        service_account_json = os.getenv("GOOGLE_SERVICE_ACCOUNT_JSON")
        if service_account_json:
            import tempfile
            from google.oauth2 import service_account
            
            # Parse the JSON credentials
            credentials_dict = json.loads(service_account_json)
            
            # Create credentials from service account info
            credentials = service_account.Credentials.from_service_account_info(credentials_dict)
            
            # Set the credentials in environment
            with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
                json.dump(credentials_dict, f)
                os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = f.name
            
            print("✅ Google Cloud service account configured")
            return True
    except Exception as e:
        print(f"⚠️ Google Cloud service account setup failed: {e}")
    
    print("⚠️ Google Cloud service account not found")
    return False

# Setup Google credentials on startup
setup_google_credentials()

# Google AI imports
try:
    import google.generativeai as genai
    from google import genai as genai_sdk
    GEMINI_AVAILABLE = True
except ImportError:
    GEMINI_AVAILABLE = False

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Get API keys - prioritize HuggingFace Secrets
GCP_KEYS = [
    # Hugging Face Secrets (these are the primary ones for HF deployment)
    os.getenv("GOOGLE_API_KEY"),
    os.getenv("GEMINI_API_KEY"), 
    os.getenv("GCP_API_KEY"),
    # Local development keys (fallback for local testing)
    os.getenv("GCP_KEY_1"),
    os.getenv("GCP_KEY_2"), 
    os.getenv("GCP_KEY_3"),
    os.getenv("GCP_KEY_4"),
    os.getenv("GCP_KEY_5"),
    os.getenv("GCP_KEY_6")
]

GOOGLE_API_KEY = next((key for key in GCP_KEYS if key), None)

if GOOGLE_API_KEY and GEMINI_AVAILABLE:
    genai.configure(api_key=GOOGLE_API_KEY)
    logger.info("✅ Google AI configured successfully")

# MCP-enabled functions for Agent1 (Image Generator)
def enhance_prompt_with_gemini(prompt: str, style: str) -> str:
    """
    Use Gemini to enhance the user's prompt for better image generation.
    
    Args:
        prompt (str): The original marketing prompt
        style (str): The desired image style
        
    Returns:
        str: Enhanced prompt optimized for image generation
    """
    if not GEMINI_AVAILABLE or not GOOGLE_API_KEY:
        # Basic enhancement without Gemini
        style_enhancers = {
            "realistic": "photorealistic, high detail, professional photography, sharp focus",
            "artistic": "artistic masterpiece, creative composition, painterly style",
            "cartoon": "cartoon style, vibrant colors, playful, animated character design",
            "photographic": "professional photograph, high quality, detailed, commercial photography",
            "illustration": "digital illustration, clean vector art, modern design"
        }
        enhancer = style_enhancers.get(style.lower(), "high quality, detailed")
        return f"{prompt}, {enhancer}"
    
    try:
        enhancement_prompt = f"""
        You are an expert prompt engineer for AI image generation. Take this marketing prompt and enhance it for optimal results.

        Original prompt: "{prompt}"
        Desired style: "{style}"

        Please provide an enhanced version that:
        1. Maintains the core marketing intent
        2. Adds specific technical details for better image quality
        3. Includes appropriate style descriptors for "{style}" style
        4. Adds professional marketing composition guidance
        5. Keeps the enhanced prompt under 150 words

        Return only the enhanced prompt without explanation.
        """
        
        model = genai.GenerativeModel('gemini-1.5-flash')
        response = model.generate_content(enhancement_prompt)
        enhanced = response.text.strip()
        
        logger.info(f"Gemini enhanced prompt: {enhanced}")
        return enhanced
        
    except Exception as e:
        logger.warning(f"Failed to enhance prompt with Gemini: {e}")
        style_enhancers = {
            "realistic": "photorealistic, high detail, professional photography",
            "artistic": "artistic masterpiece, creative composition",
            "cartoon": "cartoon style, vibrant colors, playful",
            "photographic": "professional photograph, high quality, detailed",
            "illustration": "digital illustration, clean design"
        }
        enhancer = style_enhancers.get(style.lower(), "high quality")
        return f"{prompt}, {enhancer}"

def generate_marketing_image(prompt: str, style: str = "realistic") -> str:
    """
    Generate a professional marketing image using Google Imagen3.
    
    Args:
        prompt (str): Description of the marketing image to generate
        style (str): Art style for the image (realistic, artistic, cartoon, photographic, illustration)
        
    Returns:
        str: JSON string containing image data and metadata
    """
    logger.info(f"🎨 Generating marketing image: {prompt}")
    
    try:
        # Enhance the prompt
        enhanced_prompt = enhance_prompt_with_gemini(prompt, style)
        
        # Try to generate with Google Genai SDK
        if GEMINI_AVAILABLE and GOOGLE_API_KEY:
            try:
                logger.info("🎨 Using Google Genai SDK for image generation")
                
                # Initialize the genai SDK client  
                client = genai_sdk.Client(api_key=GOOGLE_API_KEY)
                
                # Generate image using Imagen 3 via SDK
                result = client.models.generate_images(
                    model="imagen-3.0-generate-002",
                    prompt=enhanced_prompt,
                    config={
                        "number_of_images": 1,
                        "output_mime_type": "image/png"
                    }
                )
                
                # Check if we got a valid response with images
                if result and hasattr(result, 'generated_images') and len(result.generated_images) > 0:
                    generated_image = result.generated_images[0]
                    
                    if hasattr(generated_image, 'image') and hasattr(generated_image.image, 'image_bytes'):
                        # Convert image bytes to base64 data URL
                        image_bytes = generated_image.image.image_bytes
                        img_base64 = base64.b64encode(image_bytes).decode('utf-8')
                        
                        # Determine MIME type from the response or default to PNG
                        mime_type = getattr(generated_image.image, 'mime_type', 'image/png')
                        image_url = f"data:{mime_type};base64,{img_base64}"
                        
                        response_data = {
                            "success": True,
                            "image_url": image_url,
                            "prompt": prompt,
                            "enhanced_prompt": enhanced_prompt,
                            "style": style,
                            "generation_method": "google-genai-sdk",
                            "real_ai_generation": True
                        }
                        
                        logger.info("✅ Successfully generated real AI image with Google SDK!")
                        return json.dumps(response_data)
                        
            except Exception as e:
                logger.error(f"Google SDK generation failed: {e}")
        
        # Fallback: Generate a deterministic placeholder
        logger.info("🔄 Using placeholder URL fallback")
        prompt_hash = abs(hash(enhanced_prompt)) % 10000
        image_url = f"https://picsum.photos/seed/{prompt_hash}/1024/1024"
        
        response_data = {
            "success": True,
            "image_url": image_url,
            "prompt": prompt,
            "enhanced_prompt": enhanced_prompt,
            "style": style,
            "generation_method": "placeholder",
            "real_ai_generation": False
        }
        
        return json.dumps(response_data)
        
    except Exception as e:
        logger.error(f"Image generation failed: {e}")
        return json.dumps({
            "success": False,
            "error": f"Generation failed: {str(e)}",
            "prompt": prompt,
            "style": style
        })

def analyze_marketing_prompt(prompt: str, review_guidelines: str = "") -> str:
    """
    Analyze a marketing prompt for quality, relevance, and compliance.
    
    Args:
        prompt (str): The marketing prompt to analyze
        review_guidelines (str): Specific guidelines to check against
        
    Returns:
        str: JSON string containing detailed analysis and recommendations
    """
    logger.info(f"🔍 Analyzing marketing prompt: {prompt[:50]}...")
    
    try:
        word_count = len(prompt.split())
        
        # Check for marketing-specific elements
        marketing_keywords = [
            "professional", "corporate", "business", "marketing", "brand", "commercial",
            "office", "team", "collaboration", "presentation", "meeting", "workplace",
            "customer", "service", "product", "showcase", "display", "advertising"
        ]
        
        style_keywords = [
            "realistic", "photographic", "artistic", "creative", "modern", "clean",
            "minimalist", "professional", "high-quality", "detailed", "sharp"
        ]
        
        composition_keywords = [
            "lighting", "composition", "background", "foreground", "perspective",
            "angle", "framing", "focus", "depth", "contrast", "colors"
        ]
        
        # Count keyword categories
        marketing_score = sum(1 for word in marketing_keywords if word.lower() in prompt.lower()) / len(marketing_keywords)
        style_score = sum(1 for word in style_keywords if word.lower() in prompt.lower()) / len(style_keywords)
        composition_score = sum(1 for word in composition_keywords if word.lower() in prompt.lower()) / len(composition_keywords)
        
        # Base quality assessment
        if word_count < 5:
            base_quality = 0.3
            quality_issues = ["Prompt is too short and lacks detail"]
        elif word_count < 10:
            base_quality = 0.5
            quality_issues = ["Prompt could benefit from more descriptive details"]
        elif word_count < 20:
            base_quality = 0.7
            quality_issues = []
        elif word_count < 40:
            base_quality = 0.8
            quality_issues = []
        else:
            base_quality = 0.6
            quality_issues = ["Prompt might be too complex - consider simplifying"]
        
        # Adjust based on keyword presence
        quality_adjustment = (marketing_score * 0.2 + style_score * 0.15 + composition_score * 0.15)
        final_quality = min(1.0, base_quality + quality_adjustment)
        
        # Generate specific feedback
        missing_elements = []
        if marketing_score < 0.1:
            missing_elements.append("marketing context or business relevance")
        if style_score < 0.1:
            missing_elements.append("artistic style or visual quality descriptors")
        if "english" in review_guidelines.lower() and "english" not in prompt.lower():
            missing_elements.append("English language specification for text/signage")
        
        present_elements = []
        if marketing_score > 0.1:
            present_elements.append("marketing/business context")
        if style_score > 0.1:
            present_elements.append("style descriptors")
        if composition_score > 0.1:
            present_elements.append("composition guidance")
        
        # Calculate overall scores
        relevance_score = min(1.0, final_quality + (marketing_score * 0.2))
        safety_score = 0.95  # Generally high for marketing prompts
        
        # Check for potentially problematic content
        problematic_terms = ["violence", "inappropriate", "offensive", "controversial"]
        for term in problematic_terms:
            if term in prompt.lower():
                safety_score = 0.7
                break
        
        overall_score = (final_quality * 0.4 + relevance_score * 0.4 + safety_score * 0.2)
        
        # Generate recommendations
        recommendations = []
        
        if final_quality < 0.6:
            recommendations.append("Consider adding more descriptive details about the desired image")
        
        if marketing_score < 0.1:
            recommendations.append("Add marketing context (e.g., professional, business, corporate)")
        
        if "english" in review_guidelines.lower() and "english" not in prompt.lower():
            recommendations.append("Add 'English signage' or 'English text' to meet language requirements")
        
        if word_count < 10:
            recommendations.append("Expand prompt with lighting, composition, or environmental details")
        elif word_count > 50:
            recommendations.append("Consider simplifying prompt while keeping key elements")
        
        if not recommendations:
            if overall_score > 0.8:
                recommendations.append("Excellent prompt! Should generate high-quality marketing image")
            else:
                recommendations.append("Good prompt foundation - image should meet basic requirements")
        
        analysis_result = {
            "success": True,
            "quality_score": round(final_quality, 2),
            "relevance_score": round(relevance_score, 2),
            "safety_score": round(safety_score, 2),
            "overall_score": round(overall_score, 2),
            "word_count": word_count,
            "missing_elements": missing_elements,
            "present_elements": present_elements,
            "recommendations": recommendations[:5],
            "analysis_method": "prompt_analysis"
        }
        
        return json.dumps(analysis_result)
        
    except Exception as e:
        logger.error(f"Prompt analysis failed: {e}")
        return json.dumps({
            "success": False,
            "error": f"Analysis failed: {str(e)}",
            "prompt": prompt
        })

def generate_and_review_marketing_image(prompt: str, style: str = "realistic", review_guidelines: str = "") -> str:
    """
    Complete workflow: Generate a marketing image and provide quality review.
    
    Args:
        prompt (str): Description of the marketing image to generate
        style (str): Art style for the image (realistic, artistic, cartoon, photographic, illustration)
        review_guidelines (str): Specific guidelines for marketing review
        
    Returns:
        str: JSON string containing image, review, and recommendations
    """
    logger.info(f"🎭 Starting complete marketing workflow for: {prompt}")
    
    try:
        # Step 1: Generate the image
        generation_response = generate_marketing_image(prompt, style)
        generation_data = json.loads(generation_response)
        
        if not generation_data.get("success", False):
            return generation_response  # Return error
        
        # Step 2: Analyze the prompt (marketing review)
        analysis_response = analyze_marketing_prompt(prompt, review_guidelines)
        analysis_data = json.loads(analysis_response)
        
        # Combine results
        workflow_result = {
            "success": True,
            "image": {
                "url": generation_data.get("image_url", ""),
                "data": generation_data.get("image_url", ""),
                "prompt": prompt,
                "style": style
            },
            "review": {
                "quality_score": analysis_data.get("overall_score", 0.7),
                "final_status": "passed" if analysis_data.get("overall_score", 0) > 0.7 else "needs_improvement",
                "iterations": 1,
                "passed": analysis_data.get("overall_score", 0) > 0.7,
                "recommendations": analysis_data.get("recommendations", []),
                "analysis_details": analysis_data
            },
            "metadata": {
                "generation_method": generation_data.get("generation_method", "unknown"),
                "real_ai_generation": generation_data.get("real_ai_generation", False),
                "workflow_type": "gradio_mcp_server"
            }
        }
        
        logger.info("✅ Complete marketing workflow successful!")
        return json.dumps(workflow_result)
        
    except Exception as e:
        logger.error(f"Complete workflow failed: {e}")
        return json.dumps({
            "success": False,
            "error": f"Workflow failed: {str(e)}",
            "prompt": prompt,
            "style": style
        })

# Gradio interface functions
def process_generated_image_and_results(api_response_str: str) -> Tuple[Image.Image, str]:
    """Process API response and return image and review text for Gradio display"""
    try:
        response_data = json.loads(api_response_str)
        
        if not response_data.get('success', False):
            return None, f"❌ Generation failed: {response_data.get('error', 'Unknown error')}"
        
        # Extract image data
        image_info = response_data.get('image', {})
        image_data_b64 = image_info.get('data', image_info.get('url', ''))
        
        image = None
        if image_data_b64 and image_data_b64.startswith('data:image'):
            try:
                base64_data = image_data_b64.split(',')[1]
                image_bytes = base64.b64decode(base64_data)
                image = Image.open(io.BytesIO(image_bytes))
            except Exception as e:
                logger.error(f"Error processing image: {str(e)}")
        
        # Extract review data
        review_data = response_data.get('review', {})
        
        if review_data:
            quality_score = review_data.get('quality_score', 0)
            passed = review_data.get('passed', False)
            final_status = review_data.get('final_status', 'unknown')
            recommendations = review_data.get('recommendations', [])
            
            status_emoji = "🟢" if passed else "🔴"
            
            # Extract metadata about generation method
            metadata = response_data.get('metadata', {})
            generation_method = metadata.get('generation_method', 'unknown')
            
            generation_info = ""
            if generation_method == "google-genai-sdk":
                generation_info = "🎨 **Generated with**: Google Imagen3 SDK (Real AI)\n"
            elif generation_method == "placeholder":
                generation_info = "🎨 **Generated with**: Placeholder (Fallback)\n"
            
            review_text = f"""**🔍 Marketing Review Results**

{generation_info}
**Quality Score:** {quality_score:.2f}/1.0
**Status:** {status_emoji} {final_status.upper()}
**Architecture:** Gradio MCP Server

**💡 Recommendations:**
"""
            
            if recommendations:
                for i, rec in enumerate(recommendations[:5], 1):
                    review_text += f"{i}. {rec}\n"
            else:
                review_text += "• Image meets quality standards\n"
                
        else:
            review_text = "⚠️ Review data not available"
        
        return image, review_text
        
    except Exception as e:
        return None, f"❌ Error processing results: {str(e)}"

def gradio_generate_marketing_image(prompt: str, style: str, review_guidelines: str) -> Tuple[Image.Image, str]:
    """Gradio interface wrapper for complete marketing image generation"""
    if not prompt.strip():
        return None, "⚠️ Please enter a prompt to generate an image."
    
    try:
        # Use the complete workflow function
        result_json = generate_and_review_marketing_image(prompt, style, review_guidelines)
        return process_generated_image_and_results(result_json)
    except Exception as e:
        error_message = f"❌ Error: {str(e)}"
        logger.error(error_message)
        return None, error_message

# Define suggested prompts
SUGGESTED_PROMPTS = {
    "Modern office team collaboration": ("A modern office space with diverse professionals collaborating around a sleek conference table, natural lighting, professional attire, English signage visible", "realistic"),
    "Executive boardroom meeting": ("Professional executive boardroom with polished conference table, city skyline view, business documents, English presentations on screens", "realistic"),
    "Customer service excellence": ("Professional customer service representative with headset in modern call center, English signage, clean corporate environment", "realistic"),
    "Product showcase display": ("Clean product showcase on white background with professional lighting, English product labels, minimalist marketing aesthetic", "realistic"),
    "Creative workspace design": ("Creative workspace with colorful design elements, inspirational English quotes on walls, modern furniture, artistic marketing materials", "artistic"),
    "Brand presentation setup": ("Professional brand presentation setup with English branded materials, corporate colors, marketing displays, conference room setting", "realistic")
}

# Create Gradio interface
with gr.Blocks(title="Marketing Image Generator MCP", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎨 Marketing Image Generator
    ### Professional AI image generation with built-in MCP server support
    
    **Gradio MCP Server** → **Google Imagen3** → **Marketing Review** → **Results**
    
    *MCP Server available at: `/gradio_api/mcp/sse`*
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Configuration")
            
            # Main inputs
            prompt = gr.Textbox(
                label="Describe your marketing image",
                placeholder="e.g., A modern office space with natural lighting, featuring diverse professionals collaborating around a sleek conference table",
                lines=4,
                info="Be specific about the scene, style, mood, and any marketing elements you want to include"
            )
            
            style = gr.Dropdown(
                choices=["realistic", "artistic", "cartoon", "photographic", "illustration"],
                value="realistic",
                label="Art Style",
                info="Choose the artistic style for your generated image"
            )
            
            review_guidelines = gr.Textbox(
                label="🔍 Marketing Review Guidelines (Optional)",
                placeholder="e.g., All text must be in English only, focus on professional appearance, ensure brand colors are prominent",
                lines=3,
                info="Provide specific marketing guidelines for review"
            )
            
            # Generate button
            generate_btn = gr.Button("🚀 Generate Marketing Image", variant="primary", size="lg")
            
            # Status
            gr.Markdown("🔍 **Mode**: Gradio MCP Server")
            gr.Markdown(f"🔑 **API Status**: {'✅ Configured' if GOOGLE_API_KEY else '❌ No API Key'}")
        
        with gr.Column(scale=2):
            # Results display
            gr.Markdown("### 🖼️ Generated Image & Review")
            
            image_output = gr.Image(
                label="Generated Marketing Image",
                type="pil",
                height=400,
                show_download_button=True
            )
            
            review_output = gr.Markdown(
                value="Click **Generate Marketing Image** to create your marketing image with automated review",
                label="Marketing Review Results"
            )
    
    # Suggested prompts section
    gr.Markdown("---")
    gr.Markdown("### 💡 Suggested Marketing Prompts")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("**🏢 Professional/Corporate**")
            for prompt_name in ["Modern office team collaboration", "Executive boardroom meeting", "Customer service excellence"]:
                suggested_prompt, suggested_style = SUGGESTED_PROMPTS[prompt_name]
                btn = gr.Button(prompt_name, size="sm")
                btn.click(
                    fn=lambda p=suggested_prompt, s=suggested_style: (p, s),
                    outputs=[prompt, style]
                )
        
        with gr.Column():
            gr.Markdown("**🎨 Creative/Marketing**")
            for prompt_name in ["Product showcase display", "Creative workspace design", "Brand presentation setup"]:
                suggested_prompt, suggested_style = SUGGESTED_PROMPTS[prompt_name]
                btn = gr.Button(prompt_name, size="sm")
                btn.click(
                    fn=lambda p=suggested_prompt, s=suggested_style: (p, s),
                    outputs=[prompt, style]
                )
    
    # Event handlers
    generate_btn.click(
        fn=gradio_generate_marketing_image,
        inputs=[prompt, style, review_guidelines],
        outputs=[image_output, review_output],
        show_progress=True
    )
    
    # Footer
    gr.Markdown("""
    ---
    <div style='text-align: center; color: #666; font-size: 0.9rem;'>
        <p>🎨 Marketing Image Generator | Gradio MCP Server</p>
        <p>Image Generation + Marketing Review + MCP API</p>
        <p>MCP Endpoint: <code>/gradio_api/mcp/sse</code></p>
    </div>
    """)

if __name__ == "__main__":
    logger.info("🚀 Starting Marketing Image Generator with MCP Server")
    logger.info(f"🔑 Google AI: {'✅ Configured' if GOOGLE_API_KEY else '❌ No API Key'}")
    logger.info("🔗 MCP Server will be available at /gradio_api/mcp/sse")
    
    demo.launch(mcp_server=True)