import gradio as gr
from pathlib import Path
from gradio_openpose3d_editor import OpenPose3DEditor
from PIL import Image

_dwpose_detector = None

def get_dwpose_detector():
    """Get or create DWPose detector using easy-dwpose."""
    global _dwpose_detector
    if _dwpose_detector is None:
        from easy_dwpose import DWposeDetector
        _dwpose_detector = DWposeDetector(device="cpu")
    return _dwpose_detector

def handle_export(image_path: str):
    """
    Receives the PNG file path exported by the editor and returns
    the file ready for display / download.
    """
    if not image_path:
        return None, "No image exported yet."
    p = Path(image_path)
    if not p.exists():
        return None, f"File not found: image_path"
    
    return f"✅ Pose Generated!"

def detect_pose(image, detect_resolution):
    """Main pose detection function."""
    if image is None:
        return None
    
    try:
        # Convert to PIL if needed
        if isinstance(image, str):
            image = Image.open(image)
        
        # Convert to RGB if necessary
        if image.mode != "RGB":
            image = image.convert("RGB")
        
        # Resize to detect_resolution while maintaining aspect ratio
        original_size = image.size
        ratio = detect_resolution / max(original_size)
        new_size = (int(original_size[0] * ratio), int(original_size[1] * ratio))
        image_resized = image.resize(new_size, Image.Resampling.LANCZOS)
        
        # Process based on model type
        detector = get_dwpose_detector()
        result = detector(
            image_resized,
            output_type="pil",         
        )
        
        # Resize result back to original size if needed
        if result is not None and hasattr(result, 'size') and result.size != original_size:
            result = result.resize(original_size, Image.Resampling.LANCZOS)
        
        return result
            
    except Exception as e:
        print(f"Error during processing: {str(e)}")
        import traceback
        traceback.print_exc()
        return None
    
with gr.Blocks(title="OpenPose 3D Editor") as demo:
    gr.Markdown(
    """
    # 🧘 OpenPose 3D Editor Component Demo
    <span>💻 <a href='https://github.com/DEVAIEXP/gradio_component_openpose3d_editor'>Component GitHub Code</a></span>
    <p>A powerful, custom interactive 3D character posing component for Gradio apps using `gr.HTML` and Three.js. It allows users to manipulate a 3D model's bones in real-time, adjust camera angles, apply aspect ratio cropping, and export the resulting pose map directly to the Python backend..</p>    
    """
    )    
    gr.Markdown(
        "Pose the 3D model and click the **save icon** to export "
        "the cropped image directly to the Python backend."
    )

    # ── Main component ──
    with gr.Row():
        with gr.Column():
            editor = OpenPose3DEditor(        
                height=560,
                label="Pose Editor",
            )
            detect_resolution = gr.Slider(
                label="📏 Detection Resolution",
                minimum=256,
                maximum=2048,
                value=512,
                step=64,
                info="Higher = more accurate but slower"
            )
            status = gr.Textbox(label="Status", interactive=False, scale=2)

        with gr.Column():            
            preview = gr.Image(label="Export preview", type="filepath", interactive=False)
            output_image = gr.Image(label="🎨 Output Pose", type="pil", height=400, interactive=False)
            
    # Event triggered by the save button inside the component
    editor.export(
        fn=handle_export,
        inputs=editor,
        outputs=[status],
    )
    # Also shows preview when export occurs
    editor.export(
        fn=lambda p: p,
        inputs=editor,
        outputs=preview,
    ).then(
        fn=detect_pose,
        inputs=[preview, detect_resolution],
        outputs=output_image,
    )
   
if __name__ == "__main__":
    demo.launch()