import gradio as gr
import requests
import os
import json
import base64
from PIL import Image, ImageDraw, ImageFont
import io
import numpy as np
import cv2
from typing import List, Tuple, Optional
import logging
from datetime import datetime
import tempfile

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

class SalesforceConnector:
    """Handle Salesforce API connections and data retrieval"""
    
    def __init__(self):
        self.access_token = None
        self.instance_url = None
        self.session = requests.Session()
        
    def authenticate(self, client_id: str, client_secret: str, username: str, password: str, security_token: str = ""):
        """Authenticate with Salesforce using OAuth 2.0"""
        try:
            auth_url = "https://login.salesforce.com/services/oauth2/token"
            auth_data = {
                'grant_type': 'password',
                'client_id': client_id,
                'client_secret': client_secret,
                'username': username,
                'password': password + security_token
            }
            
            response = self.session.post(auth_url, data=auth_data)
            response.raise_for_status()
            
            auth_result = response.json()
            self.access_token = auth_result['access_token']
            self.instance_url = auth_result['instance_url']
            
            self.session.headers.update({
                'Authorization': f'Bearer {self.access_token}',
                'Content-Type': 'application/json'
            })
            
            logger.info("Successfully authenticated with Salesforce")
            return True
            
        except requests.exceptions.RequestException as e:
            logger.error(f"Salesforce authentication failed: {e}")
            return False
    
    def fetch_dress_catalog(self) -> List[dict]:
        """Fetch dress catalog from Salesforce"""
        try:
            # Query for dress products (adjust SOQL query based on your Salesforce schema)
            query = """
            SELECT Id, Name, Description, Product_Image_URL__c, Category__c, Price__c, Size__c, Color__c
            FROM Product2 
            WHERE Category__c = 'Dress' OR Category__c = 'Clothing'
            ORDER BY Name
            LIMIT 50
            """
            
            query_url = f"{self.instance_url}/services/data/v58.0/query"
            params = {'q': query}
            
            response = self.session.get(query_url, params=params)
            response.raise_for_status()
            
            result = response.json()
            dresses = result.get('records', [])
            
            logger.info(f"Fetched {len(dresses)} dresses from Salesforce")
            return dresses
            
        except requests.exceptions.RequestException as e:
            logger.error(f"Failed to fetch dress catalog: {e}")
            return self._get_sample_dresses()
    
    def _get_sample_dresses(self) -> List[dict]:
        """Return sample dress data for demo purposes"""
        return [
            {
                "Id": "dress_001",
                "Name": "Elegant Evening Dress",
                "Description": "Beautiful black evening dress perfect for special occasions",
                "Product_Image_URL__c": "https://images.pexels.com/photos/1926769/pexels-photo-1926769.jpeg",
                "Category__c": "Dress",
                "Price__c": 299.99,
                "Size__c": "M",
                "Color__c": "Black"
            },
            {
                "Id": "dress_002", 
                "Name": "Summer Floral Dress",
                "Description": "Light and airy floral dress for summer occasions",
                "Product_Image_URL__c": "https://images.pexels.com/photos/1926769/pexels-photo-1926769.jpeg",
                "Category__c": "Dress",
                "Price__c": 149.99,
                "Size__c": "S",
                "Color__c": "Floral"
            },
            {
                "Id": "dress_003",
                "Name": "Business Casual Dress",
                "Description": "Professional dress suitable for office wear",
                "Product_Image_URL__c": "https://images.pexels.com/photos/1926769/pexels-photo-1926769.jpeg",
                "Category__c": "Dress", 
                "Price__c": 199.99,
                "Size__c": "L",
                "Color__c": "Navy"
            }
        ]

