gradio-ui/components/visualization_display.py

"""
Visualization Display Component
Displays AI-generated architectural sketches with metadata and download functionality
"""

import base64
import logging
import tempfile
import os
from pathlib import Path
from typing import Optional, Tuple
from datetime import datetime

logger = logging.getLogger(__name__)


def format_visualization_display(visualization_data) -> Tuple[str, Optional[str], Optional[str]]:
    """
    Format visualization data for display with image and metadata
    
    Args:
        visualization_data: Visualization data (dict or VisualizationData object)
        
    Returns:
        Tuple of (HTML string, image path for gr.Image component, download path)
    """
    # Handle missing or empty visualization data
    if not visualization_data:
        return _format_no_visualization(), None, None
    
    # Handle both dict and object formats
    if isinstance(visualization_data, dict):
        success = visualization_data.get('success', True)
        if not success:
            error = visualization_data.get('error', {})
            return _format_visualization_error(error), None, None
        
        image_base64 = visualization_data.get('image_base64', '')
        prompt_used = visualization_data.get('prompt_used', '')
        model_version = visualization_data.get('model_version', 'Unknown')
        generation_timestamp = visualization_data.get('generation_timestamp', 'Unknown')
        image_format = visualization_data.get('image_format', 'PNG')
        resolution = visualization_data.get('resolution', 'Unknown')
        features_included = visualization_data.get('features_included', [])
    else:
        image_base64 = getattr(visualization_data, 'image_base64', '')
        prompt_used = getattr(visualization_data, 'prompt_used', '')
        model_version = getattr(visualization_data, 'model_version', 'Unknown')
        generation_timestamp = getattr(visualization_data, 'generation_timestamp', 'Unknown')
        image_format = getattr(visualization_data, 'image_format', 'PNG')
        resolution = getattr(visualization_data, 'resolution', 'Unknown')
        features_included = getattr(visualization_data, 'features_included', [])
    
    # Validate image data
    if not image_base64:
        return _format_no_visualization(), None, None
    
    # Convert base64 to temporary file for Gradio Image component
    image_path = None
    download_path = None
    try:
        # Decode base64 image
        image_bytes = base64.b64decode(image_base64)
        
        # Create temporary file for display
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png', mode='wb')
        temp_file.write(image_bytes)
        temp_file.close()
        image_path = temp_file.name
        
        # Create a downloadable file with a meaningful name
        download_dir = Path("gradio-ui/downloads")
        download_dir.mkdir(exist_ok=True)
        
        # Generate filename with timestamp including microseconds for uniqueness
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        download_filename = f"construction_visualization_{timestamp}.png"
        download_path = download_dir / download_filename
        
        # Write image to download location
        with open(download_path, 'wb') as f:
            f.write(image_bytes)
        
        logger.info(f"Created temporary image file: {image_path}")
        logger.info(f"Created downloadable image file: {download_path}")
    except Exception as e:
        logger.error(f"Error converting base64 to image file: {e}")
        return _format_no_visualization(), None, None
    
    # Build HTML output
    html = f"""
    <div style="padding: 20px;">
        <h2 style="margin-top: 0; color: #ffffff !important; font-weight: 700;">🎨 Architectural Visualization</h2>
        
        <div style="background: white; padding: 20px; border-radius: 8px; margin-bottom: 20px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); border: 1px solid #e2e8f0;">
            <h3 style="color: #1e40af !important; margin-top: 0; font-weight: 600;">📋 Generation Details</h3>
            <p style="color: #0f172a !important; margin: 8px 0;"><strong style="color: #0f172a !important;">Model:</strong> {model_version}</p>
            <p style="color: #0f172a !important; margin: 8px 0;"><strong style="color: #0f172a !important;">Generated:</strong> {generation_timestamp}</p>
            <p style="color: #0f172a !important; margin: 8px 0;"><strong style="color: #0f172a !important;">Format:</strong> {image_format}</p>
            <p style="color: #0f172a !important; margin: 8px 0;"><strong style="color: #0f172a !important;">Resolution:</strong> {resolution}</p>
            <p style="color: #0f172a !important; margin: 8px 0;"><strong style="color: #0f172a !important;">Download:</strong> Use the download button below the image to save</p>
        </div>
    """
    
    # Features included section
    if features_included:
        html += """
        <div style="background: white; padding: 20px; border-radius: 8px; margin-bottom: 20px; border-left: 5px solid #10b981; box-shadow: 0 2px 4px rgba(0,0,0,0.1); border: 1px solid #d1fae5;">
            <h3 style="color: #059669 !important; margin-top: 0; font-weight: 600;">🛡️ Disaster-Resistant Features</h3>
            <ul style="line-height: 1.8; color: #0f172a !important;">
        """
        for feature in features_included:
            html += f"<li style='color: #0f172a !important;'>{feature}</li>"
        html += """
            </ul>
        </div>
        """
    
    # Prompt used section
    if prompt_used:
        html += f"""
        <div style="background: white; padding: 20px; border-radius: 8px; margin-bottom: 20px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); border: 1px solid #e2e8f0;">
            <h3 style="color: #6366f1 !important; margin-top: 0; font-weight: 600;">💬 Generation Prompt</h3>
            <p style="color: #0f172a !important; line-height: 1.6; font-style: italic; background: #f8fafc; padding: 15px; border-radius: 4px; border-left: 3px solid #6366f1;">
                {prompt_used}
            </p>
        </div>
        """
    
    html += """
        <div style="background: white; padding: 15px; border-radius: 8px; margin-top: 20px; font-size: 0.9em; border: 1px solid #e2e8f0;">
            <p style="margin: 5px 0; color: #6b7280 !important; font-style: italic;">
                ℹ️ This visualization is AI-generated and represents a conceptual design. 
                Consult with licensed architects and engineers for actual construction plans.
            </p>
        </div>
    </div>
    """
    
    return html, image_path, str(download_path) if download_path else None


def _format_no_visualization() -> str:
    """Format message when no visualization is available"""
    return """
    <div style="padding: 20px;">
        <div style="background: white; padding: 30px; border-radius: 8px; text-align: center; border: 2px dashed #cbd5e1;">
            <h3 style="color: #6b7280 !important; margin-top: 0;">🎨 No Visualization Available</h3>
            <p style="color: #6b7280 !important;">
                The visualization agent did not generate an image for this construction plan.
                This may occur if the visualization service is unavailable or if there was an error during generation.
            </p>
        </div>
    </div>
    """


def _format_visualization_error(error) -> str:
    """Format error message for visualization failures"""
    if isinstance(error, dict):
        code = error.get('code', 'UNKNOWN')
        message = error.get('message', 'An unknown error occurred')
        retry_possible = error.get('retry_possible', False)
    else:
        code = getattr(error, 'code', 'UNKNOWN')
        message = getattr(error, 'message', 'An unknown error occurred')
        retry_possible = getattr(error, 'retry_possible', False)
    
    return f"""
    <div style="padding: 20px;">
        <div style="background: #fef2f2; padding: 20px; border-radius: 8px; border: 2px solid #dc2626;">
            <h3 style="color: #dc2626 !important; margin-top: 0;">❌ Visualization Generation Failed</h3>
            <p style="color: #0f172a !important;"><strong style="color: #0f172a !important;">Error Code:</strong> {code}</p>
            <p style="color: #0f172a !important;"><strong style="color: #0f172a !important;">Message:</strong> {message}</p>
            {f"<p style='color: #0f172a !important;'><strong style='color: #0f172a !important;'>Retry Possible:</strong> {'Yes' if retry_possible else 'No'}</p>" if retry_possible is not None else ''}
        </div>
    </div>
    """