import os
import io
import json
import base64
import time
import numpy as np
import logging
import gradio as gr
from PIL import Image
from gradio_client import Client

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

# ─────────  Backend connection with health monitoring ─────────
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("HF_TOKEN environment variable is required")

# Backend connection state
backend_status = {
    "client": None,
    "connected": False,
    "last_check": None,
    "error_message": ""
}

def check_backend_connection():
    """Check backend connection and update status"""
    global backend_status
    
    try:
        test_client = Client("SnapwearAI/Image_to_flatlay", hf_token=HF_TOKEN)
        backend_status["client"] = test_client
        backend_status["connected"] = True
        backend_status["error_message"] = ""
        backend_status["last_check"] = time.time()
        logger.info("✅ Backend connection established")
        return True, "🟢 Backend is ready for Image To Flatlay"
    except Exception as e:
        backend_status["client"] = None
        backend_status["connected"] = False
        backend_status["last_check"] = time.time()
        
        error_str = str(e).lower()
        if "timeout" in error_str or "read operation timed out" in error_str:
            backend_status["error_message"] = "Backend is starting up (5-6 minutes on first load)"
            return False, "🟡 Backend is starting up. Please wait 5-6 minutes and try again."
        else:
            backend_status["error_message"] = f"Connection error: {str(e)}"
            return False, f"🔴 Backend error: {str(e)}"

# Initial connection attempt
try:
    success, status_msg = check_backend_connection()
    if success:
        logger.info("Backend client established")
    else:
        logger.warning(f"Initial backend connection failed: {status_msg}")
except Exception as e:
    logger.error(f"Failed to connect to backend: {e}")
    backend_status["connected"] = False
    backend_status["error_message"] = str(e)

def update_backend_status():
    """Check and update backend status"""
    success, status_msg = check_backend_connection()
    
    if success:
        css_class = "status-ready"
    elif "starting up" in status_msg:
        css_class = "status-starting"
    else:
        css_class = "status-error"
    
    status_html = f'<div class="status-banner {css_class}">{status_msg}</div>'
    return status_html

# ─────────  Styling ─────────
css = """
body, .gradio-container {
    font-family: 'Inter', 'SF Pro Display', -apple-system, BlinkMacSystemFont, sans-serif;
}
#col-left, #col-mid, #col-right {
    margin: 0 auto;
    max-width: 430px;
}
#col-showcase {
    margin: 0 auto;
    max-width: 1100px;
}
#button {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: #ffffff;
    font-weight: 600;
    font-size: 18px;
    border: none;
    border-radius: 12px;
    padding: 12px 24px;
    transition: all 0.3s ease;
}
#button:hover {
    transform: translateY(-2px);
    box-shadow: 0 8px 25px rgba(102,126,234,0.3);
}
#button:disabled {
    background: #ccc !important;
    cursor: not-allowed;
    transform: none;
    box-shadow: none;
}
.hero-section {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 40px 20px;
    border-radius: 20px;
    margin: 20px 0;
    text-align: center;
}
.feature-box {
    background: #f8fafc;
    border: 1px solid #e2e8f0;
    padding: 20px;
    border-radius: 12px;
    margin: 10px 0;
    border-left: 4px solid #667eea;
}
.showcase-section {
    background: #ffffff;
    border: 1px solid #e2e8f0;
    padding: 30px;
    border-radius: 16px;
    box-shadow: 0 4px 20px rgba(0,0,0,0.1);
    margin: 20px 0;
}
.step-header {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 15px;
    border-radius: 12px;
    text-align: center;
    font-weight: 600;
    margin: 10px 0;
}
.social-links {
    text-align: center;
    margin: 20px 0;
}
.social-links a {
    margin: 0 10px;
    padding: 8px 16px;
    background: #667eea;
    color: white;
    text-decoration: none;
    border-radius: 8px;
    transition: all 0.3s ease;
}
.social-links a:hover {
    background: #764ba2;
    transform: translateY(-2px);
}
.error-message {
    color: #dc3545;
    font-weight: 500;
}
.success-message {
    color: #28a745;
    font-weight: 500;
}
.status-banner {
    padding: 15px;
    border-radius: 12px;
    margin: 10px 0;
    text-align: center;
    font-weight: 600;
}
.status-ready {
    background: #d4edda;
    border: 1px solid #c3e6cb;
    color: #155724;
}
.status-starting {
    background: #fff3cd;
    border: 1px solid #ffeaa7;
    color: #856404;
}
.status-error {
    background: #f8d7da;
    border: 1px solid #f5c6cb;
    color: #721c24;
}
.queue-info {
    background: #e8f4fd;
    border: 1px solid #bee5eb;
    padding: 12px;
    border-radius: 8px;
    margin: 10px 0;
    text-align: center;
    font-size: 14px;
    color: #0c5460;
}
"""