class VirtualTryOnEngine:
    """Handle virtual try-on processing"""
    
    def __init__(self):
        self.model_loaded = False
        
    def process_tryon(self, person_image: Image.Image, dress_image: Image.Image, dress_info: dict) -> Tuple[Image.Image, str]:
        """Process virtual try-on between person and dress images"""
        try:
            # For demo purposes, we'll create a simple overlay effect
            # In production, you would use a proper virtual try-on model
            result_image = self._create_demo_tryon(person_image, dress_image, dress_info)
            
            confidence_score = "85%"
            return result_image, confidence_score
            
        except Exception as e:
            logger.error(f"Try-on processing failed: {e}")
            return person_image, "Error"
    
    def _create_demo_tryon(self, person_image: Image.Image, dress_image: Image.Image, dress_info: dict) -> Image.Image:
        """Create a demo try-on result by overlaying dress info"""
        # Resize images to standard size
        person_img = person_image.convert("RGB")
        person_img = person_img.resize((512, 768))
        
        # Create a copy for the result
        result_img = person_img.copy()
        draw = ImageDraw.Draw(result_img)
        
        # Add dress information overlay
        try:
            font = ImageFont.load_default()
        except:
            font = None
            
        # Add semi-transparent overlay
        overlay = Image.new('RGBA', result_img.size, (0, 0, 0, 128))
        result_img = Image.alpha_composite(result_img.convert('RGBA'), overlay).convert('RGB')
        
        # Add dress information text
        dress_name = dress_info.get('Name', 'Selected Dress')
        price = dress_info.get('Price__c', 0)
        color = dress_info.get('Color__c', 'N/A')
        size = dress_info.get('Size__c', 'N/A')
        
        text_lines = [
            f"Virtual Try-On Result",
            f"Dress: {dress_name}",
            f"Price: ${price}",
            f"Color: {color}",
            f"Size: {size}",
            f"Confidence: 85%"
        ]
        
        y_offset = 20
        for line in text_lines:
            draw.text((20, y_offset), line, fill="white", font=font)
            y_offset += 25
            
        return result_img

class VirtualTryOnApp:
    """Main application class"""
    
    def __init__(self):
        self.sf_connector = SalesforceConnector()
        self.tryon_engine = VirtualTryOnEngine()
        self.dress_catalog = []
        self.selected_dress = None
        
    def initialize_salesforce(self, client_id: str, client_secret: str, username: str, password: str, security_token: str = ""):
        """Initialize Salesforce connection"""
        success = self.sf_connector.authenticate(client_id, client_secret, username, password, security_token)
        if success:
            self.dress_catalog = self.sf_connector.fetch_dress_catalog()
            return "✅ Successfully connected to Salesforce and loaded dress catalog!", self._get_dress_options()
        else:
            # Load sample data for demo
            self.dress_catalog = self.sf_connector._get_sample_dresses()
            return "⚠️ Using sample data (Salesforce connection failed)", self._get_dress_options()
    
    def _get_dress_options(self) -> List[str]:
        """Get dress options for dropdown"""
        return [f"{dress['Name']} - ${dress['Price__c']} ({dress['Color__c']})" for dress in self.dress_catalog]
    
    def select_dress(self, dress_selection: str) -> Tuple[str, str]:
        """Handle dress selection"""
        if not dress_selection or not self.dress_catalog:
            return "Please select a dress", ""
            
        # Find selected dress
        dress_name = dress_selection.split(' - ')[0]
        self.selected_dress = next((dress for dress in self.dress_catalog if dress['Name'] == dress_name), None)
        
        if self.selected_dress:
            dress_info = f"""
            **Selected Dress Details:**
            - **Name:** {self.selected_dress['Name']}
            - **Description:** {self.selected_dress['Description']}
            - **Price:** ${self.selected_dress['Price__c']}
            - **Color:** {self.selected_dress['Color__c']}
            - **Size:** {self.selected_dress['Size__c']}
            """
            
            image_url = self.selected_dress.get('Product_Image_URL__c', '')
            return dress_info, image_url
        
        return "Dress not found", ""
    
    def process_virtual_tryon(self, person_image: Image.Image) -> Tuple[Image.Image, str]:
        """Process the virtual try-on"""
        if person_image is None:
            return None, "Please upload a person image"
            
        if self.selected_dress is None:
            return person_image, "Please select a dress first"
        
        try:
            # Load dress image from URL
            dress_url = self.selected_dress.get('Product_Image_URL__c', '')
            if dress_url:
                dress_response = requests.get(dress_url, timeout=10)
                dress_image = Image.open(io.BytesIO(dress_response.content))
            else:
                # Create placeholder dress image
                dress_image = Image.new('RGB', (300, 400), color='lightblue')
            
            # Process virtual try-on
            result_image, confidence = self.tryon_engine.process_tryon(person_image, dress_image, self.selected_dress)
            
            status_message = f"✅ Virtual try-on completed! Confidence: {confidence}"
            return result_image, status_message
            
        except Exception as e:
            logger.error(f"Virtual try-on failed: {e}")
            return person_image, f"❌ Try-on failed: {str(e)}"

