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

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

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>"
        
        # Make API request
        response = requests.post(url=url, files=files)
        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: {str(e)}</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 class="matching-section">
            <div class="matches-table-container">
                <table class="matches-table">
                    <thead>
                        <tr>
                            <th>First Image Face</th>
                            <th>Second Image 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 """
    footer {visibility: hidden}
    
    /* Full width container */
    .gradio-container {
        max-width: 95% !important;
        width: 95% !important;
        margin: 0 auto !important;
        font-size: 16px !important;
        height: 600px !important;
    }
    
    /* Center everything */
    .container {
        display: flex;
        flex-direction: column;
        align-items: center;
        justify-content: center;
        width: 100%;
        height: 100% !important;
    }
    
    /* Main content layout */
    .main-content-row {
        display: flex;
        gap: 25px;
        width: 100%;
        margin-bottom: 25px;
        height: 100% !important;
        min-height: 500px;
    }
    
    .upload-section {
        flex: 2;
        display: flex;
        flex-direction: column;
        gap: 20px;
        height: 100%;
    }
    
    .result-section {
        flex: 1.2;
        display: flex;
        flex-direction: column;
        height: 100%;
        min-height: 400px;
    }
    
    .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;
        margin: 20px 0 !important;
        width: 100% !important;
    }
    
    .button-primary:hover {
        background: var(--button-primary-background-fill-hover) !important;
    }
    
    /* Result container styling with scroll - FIXED */
    .result-container {
        background: var(--background-fill-primary);
        padding: 8px;
        border-radius: 10px;
        margin-top: 0;
        width: 100%;
        text-align: center;
        height: 100% !important;
        display: flex;
        flex-direction: column;
        min-height: 400px;
    }
    
    .result-content {
        width: 100%;
        height: 100% !important;
        overflow-y: auto !important;
        flex: 1;
        display: flex;
        flex-direction: column;
    }
    
    /* Scrollbar styling for result content */
    .result-content::-webkit-scrollbar {
        width: 8px;
    }
    
    .result-content::-webkit-scrollbar-track {
        background: var(--background-fill-secondary);
        border-radius: 4px;
    }
    
    .result-content::-webkit-scrollbar-thumb {
        background: var(--border-color-primary);
        border-radius: 4px;
    }
    
    .result-content::-webkit-scrollbar-thumb:hover {
        background: var(--body-text-color);
        opacity: 0.7;
    }
    
    /* Matching table container with scroll - FIXED */
    .matches-table-container {
        width: 100%;
        overflow-x: auto;
        overflow-y: auto;
        max-height: 100% !important;
        flex: 1;
        display: block;
    }
    
    /* Matching table - FIXED LAYOUT */
    .matches-table {
        width: 100% !important;
        border-collapse: collapse;
        margin-top: 8px;
        font-size: 1em !important;
        min-width: 100% !important;
        table-layout: fixed;
    }
    
    .matches-table th {
        background: var(--background-fill-secondary);
        color: var(--body-text-color);
        padding: 12px 8px !important;
        text-align: center;
        font-size: 1em !important;
        font-weight: 700;
        border-bottom: 2px solid var(--border-color-primary);
        position: sticky;
        top: 0;
        z-index: 10;
    }
    
    .matches-table td {
        padding: 10px 8px !important;
        border-bottom: 1px solid var(--border-color-primary);
        text-align: center;
        font-size: 0.95em !important;
        color: var(--body-text-color);
        word-wrap: break-word;
    }
    
    /* Fixed column widths to prevent cutting */
    .matches-table th:nth-child(1),
    .matches-table td:nth-child(1) {
        width: 25% !important;
        min-width: 120px;
    }
    
    .matches-table th:nth-child(2),
    .matches-table td:nth-child(2) {
        width: 25% !important;
        min-width: 120px;
    }
    
    .matches-table th:nth-child(3),
    .matches-table td:nth-child(3) {
        width: 20% !important;
        min-width: 100px;
    }
    
    .matches-table th:nth-child(4),
    .matches-table td:nth-child(4) {
        width: 30% !important;
        min-width: 120px;
    }
    
    .row-number {
        font-weight: 600;
        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: 60px !important;
        height: 60px !important;
        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: 6px 10px !important;
        border-radius: 12px;
        font-size: 1em !important;
        font-weight: 700;
        text-transform: capitalize;
        display: inline-block;
        min-width: 120px;
    }
    
    .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;
    }
    
    /* Ensure the HTML output container has proper scroll */
    .gr-html {
        height: 100% !important;
        max-height: 100% !important;
        overflow: hidden !important;
        display: flex !important;
        flex-direction: column !important;
    }
    
    .gr-html > div {
        height: 100% !important;
        max-height: 100% !important;
        overflow: hidden !important;
        display: flex !important;
        flex-direction: column !important;
    }
    
    /* Force scrollbars to always show for consistent layout */
    .result-content {
        overflow-y: scroll !important;
    }
    
    .matches-table-container {
        overflow-y: scroll !important;
        overflow-x: auto !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"):
        # 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=320,
                            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=320,
                            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: 40px; 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_error=True,
        server_name="0.0.0.0",
        server_port=7860
    )