def image_to_base64(image: Image.Image) -> str:
    """
    Convert a PIL Image to a base64‐encoded PNG string.
    """
    if image is None:
        return ""
    if image.mode not in ("RGB", "RGBA"):
        image = image.convert("RGB")
    buffer = io.BytesIO()
    image.save(buffer, format="PNG", optimize=True)
    buffer.seek(0)
    return base64.b64encode(buffer.getvalue()).decode("utf-8")


def base64_to_image(b64_str: str) -> Image.Image:
    """
    Decode a base64 string (with or without data URL prefix) into a PIL Image.
    """
    if not b64_str:
        return None
    try:
        if b64_str.startswith("data:"):
            b64_str = b64_str.split(",", 1)[1]
        data = base64.b64decode(b64_str)
        return Image.open(io.BytesIO(data)).convert("RGBA")
    except Exception as e:
        logger.error(f"Failed to decode base64 image: {e}")
        return None

# ───────── Section 2: Flatley Image Generation ─────────
def generate_flatlay(image, prompt):
        """
        1. Convert ImageEditor data to JSON payload.
        2. Use `mask_b64` directly.
        3. Call backend `/predict` endpoint.
        4. Decode returned base64 and return as PIL Image.
        """
        # Check backend connection first
        if not backend_status["connected"] or not backend_status["client"]:
            success, status_msg = check_backend_connection()
            if not success:
                return None, 0, status_msg
        current_client = backend_status["client"]
        # Validate inputs
        if not image:
            return None
        if not prompt:
            return None
            
        # 1) Prepare JSON payload
        payload_str = image_to_base64(image)

        # 2) Invoke backend
        from gradio_client import Client
        HF_TOKEN = os.getenv("HF_TOKEN")
        client = Client("SnapwearAI/Image_to_flatlay", hf_token=HF_TOKEN)

        try:
            result_b64 = current_client.predict(
                payload_str,
                prompt,
                api_name="/predict"
            )
        except Exception as e:
            logger.error(f"Image generation call failed: {e}")
            return None

        # 3) Decode and return
        result_img = base64_to_image(result_b64) if result_b64 else None
        return result_img