# Initialize the app
app = VirtualTryOnApp()

def setup_salesforce_connection(client_id, client_secret, username, password, security_token):
    """Setup Salesforce connection and return status and dress options"""
    status, dress_options = app.initialize_salesforce(client_id, client_secret, username, password, security_token)
    return status, gr.Dropdown(choices=dress_options, value=None)

def handle_dress_selection(dress_selection):
    """Handle dress selection and return dress info and image"""
    dress_info, image_url = app.select_dress(dress_selection)
    
    # Try to load dress image
    dress_image = None
    if image_url:
        try:
            response = requests.get(image_url, timeout=10)
            dress_image = Image.open(io.BytesIO(response.content))
        except:
            dress_image = None
    
    return dress_info, dress_image

def process_tryon(person_image):
    """Process virtual try-on"""
    return app.process_virtual_tryon(person_image)

# Custom CSS for better styling
custom_css = """
.gradio-container {
    font-family: 'Inter', sans-serif;
    max-width: 1200px !important;
    margin: 0 auto;
}

.header-text {
    text-align: center;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    font-size: 2.5rem;
    font-weight: bold;
    margin-bottom: 1rem;
}

.section-header {
    background: linear-gradient(90deg, #4facfe 0%, #00f2fe 100%);
    padding: 1rem;
    border-radius: 8px;
    color: white;
    font-weight: bold;
    margin: 1rem 0;
}

.status-success {
    background-color: #d4edda;
    color: #155724;
    padding: 1rem;
    border-radius: 8px;
    border: 1px solid #c3e6cb;
}

.status-error {
    background-color: #f8d7da;
    color: #721c24;
    padding: 1rem;
    border-radius: 8px;
    border: 1px solid #f5c6cb;
}

.image-container {
    border: 2px solid #e1e5e9;
    border-radius: 12px;
    padding: 1rem;
    background: white;
    box-shadow: 0 2px 8px rgba(0,0,0,0.1);
}

.tryon-results {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
    padding: 2rem;
    border-radius: 12px;
    margin-top: 1rem;
}
"""

