"""
All-in-One Marketing Image Generator with Proper A2A Architecture
Includes A2A Orchestrator + Agent1 (MCP-Imagen3) + Agent2 (Reviewer) + Gradio Interface
"""

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

# 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()

# A2A imports
try:
    from a2a_orchestrator import A2AOrchestrator
    A2A_AVAILABLE = True
except ImportError:
    A2A_AVAILABLE = False

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

# Initialize A2A Orchestrator
if A2A_AVAILABLE:
    orchestrator = A2AOrchestrator()
    logger.info("✅ A2A Orchestrator initialized")
else:
    orchestrator = None
    logger.warning("⚠️ A2A not available - will use fallback mode")

def process_generated_image_and_results(api_response):
    """Process API response and return image and review text for Gradio display"""
    try:
        response_data = api_response
        
        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')
            iterations = review_data.get('iterations', 0)
            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')
            real_ai = metadata.get('real_ai_generation', False)
            
            generation_info = ""
            if generation_method == "imagen3-mcp":
                generation_info = "🎨 **Generated with**: Imagen3-MCP (Real AI)\n"
            elif generation_method == "google-genai-sdk":
                generation_info = "🎨 **Generated with**: Google Genai SDK (Real AI)\n"
            elif generation_method == "placeholder":
                generation_info = "🎨 **Generated with**: Placeholder (Fallback)\n"
            
            review_text = f"""**🔍 A2A Marketing Review Results**

{generation_info}
**Quality Score:** {quality_score:.2f}/1.0
**Status:** {status_emoji} {final_status.upper()}
**Iterations:** {iterations}
**Architecture:** A2A Agent Communication

**💡 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"
            
            # Add workflow history
            workflow_history = review_data.get('workflow_history', [])
            if workflow_history and len(workflow_history) > 1:
                review_text += "\n**📋 A2A Workflow History:**\n"
                for item in workflow_history:
                    iteration = item.get('iteration', 'N/A')
                    score = item.get('review_score', 'N/A')
                    review_text += f"• Iteration {iteration}: Score {score:.2f}\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 generate_marketing_image_a2a(prompt, style, max_retries, review_threshold, review_guidelines):
    """Main function for the Gradio interface using A2A Orchestrator"""
    if not prompt.strip():
        return None, "⚠️ Please enter a prompt to generate an image."
    
    if not A2A_AVAILABLE or not orchestrator:
        return None, "❌ A2A Orchestrator not available. Please check setup."
    
    try:
        logger.info(f"🎭 Starting A2A workflow for: {prompt}")
        
        # Use the A2A orchestrator to coordinate agents
        result = asyncio.run(
            orchestrator.generate_image_with_review(
                prompt=prompt,
                style=style,
                max_retries=max_retries,
                review_threshold=review_threshold,
                review_guidelines=review_guidelines
            )
        )
        
        if result["success"]:
            image, review_text = process_generated_image_and_results(result)
            success_message = f"✅ A2A Image generation successful!\n\n{review_text}"
            return image, success_message
        else:
            error_message = f"❌ A2A Generation failed: {result.get('error', 'Unknown error')}"
            return None, error_message
            
    except Exception as e:
        error_message = f"❌ A2A Error: {str(e)}"
        logger.error(error_message)
        return None, error_message

async def check_a2a_agents():
    """Check if A2A agents are available"""
    if not orchestrator:
        return {"agent1": False, "agent2": False, "orchestrator": False}
    
    try:
        health_status = await orchestrator.check_agents_health()
        health_status["orchestrator"] = True
        return health_status
    except Exception as e:
        logger.error(f"Error checking agent health: {e}")
        return {"agent1": False, "agent2": False, "orchestrator": False}

def get_a2a_system_status():
    """Get A2A system status for display"""
    if not A2A_AVAILABLE:
        return "**🔧 A2A System Status:**\n\n❌ A2A not available\n⚠️ Install google-a2a to enable A2A functionality"
    
    try:
        health_status = asyncio.run(check_a2a_agents())
        
        status_text = "**🔧 A2A System Status:**\n\n"
        status_text += f"{'✅' if health_status.get('orchestrator', False) else '❌'} A2A Orchestrator\n"
        status_text += f"{'✅' if health_status.get('agent1', False) else '❌'} Agent1 (Image Generator)\n"
        status_text += f"{'✅' if health_status.get('agent2', False) else '❌'} Agent2 (Marketing Reviewer)\n"
        
        all_healthy = all(health_status.values())
        if all_healthy:
            status_text += "\n🎉 All A2A agents are running!"
        else:
            status_text += "\n⚠️ Some A2A agents are not responding."
            status_text += "\nRun: `python a2a_servers.py` to start agents"
        
        return status_text
        
    except Exception as e:
        return f"**🔧 A2A System Status:**\n\n❌ Error checking status: {str(e)}"

# Define suggested prompts (same as before)
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="A2A Marketing Image Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎨 A2A Marketing Image Generator
    ### Professional AI image generation with proper A2A agent architecture
    
    **Agent1** (MCP-Imagen3) ↔ **A2A** ↔ **Agent2** (Marketing Reviewer) ↔ **A2A** ↔ **Orchestrator**
    """)
    
    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, check accessibility compliance, verify readability",
                lines=3,
                info="Provide specific marketing guidelines for the Marketing Reviewer to evaluate against your brand standards"
            )
            
            # Advanced settings
            with gr.Accordion("🔧 Advanced Settings", open=False):
                max_retries = gr.Slider(
                    minimum=1,
                    maximum=5,
                    value=3,
                    step=1,
                    label="Max Retries",
                    info="Maximum number of retry attempts if generation fails"
                )
                
                review_threshold = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.8,
                    step=0.1,
                    label="Quality Threshold",
                    info="Minimum quality score required for auto-approval"
                )
            
            # Generate button
            generate_btn = gr.Button("🚀 Generate with A2A Architecture", variant="primary", size="lg")
            
            # Status
            gr.Markdown("🔍 **Mode**: A2A Agent-to-Agent Communication")
            
            # System status
            with gr.Accordion("📊 A2A System Status", open=False):
                status_display = gr.Markdown(get_a2a_system_status())
                refresh_status_btn = gr.Button("🔄 Refresh A2A Status", size="sm")
        
        with gr.Column(scale=2):
            # Results display
            gr.Markdown("### 🖼️ Generated Image & A2A Review")
            
            image_output = gr.Image(
                label="Generated Marketing Image",
                type="pil",
                height=400,
                show_download_button=True
            )
            
            review_output = gr.Markdown(
                value="Click **Generate with A2A Architecture** to create your marketing image using agent-to-agent communication",
                label="A2A 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=generate_marketing_image_a2a,
        inputs=[prompt, style, max_retries, review_threshold, review_guidelines],
        outputs=[image_output, review_output],
        show_progress=True
    )
    
    refresh_status_btn.click(
        fn=get_a2a_system_status,
        outputs=status_display
    )
    
    # Footer
    gr.Markdown("""
    ---
    <div style='text-align: center; color: #666; font-size: 0.9rem;'>
        <p>🎨 A2A Marketing Image Generator | Proper Agent-to-Agent Architecture</p>
        <p>Agent1 (MCP-Imagen3) + Agent2 (Reviewer) + A2A Orchestrator + Gradio Interface</p>
        <p>Start agents with: <code>python a2a_servers.py</code></p>
    </div>
    """)

if __name__ == "__main__":
    logger.info("🚀 Starting A2A Marketing Image Generator")
    logger.info(f"🤖 A2A Available: {'✅ Yes' if A2A_AVAILABLE else '❌ No'}")
    
    if A2A_AVAILABLE:
        logger.info("💡 Using proper A2A agent-to-agent communication")
        logger.info("🔗 Architecture: Gradio → Orchestrator → Agent1 (MCP-Imagen3) → Agent2 (Reviewer)")
        logger.info("📋 Make sure to start agents: python a2a_servers.py")
    else:
        logger.warning("💡 A2A not available - install google-a2a for full functionality")
    
    demo.launch()