import numpy as np
import gradio as gr
import cv2
from PIL import Image
import time
import os
import insightface
from insightface.app import FaceAnalysis
from insightface.model_zoo import get_model as get_insightface_model

# Global models (load once)
app = None
swapper = None

def init_models():
    global app, swapper
    if app is None:
        app = FaceAnalysis(name='buffalo_l', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
        app.prepare(ctx_id=0, det_size=(640, 640))
    if swapper is None:
        # Download model if not present
        model_path = os.path.expanduser('~/.insightface/models/inswapper_128.onnx')
        if not os.path.exists(model_path):
            # Provide a fallback; user must download manually or ensure it exists
            # In production, you would download from a known URL. Here we assume it's already available.
            pass
        swapper = get_insightface_model(model_path)

def swap_face_image(source_img_np, target_img_np, do_enhance=False):
    """Swap face from source image onto target image. Returns output image as numpy array (RGB)."""
    init_models()
    
    # Convert numpy (RGB) to BGR for insightface (which expects BGR)
    source_bgr = cv2.cvtColor(source_img_np, cv2.COLOR_RGB2BGR)
    target_bgr = cv2.cvtColor(target_img_np, cv2.COLOR_RGB2BGR)
    
    # Detect faces
    source_faces = app.get(source_bgr)
    target_faces = app.get(target_bgr)
    
    if len(source_faces) == 0:
        raise ValueError("No face detected in source image")
    if len(target_faces) == 0:
        raise ValueError("No face detected in target image")
    
    # Use the first face in source (largest by default)
    source_face = source_faces[0]
    # For target, we can pick the largest face (index 0 after sorting)
    target_faces = sorted(target_faces, key=lambda x: (x.bbox[2]-x.bbox[0])*(x.bbox[3]-x.bbox[1]), reverse=True)
    target_face = target_faces[0]
    
    # Perform swap
    result_bgr = swapper.get(target_bgr, target_face, source_face, paste_back=True)
    
    # Convert back to RGB
    result_rgb = cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB)
    
    # Optional face enhancement (placeholder - can add GFPGAN later)
    if do_enhance:
        # For now, just return as is (could integrate with roop's face_enhancer or GFPGAN)
        pass
    
    return result_rgb

def process_swap(source_file, target_file, do_face_enhancer):
    """Gradio interface function: takes source and target images, returns output image path."""
    if source_file is None or target_file is None:
        yield "❌ Please upload both source image and target image", None, None, gr.update(visible=False)
        return
    
    yield "🟡 Processing: Analyzing faces and detecting features...", None, None, gr.update(visible=False)
    time.sleep(1)
    
    try:
        # Ensure inputs are numpy arrays (they already are for gr.Image(type='numpy'))
        result_img = swap_face_image(source_file, target_file, do_face_enhancer)
        
        # Save to a temporary PNG file
        output_path = "output_swapped.png"
        Image.fromarray(result_img).save(output_path)
        
        yield "✅ Processing complete! Result is ready for download.", output_path, output_path, gr.update(visible=True)
    except Exception as e:
        yield f"❌ Error: {str(e)}", None, None, gr.update(visible=False)