# Create the Gradio interface
with gr.Blocks(css=custom_css, title="Virtual Try-On for Clothing") as demo:
    
    # Header
    gr.HTML("""
    <div class="header-text">
        👗 Virtual Try-On for Clothing
    </div>
    <p style="text-align: center; font-size: 1.2rem; color: #666; margin-bottom: 2rem;">
        Experience the future of fashion with AI-powered virtual try-on technology
    </p>
    """)
    
    # Salesforce Configuration Section
    with gr.Accordion("🔧 Salesforce Configuration", open=True):
        gr.HTML('<div class="section-header">Connect to Salesforce to access dress catalog</div>')
        
        with gr.Row():
            with gr.Column(scale=2):
                sf_client_id = gr.Textbox(
                    label="Client ID",
                    placeholder="Enter your Salesforce Client ID",
                    type="password"
                )
                sf_client_secret = gr.Textbox(
                    label="Client Secret", 
                    placeholder="Enter your Salesforce Client Secret",
                    type="password"
                )
            with gr.Column(scale=2):
                sf_username = gr.Textbox(
                    label="Username",
                    placeholder="Enter your Salesforce username"
                )
                sf_password = gr.Textbox(
                    label="Password",
                    placeholder="Enter your Salesforce password", 
                    type="password"
                )
                sf_security_token = gr.Textbox(
                    label="Security Token (Optional)",
                    placeholder="Enter security token if required"
                )
        
        sf_connect_btn = gr.Button("🔗 Connect to Salesforce", variant="primary", size="lg")
        sf_status = gr.Textbox(label="Connection Status", interactive=False)
    
    # Main Application Interface
    with gr.Row():
        # Left Column - Input Section
        with gr.Column(scale=1):
            gr.HTML('<div class="section-header">📷 Upload Your Photo</div>')
            
            person_image = gr.Image(
                label="Upload your photo",
                type="pil",
                height=400,
                elem_classes=["image-container"]
            )
            
            gr.HTML('<div class="section-header">👗 Select Dress</div>')
            
            dress_dropdown = gr.Dropdown(
                label="Choose a dress from catalog",
                choices=[],
                interactive=True
            )
            
            dress_info = gr.Markdown(
                label="Dress Information",
                value="Select a dress to see details"
            )
            
            dress_image = gr.Image(
                label="Selected Dress",
                type="pil",
                height=300,
                elem_classes=["image-container"]
            )
        
        # Right Column - Results Section  
        with gr.Column(scale=1):
            gr.HTML('<div class="section-header">✨ Virtual Try-On Results</div>')
            
            tryon_btn = gr.Button(
                "🎭 Generate Virtual Try-On",
                variant="primary",
                size="lg",
                elem_classes=["tryon-button"]
            )
            
            result_image = gr.Image(
                label="Try-On Result",
                type="pil", 
                height=500,
                elem_classes=["image-container", "tryon-results"]
            )
            
            result_status = gr.Textbox(
                label="Processing Status",
                interactive=False
            )
            
            # Action buttons
            with gr.Row():
                download_btn = gr.DownloadButton(
                    "💾 Download Result",
                    variant="secondary"
                )
                share_btn = gr.Button(
                    "📤 Share Result", 
                    variant="secondary"
                )
    
    # Advanced Options
    with gr.Accordion("⚙️ Advanced Options", open=False):
        with gr.Row():
            fit_adjustment = gr.Slider(
                label="Fit Adjustment",
                minimum=-10,
                maximum=10,
                value=0,
                step=1
            )
            brightness_adjustment = gr.Slider(
                label="Brightness",
                minimum=-50,
                maximum=50, 
                value=0,
                step=5
            )
            pose_adjustment = gr.Dropdown(
                label="Pose Adjustment",
                choices=["Natural", "Straighten", "Slight Turn"],
                value="Natural"
            )
    
    # Footer Information
    gr.HTML("""
    <div style="margin-top: 3rem; padding: 2rem; background: #f8f9fa; border-radius: 12px; text-align: center;">
        <h3>🌟 Features</h3>
        <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(250px, 1fr)); gap: 1rem; margin-top: 1rem;">
            <div>
                <strong>🔗 Salesforce Integration</strong><br>
                Direct access to your dress catalog
            </div>
            <div>
                <strong>🤖 AI-Powered Try-On</strong><br>
                Advanced virtual try-on technology
            </div>
            <div>
                <strong>📱 Responsive Design</strong><br>
                Works on all devices
            </div>
            <div>
                <strong>💾 Export Options</strong><br>
                Save and share your results
            </div>
        </div>
        <p style="margin-top: 1.5rem; color: #666;">
            Upload your photo, select a dress from your Salesforce catalog, and see how it looks on you instantly!
        </p>
    </div>
    """)
    
    # Event handlers
    sf_connect_btn.click(
        fn=setup_salesforce_connection,
        inputs=[sf_client_id, sf_client_secret, sf_username, sf_password, sf_security_token],
        outputs=[sf_status, dress_dropdown]
    )
    
    dress_dropdown.change(
        fn=handle_dress_selection,
        inputs=[dress_dropdown],
        outputs=[dress_info, dress_image]
    )
    
    tryon_btn.click(
        fn=process_tryon,
        inputs=[person_image],
        outputs=[result_image, result_status]
    )
    
    # Initialize with sample data on load
    demo.load(
        fn=lambda: app.initialize_salesforce("", "", "", "", ""),
        outputs=[sf_status, dress_dropdown]
    )

