import gradio as gr
import cv2
import mediapipe as mp
import numpy as np
import time
from pathlib import Path
import tempfile
import os
from face_reconstruction_main import FaceMeshDetector
import plotly.graph_objects as go
import json

# Initialize the detector
detector = FaceMeshDetector()

def process_frame(frame):
    """Process a single frame and return the image with face mesh overlay"""
    if frame is None:
        return None
    
    img, faces = detector.detect_faces(frame)
    return img, faces

def create_3d_plot(vertices):
    """Create a 3D plotly figure from vertices"""
    x, y, z = zip(*vertices)
    
    fig = go.Figure(data=[go.Scatter3d(
        x=x, y=y, z=z,
        mode='markers',
        marker=dict(
            size=2,
            color=z,
            colorscale='Viridis',
        )
    )])
    
    fig.update_layout(
        scene=dict(
            xaxis_title='X',
            yaxis_title='Y',
            zaxis_title='Z',
            aspectmode='cube'
        ),
        margin=dict(l=0, r=0, b=0, t=0),
        showlegend=False
    )
    
    return fig

def save_obj_and_display(frame):
    """Save frame as OBJ and return both file path and 3D visualization"""
    try:
        if frame is None:
            return None, None, "No frame provided"
        
        img, faces = process_frame(frame)
        
        if not faces:
            return None, None, "No face detected in the frame"
        
        # Scale vertices
        vertices = np.array(faces[0]) * 150
        
        # Create temporary directory if it doesn't exist
        temp_dir = Path('temp')
        temp_dir.mkdir(exist_ok=True)
        
        # Generate timestamp
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        filename = f"face_mesh_{timestamp}"
        
        # Save OBJ file
        obj_path = temp_dir / f"{filename}.obj"
        success = detector.save_obj_file(vertices, filename)
        
        if not success:
            return None, None, "Failed to save OBJ file"
        
        # Create 3D visualization
        fig = create_3d_plot(vertices)
        
        return str(obj_path), fig, "Model generated successfully"
    
    except Exception as e:
        return None, None, f"Error: {str(e)}"

def cleanup_old_files():
    """Remove files older than 1 hour"""
    temp_dir = Path('temp')
    if temp_dir.exists():
        current_time = time.time()
        for file in temp_dir.glob('*.obj'):
            if current_time - file.stat().st_mtime > 3600:  # 1 hour
                file.unlink()

def create_interface():
    """Create the Gradio interface"""
    with gr.Blocks() as interface:
        gr.Markdown("""
        # 3D Face Reconstruction
        Upload an image or use your webcam to generate a 3D face model.
        
        Instructions:
        1. Upload an image or use the webcam
        2. Click 'Generate 3D Model'
        3. View the 3D model in the interactive viewer
        4. Download the OBJ file for use in 3D software
        """)
        
        with gr.Row():
            with gr.Column():
                # Input methods - using webcam and image input
                with gr.Tab("Webcam"):
                    input_webcam = gr.Image(
                        label="Webcam Input", 
                        type="numpy",
                        sources="webcam"
                    )
                
                with gr.Tab("Upload"):
                    input_image = gr.Image(
                        label="Upload Image",
                        type="numpy"
                    )
                
                generate_button = gr.Button("Generate 3D Model")
            
            with gr.Column():
                # Output displays
                obj_file = gr.File(label="Download OBJ File")
                plot_3d = gr.Plot(label="3D Preview")
                status_text = gr.Textbox(label="Status")
        
        # Set up event handlers for both webcam and uploaded image
        generate_button.click(
            fn=save_obj_and_display,
            inputs=[input_webcam],
            outputs=[obj_file, plot_3d, status_text]
        )
        
        generate_button.click(
            fn=save_obj_and_display,
            inputs=[input_image],
            outputs=[obj_file, plot_3d, status_text]
        )
        
        # Example section
        gr.Markdown("""
        ## Tips for Best Results
        - Ensure good lighting
        - Face the camera directly
        - Keep a neutral expression
        - Avoid extreme head rotations
        - Stay within arm's length of the camera
        """)
        
        # Cleanup old files periodically
        cleanup_old_files()
    
    return interface

# Create and launch the interface
interface = create_interface()

# Launch locally for testing
if __name__ == "__main__":
    interface.launch()