import gradio as gr
import os
import requests
from PIL import Image
import json

# Configuration
API_BASE_URL = os.getenv("API_BASE_URL")
API_TOKEN = os.getenv("API_TOKEN")

def face_compare(frame1, frame2, request: gr.Request = None):
    """Face comparison with enhanced result display"""
    try:
        url = f"{API_BASE_URL}"
        
        # Prepare files
        files = {}
        if frame1:
            files['file1'] = open(frame1, 'rb')
        if frame2:
            files['file2'] = open(frame2, 'rb')
        
        if not files:
            return "<div class='error-message'>Please upload both images</div>"
        
        # Add Bearer token to headers
        headers = {
            "Authorization": f"Bearer {API_TOKEN}"
        }

        # Make API request
        response = requests.post(url=url, files=files, headers=headers)
        result = response.json()
        
        # Close files
        for file in files.values():
            file.close()
        
        # Enhanced result processing
        return format_face_comparison_result(result, frame1, frame2)
        
    except Exception as e:
        return f"<div class='error-message'>Error processing request</div>"

def format_face_comparison_result(result, img1_path, img2_path):
    """Format face comparison results with professional styling"""
    
    detections = result.get("detections", [])
    matches = result.get("match", [])
    
    # Create result HTML
    html = "<div class='result-content'>"
    
    # Detection results - show all detected faces
    if detections:
        for i, detection in enumerate(detections):
            face_image = detection.get("face", "")
            first_face_index = detection.get("firstFaceIndex")
            second_face_index = detection.get("secondFaceIndex")
                            
    # Matching results in the new table format
    if matches:
        html += """
        <div>
            <div class="matches-table">
                <table>
                    <thead>
                        <tr>
                            <th>First Face</th>
                            <th>Second Face</th>
                            <th>Similarity Score</th>
                            <th>Result</th>
                        </tr>
                    </thead>
                    <tbody>
        """
        
        # Group matches by first image face index for better organization
        match_groups = {}
        for match in matches:
            first_face_index = match.get("firstFaceIndex", "N/A")
            if first_face_index not in match_groups:
                match_groups[first_face_index] = []
            match_groups[first_face_index].append(match)
        
        row_number = 1
        for first_face_index in sorted(match_groups.keys()):
            for match in match_groups[first_face_index]:
                first_face_index = match.get("firstFaceIndex", "N/A")
                second_face_index = match.get("secondFaceIndex", "N/A")
                similarity = match.get("similarity", 0)
                
                # Get face images for display
                first_face_img = ""
                second_face_img = ""
                
                for detection in detections:
                    if detection.get("firstFaceIndex") == first_face_index:
                        first_face_img = detection.get("face", "")
                    if detection.get("secondFaceIndex") == second_face_index:
                        second_face_img = detection.get("face", "")
                
                # Determine result and color
                if similarity >= 0.6:  # Threshold for same person
                    result_text = "same person"
                    result_class = "result-same"
                else:
                    result_text = "different person"
                    result_class = "result-different"
                
                first_face_display = f"<img src='data:image/png;base64,{first_face_img}' class='table-face-thumbnail' />" if first_face_img else f"Face {first_face_index}"
                second_face_display = f"<img src='data:image/png;base64,{second_face_img}' class='table-face-thumbnail' />" if second_face_img else f"Face {second_face_index}"
                
                html += f"""
                <tr>
                    <td class="face-cell">
                        <div class="face-display">
                            {first_face_display}
                            <div class="face-label">Face {first_face_index}</div>
                        </div>
                    </td>
                    <td class="face-cell">
                        <div class="face-display">
                            {second_face_display}
                            <div class="face-label">Face {second_face_index}</div>
                        </div>
                    </td>
                    <td class="similarity-score">{similarity:.4f}</td>
                    <td><span class="result-text {result_class}">{result_text}</span></td>
                </tr>
                """
                row_number += 1
        
        html += """
                    </tbody>
                </table>
            </div>
        </div>
        """
    else:
        html += "<div class='no-results'>No face matches found.</div>"
    
    html += "</div>"
    return html