# Additional utility functions
def create_sample_person_image():
    """Create a sample person image for testing"""
    img = Image.new('RGB', (512, 768), color='lightgray')
    draw = ImageDraw.Draw(img)
    
    # Draw a simple person silhouette
    draw.ellipse([206, 100, 306, 200], fill='darkgray')  # Head
    draw.rectangle([231, 200, 281, 400], fill='darkgray')  # Body
    draw.rectangle([206, 400, 231, 600], fill='darkgray')  # Left leg
    draw.rectangle([281, 400, 306, 600], fill='darkgray')  # Right leg
    draw.rectangle([181, 220, 206, 350], fill='darkgray')  # Left arm
    draw.rectangle([306, 220, 331, 350], fill='darkgray')  # Right arm
    
    return img

def validate_image_input(image):
    """Validate uploaded image"""
    if image is None:
        return False, "No image uploaded"
    
    # Check image size
    if image.size[0] < 100 or image.size[1] < 100:
        return False, "Image too small (minimum 100x100 pixels)"
    
    # Check image format
    if image.format not in ['JPEG', 'PNG', 'JPG']:
        return False, "Unsupported image format (use JPEG or PNG)"
        
    return True, "Image is valid"

def enhance_image_quality(image: Image.Image) -> Image.Image:
    """Enhance image quality for better try-on results"""
    try:
        # Convert to RGB if necessary
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # Resize to optimal dimensions
        max_size = 1024
        if max(image.size) > max_size:
            ratio = max_size / max(image.size)
            new_size = tuple(int(dim * ratio) for dim in image.size)
            image = image.resize(new_size, Image.Resampling.LANCZOS)
        
        # Enhance contrast and sharpness
        from PIL import ImageEnhance
        
        enhancer = ImageEnhance.Contrast(image)
        image = enhancer.enhance(1.1)
        
        enhancer = ImageEnhance.Sharpness(image)
        image = enhancer.enhance(1.1)
        
        return image
        
    except Exception as e:
        logger.error(f"Image enhancement failed: {e}")
        return image

# Error handling and logging
def log_user_interaction(action: str, details: dict = None):
    """Log user interactions for analytics"""
    timestamp = datetime.now().isoformat()
    log_entry = {
        'timestamp': timestamp,
        'action': action,
        'details': details or {}
    }
    logger.info(f"User interaction: {json.dumps(log_entry)}")

# Performance monitoring
def measure_processing_time(func):
    """Decorator to measure processing time"""
    def wrapper(*args, **kwargs):
        start_time = datetime.now()
        result = func(*args, **kwargs)
        end_time = datetime.now()
        processing_time = (end_time - start_time).total_seconds()
        logger.info(f"Processing time for {func.__name__}: {processing_time:.2f} seconds")
        return result
    return wrapper

# Security utilities
def sanitize_user_input(input_text: str) -> str:
    """Sanitize user input to prevent security issues"""
    if not input_text:
        return ""
    
    # Remove potentially dangerous characters
    dangerous_chars = ['<', '>', '"', "'", '&', ';']
    for char in dangerous_chars:
        input_text = input_text.replace(char, '')
    
    return input_text.strip()

def validate_salesforce_credentials(client_id: str, client_secret: str, username: str, password: str) -> bool:
    """Validate Salesforce credentials format"""
    if not all([client_id, client_secret, username, password]):
        return False
    
    # Basic format validation
    if len(client_id) < 10 or len(client_secret) < 10:
        return False
        
    if '@' not in username:
        return False
        
    return True

# Configuration management
class AppConfig:
    """Application configuration management"""
    
    def __init__(self):
        self.max_image_size = 10 * 1024 * 1024  # 10MB
        self.supported_formats = ['JPEG', 'PNG', 'JPG']
        self.max_processing_time = 30  # seconds
        self.cache_duration = 3600  # 1 hour
        
    def get_huggingface_config(self):
        """Get configuration for Hugging Face deployment"""
        return {
            'title': 'Virtual Try-On for Clothing',
            'description': 'AI-powered virtual try-on application with Salesforce integration',
            'tags': ['fashion', 'ai', 'virtual-try-on', 'salesforce'],
            'requirements': 'requirements.txt'
        }

# Initialize configuration
config = AppConfig()

# Main execution
if __name__ == "__main__":
    # Launch the Gradio app
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        debug=True,
        show_error=True,
        favicon_path=None,
        ssl_verify=False
    )