# Custom CSS (same as original, with small adjustments for image output)
custom_css = """
:root {
    --neon-primary: #00f3ff;
    --neon-secondary: #ff00ff;
    --neon-accent: #00ff87;
    --neon-warning: #ffcc00;
    --dark-bg: #0a0a1a;
    --dark-panel: #13132b;
    --darker-panel: #0c0c1f;
    --text-primary: #ffffff;
    --text-secondary: #a0a0c0;
}
body {
    background: var(--dark-bg) !important;
    color: var(--text-primary) !important;
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
}
.gr-block {
    background: var(--dark-panel) !important;
    border-radius: 12px !important;
    border: 1px solid rgba(0, 243, 255, 0.2) !important;
    box-shadow: 0 0 15px rgba(0, 243, 255, 0.1) !important;
}
.gr-box {
    border-color: rgba(0, 243, 255, 0.3) !important;
    color: var(--text-primary) !important;
    background: rgba(10, 10, 26, 0.7) !important;
}
h1, h2, h3, h4, label, .gr-label {
    color: var(--text-primary) !important;
    text-shadow: 0 0 5px rgba(0, 243, 255, 0.5);
}
.gr-button {
    background: linear-gradient(45deg, var(--neon-primary), var(--neon-secondary)) !important;
    color: black !important;
    border: none !important;
    border-radius: 8px !important;
    padding: 12px 28px !important;
    font-weight: 600 !important;
    text-transform: uppercase !important;
    letter-spacing: 1px !important;
    box-shadow: 0 0 10px var(--neon-primary), 0 0 20px rgba(0, 243, 255, 0.3) !important;
    transition: all 0.3s ease !important;
    width: 100% !important;
    margin: 10px 0 !important;
}
.gr-button:not(:disabled):hover {
    transform: translateY(-2px);
    box-shadow: 0 0 15px var(--neon-primary), 0 0 30px rgba(0, 243, 255, 0.5) !important;
}
.gr-button:disabled {
    background: #4b5563 !important;
    box-shadow: none !important;
}
.status-indicator {
    display: inline-block;
    width: 12px;
    height: 12px;
    border-radius: 50%;
    margin-right: 10px;
    background-color: var(--neon-accent);
    box-shadow: 0 0 0 0 rgba(0, 255, 135, 0.7);
    animation: pulse 2s infinite;
}
@keyframes pulse {
    0% {
        box-shadow: 0 0 0 0 rgba(0, 255, 135, 0.7);
    }
    70% {
        box-shadow: 0 0 0 10px rgba(0, 255, 135, 0);
    }
    100% {
        box-shadow: 0 0 0 0 rgba(0, 255, 135, 0);
    }
}
.status-text {
    color: var(--text-secondary);
    font-size: 0.9rem;
    display: flex;
    align-items: center;
}
.header {
    display: flex;
    justify-content: space-between;
    align-items: center;
    margin-bottom: 2rem;
    padding: 1.5rem;
    background: linear-gradient(90deg, rgba(0,243,255,0.1) 0%, rgba(255,0,255,0.1) 100%);
    border-radius: 12px;
    border: 1px solid rgba(0, 243, 255, 0.3);
    box-shadow: 0 0 20px rgba(0, 243, 255, 0.2);
}
.title-section h1 {
    margin-bottom: 0.25rem;
    font-weight: 800;
    background: linear-gradient(45deg, var(--neon-primary), var(--neon-secondary));
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    text-shadow: 0 0 10px rgba(0, 243, 255, 0.5);
}
.title-section p {
    color: var(--text-secondary);
    margin: 0;
}
.instructions {
    background: linear-gradient(90deg, rgba(0,243,255,0.05) 0%, rgba(255,0,255,0.05) 100%) !important;
    padding: 1.5rem !important;
    margin-bottom: 2rem !important;
    border: 1px solid rgba(0, 243, 255, 0.2) !important;
}
.instructions h3 {
    margin-top: 0;
    margin-bottom: 0.75rem;
    color: var(--neon-primary) !important;
}
.instructions ul {
    margin-bottom: 0;
    padding-left: 1.5rem;
}
.instructions li {
    margin-bottom: 0.5rem;
    color: var(--text-secondary);
}
.instructions li:last-child {
    margin-bottom: 0;
}
.footer {
    text-align: center;
    margin-top: 2rem;
    padding-top: 1.5rem;
    border-top: 1px solid rgba(0, 243, 255, 0.2);
    color: var(--text-secondary);
    font-size: 0.875rem;
}
.image-container {
    border: 2px solid rgba(0, 243, 255, 0.3);
    border-radius: 12px;
    padding: 8px;
    background: rgba(0, 0, 0, 0.2);
    margin-bottom: 1.5rem;
    box-shadow: 0 0 15px rgba(0, 243, 255, 0.1);
}
.image-container .gr-label {
    background: rgba(0, 243, 255, 0.1);
    padding: 8px 12px;
    border-radius: 8px;
    margin-bottom: 10px;
    display: inline-block;
}
.control-panel {
    background: linear-gradient(90deg, rgba(0,243,255,0.08) 0%, rgba(255,0,255,0.08) 100%) !important;
    padding: 1.5rem !important;
    border: 1px solid rgba(0, 243, 255, 0.3) !important;
    border-radius: 12px !important;
    margin-bottom: 1.5rem !important;
}
.status-panel {
    background: var(--darker-panel) !important;
    padding: 1.5rem !important;
    border: 1px solid rgba(0, 255, 135, 0.3) !important;
    border-radius: 12px !important;
    box-shadow: 0 0 15px rgba(0, 255, 135, 0.1) !important;
}
.status-header {
    display: flex;
    align-items: center;
    margin-bottom: 1rem;
    padding-bottom: 0.5rem;
    border-bottom: 1px solid rgba(0, 255, 135, 0.2);
}
.status-content {
    min-height: 100px;
}
.output-highlight {
    border: 2px solid var(--neon-accent) !important;
    box-shadow: 0 0 20px rgba(0, 255, 135, 0.3) !important;
}
.upload-text {
    color: var(--text-secondary);
    text-align: center;
    padding: 20px;
}
.progress-bar {
    height: 6px;
    background: rgba(0, 243, 255, 0.2);
    border-radius: 3px;
    margin: 10px 0;
    overflow: hidden;
}
.progress-fill {
    height: 100%;
    background: linear-gradient(90deg, var(--neon-primary), var(--neon-accent));
    border-radius: 3px;
    width: 0%;
    transition: width 0.3s ease;
}
.control-item {
    margin-bottom: 1rem;
    padding: 1rem;
    background: rgba(0, 0, 0, 0.2);
    border-radius: 8px;
    border: 1px solid rgba(0, 243, 255, 0.1);
}
.control-item:last-child {
    margin-bottom: 0;
}
.download-btn {
    background: linear-gradient(45deg, var(--neon-accent), #00cc70) !important;
    margin-top: 15px !important;
}
.image-output-container {
    border: 2px solid rgba(0, 243, 255, 0.3);
    border-radius: 12px;
    padding: 8px;
    background: rgba(0, 0, 0, 0.2);
    margin-bottom: 1.5rem;
    box-shadow: 0 0 15px rgba(0, 243, 255, 0.1);
}
"""