def get_custom_css():
    """Return simplified CSS styling that works for both light and dark themes"""
    return """

    /* Center everything */
    .container {
        display: flex;
        flex-direction: column;
        align-items: center;
        justify-content: center;
        width: 100%;
    }
    
    /* Header styling - logo and text in same line */
    .company-header {
        background: var(--background-fill-primary);
        padding: 10px;
        text-align: center;
        width: 100%;
        display: flex;
        align-items: center;
        justify-content: center;
        gap: 25px;
        flex-wrap: wrap;
    }
    
    .header-logo {
        flex-shrink: 0;
    }
    
    .header-logo img {
        width: 80px;
        height: auto;
    }
    
    .header-text {
        text-align: center;
    }
    
    .header-text h1 {
        font-size: 2.4em !important;
        font-weight: 700;
        color: var(--body-text-color);
    }
    
    .header-text p {
        font-size: 1.3em !important;
        color: var(--body-text-color);
        opacity: 0.8;
    }
    
    /* Main content layout */
    .main-content-row {
        display: flex;
        gap: 25px;
        width: 100%;
    }
    
    .upload-section {
        flex: 2;
        display: flex;
        flex-direction: column;
        gap: 20px;
    }
    
    .result-section {
        flex: 1.2;
    }
    
    .upload-images-row {
        display: flex;
        gap: 20px;
        width: 100%;
    }
    
    .upload-image-col {
        flex: 1;
    }
    
    /* Button styling */
    .button-primary {
        background: var(--button-primary-background-fill) !important;
        border: none !important;
        padding: 6px 12px !important;
        font-size: 1.2em !important;
        font-weight: 600 !important;
        color: var(--button-primary-text-color) !important;
        border-radius: 8px !important;
        cursor: pointer !important;
        transition: background-color 0.2s ease !important;
        width: 100% !important;
    }
    
    .button-primary:hover {
        background: var(--button-primary-background-fill-hover) !important;
    }
        
    .result-content {
        width: 100%;
    }
    
    /* Detection cards */
    .detections-grid {
        display: grid;
        grid-template-columns: repeat(auto-fit, minmax(160px, 1fr));
        gap: 15px;
        justify-content: center;
    }
    
    .detection-card {
        background: var(--background-fill-secondary);
        padding: 4px;
        border-radius: 8px;
        text-align: center;
        display: flex;
        flex-direction: column;
        align-items: center;
    }
    
    .face-thumbnail {
        width: 60px;
        height: 60px;
        border-radius: 50%;
        object-fit: cover;
    }
        
    /* Matching table - NEW STYLING */
    .matches-table {
        display: flex;
        justify-content: center;
        width: 100%;
        overflow-x: auto;
    }
    
    .matches-table table {
        width: 100%;
        border-collapse: collapse;
        font-size: 1em !important;
        min-width: 450px;
    }
    
    .matches-table th {
        background: var(--background-fill-secondary);
        color: var(--body-text-color);
        padding: 4px 2px !important;
        text-align: center;
        font-size: 1em !important;
        font-weight: 700;
        border-bottom: 2px solid var(--border-color-primary);
    }
    
    .matches-table td {
        padding: 4px 2px !important;
        border-bottom: 1px solid var(--border-color-primary);
        text-align: center;
        font-size: 0.95em !important;
        color: var(--body-text-color);
    }
    
    .face-cell {
        vertical-align: middle;
    }
    
    .face-display {
        display: flex;
        flex-direction: column;
        align-items: center;
        gap: 5px;
    }
    
    .table-face-thumbnail {
        width: 70px;
        height: 70px;
        border-radius: 50%;
        object-fit: cover;
        border: 2px solid var(--border-color-primary);
    }
    
    .face-label {
        font-size: 0.9em !important;
        color: var(--body-text-color);
        opacity: 1;
        font-weight: 600;
    }
    
    .similarity-score {
        font-weight: 700;
        color: var(--body-text-color);
        font-size: 1.05em !important;
    }
    
    .result-text {
        padding: 8px 12px !important;
        border-radius: 12px;
        font-size: 1.1em !important;
        font-weight: 700;
        text-transform: capitalize;
    }
    
    .result-same {
        background: #d4edda;
        color: #155724;
    }
    
    .result-different {
        background: #f8d7da;
        color: #721c24;
    }
    
    .no-results {
        text-align: center;
        padding: 40px;
        color: var(--body-text-color);
        opacity: 0.7;
        font-style: italic;
        font-size: 1.1em !important;
    }
    
    /* Error messages */
    .error-message {
        background: var(--background-fill-secondary);
        color: var(--body-text-color);
        padding: 20px;
        border-radius: 8px;
        text-align: center;
        width: 100%;
        opacity: 0.9;
        font-size: 1.1em !important;
    }
    
    """

# Create Gradio interface
with gr.Blocks(
    title="MiniAiLive - Face Recognition WebAPI Playground", 
    css=get_custom_css()
) as demo:
    
    with gr.Column(elem_classes="container"):
        # Header Section - Logo and text in same line, centered
        gr.HTML("""
        <div class="company-header">
            <div class="header-logo">
                <img src="https://miniai.live/wp-content/uploads/2025/11/logo_new.png" alt="MiniAiLive Logo">
            </div>
            <div class="header-text">
                <h1>MiniAiLive Face Recognition WebAPI Playground</h1>
                <p>Experience our NIST FRVT Top Ranked 1:1 & 1:N Face Matching Technology</p>
            </div>
        </div>
        """)
        
        # Main Content - Upload and Results
        with gr.Row(elem_classes="main-content-row"):
            # Upload Section
            with gr.Column(scale=0.6, elem_classes="upload-section"):
                with gr.Row(elem_classes="upload-images-row"):
                    # First Image
                    with gr.Column(scale=1, elem_classes="upload-image-col"):
                        im_match_in1 = gr.Image(
                            type='filepath', 
                            height=380,
                            label="First Image",
                            show_download_button=False
                        )
                        gr.Examples(
                            examples=[
                                "assets/1.jpg",
                                "assets/2.jpg",
                                "assets/3.jpg",
                                "assets/4.jpg",
                            ],
                            inputs=im_match_in1,
                            label="First Image Examples"
                        )
                    
                    # Second Image
                    with gr.Column(scale=1, elem_classes="upload-image-col"):
                        im_match_in2 = gr.Image(
                            type='filepath', 
                            height=380,
                            label="Second Image",
                            show_download_button=False
                        )
                        gr.Examples(
                            examples=[
                                "assets/1-1.jpg",
                                "assets/2-1.jpg",
                                "assets/3-1.jpg",
                                "assets/4-1.jpg",
                            ],
                            inputs=im_match_in2,
                            label="Second Image Examples"
                        )
                
                btn_f_match = gr.Button(
                    "Compare Faces 🚀", 
                    variant='primary',
                    elem_classes="button-primary"
                )
            
            # Results Section
            with gr.Column(scale=0.4, elem_classes="result-section"):
                txt_compare_out = gr.HTML(
                    value="<div style='text-align: center; padding: 10px; font-size: 1.1em;'>Results will appear here after comparison</div>"
                )
    
    # Connect the function
    btn_f_match.click(
        face_compare, 
        inputs=[im_match_in1, im_match_in2], 
        outputs=txt_compare_out
    )

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