# ───────── Gradio App (Single Canvas) ─────────
# ─────────  Main UI ─────────
with gr.Blocks(css=css, title="Snapwear Image to Flatlay") as demo:
    
    # ──────── Hero Section ────────
    gr.HTML("""
        <div class="hero-section">
            <h1 style="font-size:48px;margin:0;background:linear-gradient(45deg,#fff,#f0f8ff);-webkit-background-clip:text;-webkit-text-fill-color:transparent;">
                👕 Snapwear Image To Flatlay
            </h1>
            <h2 style="font-size:24px; margin:10px 0; color:#333; opacity:0.9; text-align:center;">
                Transform your model’s cloth into a flatlay.
            </h2>
            <div class="social-links">
                <a href="https://snapwear.io" target="_blank">🌐 Official Website</a>
                <a href="https://www.instagram.com/snapwearai/" target="_blank">📸 Instagram</a>
                <a href="https://huggingface.co/spaces/SnapwearAI/Snapwear-Texture-Transfer" target="_blank">🎨 Pattern Transfer</a>
                <a href="https://huggingface.co/spaces/SnapwearAI/Snapwear-Virtual-Try-On" target="_blank">👗 Snapwear Virtual TryOn</a>
                <a href="https://huggingface.co/spaces/SnapwearAI/Snapwear_background_ai_model" target="_blank">🌄 Snapwear Background AI</a>
            </div>
            <p style="font-size:13px; margin-top:15px; opacity:0.7;">
                <b>Disclaimer:</b> This demo is free for trials only. Any solicitation 
                for payment based on the free features we provide on this HuggingFace Space 
                is a fraudulent act.
            </p>
        </div>
    """)
    # ──────── Backend Status Section ────────
    with gr.Row():
        with gr.Column():
            # Initial status display
            if backend_status["connected"]:
                initial_status = '<div class="status-banner status-ready">🟢 Image To Flatlay is ready!</div>'
            else:
                initial_status = '<div class="status-banner status-starting">🟡 Model may be starting up. Click "Check Status" to verify.</div>'
            
            status_display = gr.HTML(value=initial_status)
            
            # Status check button
            check_status_btn = gr.Button("🔄 Check Status", size="sm")

    # ──────── Key Features ────────
    gr.HTML("""
            <div style="display:grid;grid-template-columns:repeat(auto-fit,minmax(250px,1fr));gap:20px;margin:30px 0;">
                <div class="feature-box">
                    <h3>👗 Model-to-Flatlay Conversion</h3>
                    <p>Automatically transform a model-worn garment into a professional flatlay presentation</p>
                </div>
                <div class="feature-box">
                    <h3>⚙️ Automated Garment Extraction</h3>
                    <p>Isolate clothing from the model image with pixel-perfect precision for crisp layouts</p>
                </div>
                <div class="feature-box">
                    <h3>🚀 Rapid Processing</h3>
                    <p>Generate flatlay images in under 60 seconds with one click</p>
                </div>
            </div>
        """)

    # ──────── Main Interface ────────
    with gr.Row():
        with gr.Column(scale=1, elem_id="col-left"):
            gr.HTML('<div class="step-header">Step 1: Upload Model Image 🖼️</div>')
            model_image = gr.Image(
                label="Model Image",
                type="pil",
                image_mode="RGBA",
                height=600
            )
            gr.HTML('<div style="font-size:14px; color:#666; margin-top:8px; text-align:center;">'
            '⚠️ <b>Important:</b> First upload the model image wearing clothes.<br/>')
            
            gr.Examples(
                label="Example Model Images",
                inputs=model_image,
                examples_per_page=12,
                examples=[f"examples/{i}.jpg" for i in range(1, 3)] if os.path.exists("examples") else [],
            )
        
        
        # ③ Generated Image
        with gr.Column(scale=1, elem_id="col-right"):
            gr.HTML('<div class="step-header">Step 2: Prompt and Generate ✨👗</div>')
            result_preview = gr.Image(label="Generated flatlay",show_share_button=True, height=600)
            with gr.Column():
                prompt_box  = gr.Textbox(label="Clothing Prompt", placeholder="Describe the garment to flatlay...")
                # ✅ Adding prompt examples here
                gr.Examples(
                    label="Prompt Examples",
                    examples=[
                        "t-shirt",
                        "dress"
                    ],
                    inputs=prompt_box
                )
                gen_button  = gr.Button("Generate Flatlay", elem_id="button")


    # ──────── Event Handlers ────────
        
    # Status check button
    check_status_btn.click(
            fn=update_backend_status,
            outputs=[status_display]
        )
    
    gen_button.click(
            fn=generate_flatlay,
            inputs=[model_image, prompt_box],
            outputs=[result_preview],
            concurrency_limit=1,  # Match backend queue system
            show_progress=True
        )
    # ──────── Look-Book Grid ────────
    # Virtual try-on examples
    lookbook_rows = [
        [f"lookbook/model{i}.jpg",
         f"lookbook/result{i}.jpg"]
        for i in range(1, 4) if os.path.exists("lookbook")  # adjust range to your file count
    ]

    if lookbook_rows:
        gr.HTML("""
            <div class="showcase-section">
                <h2 style="text-align:center;color:#333;margin-bottom:30px;">
                    🌟 Image to Flatlay Showcase
                </h2>
            </div>
        """)

        gr.Examples(
            examples=lookbook_rows,
            inputs=[model_image, result_preview],
            label=None,
            examples_per_page=4,
        )

    # ──────── Model Comparison Grid ────────
    if os.path.exists("examples/Grid.jpg"):
        gr.HTML("""
            <div class="showcase-section">
                <h2 style="text-align:center;color:#333;margin-bottom:20px;">
                    🔬 Model Comparison Analysis
                </h2>
                <p style="text-align:center;color:#666;margin-bottom:30px;font-size:16px;">
                    See how Snapwear Image to Flatlay compares against leading Models
                </p>
            </div>
        """)
        
        # Display the comparison grid image
        with gr.Row():
            with gr.Column():
                comparison_image = gr.Image(
                    value="examples/Grid.jpg",
                    label="Image to Flatlay Model Comparison",
                    show_label=True,
                    interactive=False,
                    height=600,
                    show_download_button=True,
                    show_share_button=False
                )
    
    gr.HTML("""
            <div style="background:#f8fafc;border:1px solid #e2e8f0;padding:30px;border-radius:16px;margin:30px 0;">
                <h2 style="text-align:center;color:#333;margin-bottom:25px;">🎯 Perfect For</h2>
                <div style="display:grid;grid-template-columns:repeat(auto-fit,minmax(200px,1fr));gap:20px;">
                    <div style="text-align:center;padding:15px;">
                        <h3 style="color:#667eea;">🛍️ Online Retailers</h3>
                        <p style="color:#666;">Generate consistent flatlay product shots for catalogs and listings</p>
                    </div>
                    <div style="text-align:center;padding:15px;">
                        <h3 style="color:#667eea;">👗 Fashion Designers</h3>
                        <p style="color:#666;">Showcase apparel collections in clean, styled layouts</p>
                    </div>
                    <div style="text-align:center;padding:15px;">
                        <h3 style="color:#667eea;">🛒 E-Commerce Sellers</h3>
                        <p style="color:#666;">Create professional flatlays to boost click-through and sales</p>
                    </div>
                    <div style="text-align:center;padding:15px;">
                        <h3 style="color:#667eea;">📱 Social Media Marketers</h3>
                        <p style="color:#666;">Design eye-catching flatlay visuals for Instagram and Pinterest</p>
                    </div>
                </div>
            </div>
        """)


    # ──────── Footer ────────
    gr.HTML("""
        <div style="text-align:center;padding:40px 20px;background:#f8fafc;border:1px solid #e2e8f0;border-radius:16px;margin:30px 0;">
            <h3 style="color:#333;">🚀 Powered by Snapwear AI</h3>
            <p style="color:#666;">
                Experience the future of virtual Photoshoot.
            </p>
            <div class="social-links">
                <a href="https://snapwear.io" target="_blank">🌐 Website</a>
                <a href="https://www.instagram.com/snapwearai/" target="_blank">📸 Instagram</a>
                <a href="https://huggingface.co/spaces/SnapwearAI/Snapwear-Texture-Transfer" target="_blank">🎨 Pattern Transfer</a>
                <a href="https://huggingface.co/spaces/SnapwearAI/Snapwear-Virtual-Try-On" target="_blank">👗 Snapwear Virtual TryOn</a>
            </div>
            <p style="font-size:12px;color:#999;margin-top:20px;">
                © 2024 Snapwear AI. Professional AI tools for fashion and design.
            </p>
        </div>
    """)

    
    
# ───────── Launch App ─────────
if __name__ == "__main__":
    demo.queue(
        max_size=20,
        default_concurrency_limit=1,  # Single concurrent request to match backend
        api_open=False
    ).launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_api=False
    )