# Build Gradio UI
with gr.Blocks(css=custom_css, title="Neon Face Swap AI - Image to Image", theme=gr.themes.Default(primary_hue="cyan", secondary_hue="pink")) as demo:
    with gr.Row(elem_classes="header"):
        with gr.Column(scale=3):
            with gr.Row(elem_classes="title-section"):
                gr.Markdown("""
                # 🌌 AI FACE SWAPPER (Image)
                ### Next-Generation AI Face Swapping for Images
                """)
        with gr.Column(scale=1):
            with gr.Row():
                gr.Markdown("""
                <div class="status-text">
                    <div class="status-indicator"></div>
                    System Online
                </div>
                """)
    
    with gr.Row():
        with gr.Column(scale=1, min_width=400):
            gr.Markdown("### 📤 SOURCE IMAGE")
            with gr.Group(elem_classes="image-container"):
                source_image = gr.Image(label="Upload Source Face", type="numpy", height=250, elem_classes="gr-box")
            
            gr.Markdown("### 🎯 TARGET IMAGE")
            with gr.Group(elem_classes="image-container"):
                target_image = gr.Image(label="Upload Target Image (where to swap the face)", type="numpy", height=250, elem_classes="gr-box")
        
        with gr.Column(scale=1, min_width=350):
            with gr.Group(elem_classes="control-panel"):
                gr.Markdown("### ⚙️ PROCESSING CONTROLS")
                with gr.Group(elem_classes="control-item"):
                    face_enhancer = gr.Checkbox(label="Enable Face Enhancer", value=False, info="(Experimental) Improves face quality but may not be stable")
                with gr.Group(elem_classes="control-item"):
                    submit = gr.Button("🔄 START SWAP", variant="primary")
            
            with gr.Group(elem_classes="status-panel"):
                with gr.Column():
                    gr.Markdown("""
                    <div class="status-header">
                        <h4 style="margin: 0;">📊 SYSTEM STATUS</h4>
                    </div>
                    """)
                    with gr.Group(elem_classes="status-content"):
                        info_text = gr.Textbox(label="Current Status", value="🟢 Ready to process image", interactive=False, lines=3)
                        with gr.Row():
                            gr.Markdown("**GPU:** ✅ Active")
                            gr.Markdown("**Memory:** 🟡 Stable")
        
        with gr.Column(scale=1, min_width=400):
            gr.Markdown("### 📥 OUTPUT RESULT")
            with gr.Group(elem_classes="image-output-container output-highlight"):
                output_image = gr.Image(label="Swapped Result", type="filepath", interactive=False, height=450, elem_classes="gr-box")
                download_btn = gr.Button("💾 DOWNLOAD RESULT", visible=False, elem_classes="download-btn")
                download_file = gr.File(label="Download Result", visible=False, interactive=False)
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("""
            <div class="instructions">
                <h3>🚀 HOW TO USE</h3>
                <ul>
                    <li><strong>Source Face:</strong> Select an image containing the face you want to use</li>
                    <li><strong>Target Image:</strong> Select an image where you want to place the source face</li>
                    <li><strong>Face Enhancer:</strong> (Experimental) Enable for higher quality results (takes longer)</li>
                    <li>Click <strong>START SWAP</strong> to begin processing</li>
                    <li>Monitor progress in the <strong>SYSTEM STATUS</strong> panel</li>
                    <li>Download your result when processing is complete</li>
                </ul>
                <p><strong>Note:</strong> The first run may take extra time to download the face analysis model.</p>
            </div>
            """)
    
    gr.Markdown("""
    <div class="footer">
        <p>Powered by InsightFace • Face Swapper AI (Image Mode) • GPU Accelerated</p>
    </div>
    """)
    
    # Event handlers
    def show_download():
        return gr.update(visible=True)
    
    submit.click(
        fn=process_swap,
        inputs=[source_image, target_image, face_enhancer],
        outputs=[info_text, output_image, download_file, download_btn]
    )
    
    download_btn.click(
        fn=lambda: gr.update(visible=True),
        inputs=None,
        outputs=download_file
    )

if __name__ == "__main__":
    demo.launch(share=False, server_name="0.0.0.0")