# Standard library imports
import os
import gradio as gr
import spaces



# Local imports
from age_estimation.age_estimation import age_estimation as _age_estimation
from detection.face_detection import face_detection 
from detection.object_detection import object_detection as _object_detection
from utils.ui_utils import update_input_visibility

@spaces.GPU
def age_estimation(input_type, uploaded_image, image_url, base64_string):
    return _age_estimation(input_type, uploaded_image, image_url, base64_string)

@spaces.GPU
def object_detection(input_type, uploaded_image, image_url, base64_string):
    return _object_detection(input_type, uploaded_image, image_url, base64_string)

with gr.Blocks() as demo:
    # Add a title to the interface
    gr.Markdown("# Computer Vision Tools")
    # Create a tab for face detection
    with gr.Tab("Face Detection"):
        # Input Method Selection
        face_input_type = gr.Radio(
            ["Upload File", "Enter URL", "Enter Base64"],
            label="Input Method",
            value="Upload File",  # Default selection
        )

        # Face Detection Method Selection
        face_detection_method = gr.Radio(
            ["OpenCV", "dlib"],
            label="Face Detection Method",
            value="OpenCV",  # Default selection
        )

        # Input Components (initially only file upload is visible)
        with gr.Row():
            face_img_upload = gr.Image(type="pil", label="Upload Image", visible=True)
            face_url_input = gr.Textbox(
                label="Enter Image URL", placeholder="e.g., https://...", visible=False
            )
            face_base64_input = gr.Textbox(
                label="Enter Base64 String",
                placeholder="Enter base64 string here...",
                visible=False,
            )

        # Process Button
        face_process_btn = gr.Button("Process Image")

        # Output Components
        face_image_output = gr.Image(label="Detected Faces Image")
        face_bbox_output = gr.JSON(label="Raw Bounding Box Data")

        # Link radio button change to visibility update function
        face_input_type.change(
            fn=update_input_visibility,
            inputs=[face_input_type],
            outputs=[face_img_upload, face_url_input, face_base64_input],
            queue=False,
            api_name=False,
        )

        # Link process button to the face detection function
        # The face_detection function will now return a tuple
        face_process_btn.click(
            fn=face_detection,
            inputs=[
                face_input_type,
                face_img_upload,
                face_url_input,
                face_base64_input,
                face_detection_method,
            ],
            outputs=[face_image_output, face_bbox_output],
        )
    # Create a tab for age estimation
    with gr.Tab("Age Estimation"):
        # Input Method Selection
        age_input_type = gr.Radio(
            ["Upload File", "Enter URL", "Enter Base64"],
            label="Input Method",
            value="Upload File",  # Default selection
        )

        # Input Components (initially only file upload is visible)
        with gr.Row():
            age_img_upload = gr.Image(type="pil", label="Upload Image", visible=True)
            age_url_input = gr.Textbox(
                label="Enter Image URL", placeholder="e.g., https://...", visible=False
            )
            age_base64_input = gr.Textbox(
                label="Enter Base64 String",
                placeholder="Enter base64 string here...",
                visible=False,
            )

        # Process Button
        age_process_btn = gr.Button("Estimate Age")

        # Output Components
        age_text_output = gr.Textbox(label="Estimated Age Summary")
        age_raw_output = gr.JSON(label="Raw Age Estimation Data")

        # Link radio button change to visibility update function
        age_input_type.change(
            fn=update_input_visibility,
            inputs=[age_input_type],
            outputs=[age_img_upload, age_url_input, age_base64_input],
            queue=False,
            api_name=False,
        )

        # Link process button to the age estimation function
        # The age_estimation function will now return a tuple
        age_process_btn.click(
            fn=age_estimation,
            inputs=[age_input_type, age_img_upload, age_url_input, age_base64_input],
            outputs=[age_text_output, age_raw_output],
        )
    # Create a tab for object detection
    with gr.Tab("Object Detection"):
        # Input Method Selection
        obj_input_type = gr.Radio(
            ["Upload File", "Enter URL", "Enter Base64"],
            label="Input Method",
            value="Upload File",  # Default selection
        )

        # Input Components (initially only file upload is visible)
        with gr.Row():
            obj_img_upload = gr.Image(type="pil", label="Upload Image", visible=True)
            obj_url_input = gr.Textbox(
                label="Enter Image URL", placeholder="e.g., https://...", visible=False
            )
            obj_base64_input = gr.Textbox(
                label="Enter Base64 String",
                placeholder="Enter base64 string here...",
                visible=False,
            )

        # Process Button
        obj_process_btn = gr.Button("Detect Objects")

        # Output Components
        obj_image_output = gr.Image(
            label="Detected Objects Image"
        )  # Updated label for clarity
        obj_raw_output = gr.JSON(label="Raw Object Detection Data")  # Added JSON output

        # Link radio button change to visibility update function
        obj_input_type.change(
            fn=update_input_visibility,
            inputs=[obj_input_type],
            outputs=[obj_img_upload, obj_url_input, obj_base64_input],
            queue=False,
            api_name=False,
        )

        # Link process button to the object detection function
        # The object_detection function now returns a tuple (image, raw_data)
        obj_process_btn.click(
            fn=object_detection,
            inputs=[obj_input_type, obj_img_upload, obj_url_input, obj_base64_input],
            outputs=[obj_image_output, obj_raw_output],  # Updated outputs
        )

    # Launch the Gradio demo
    port = int(os.environ.get("GRADIO_SERVER_PORT", 7860))
    import sys

    if "--server_port" in sys.argv:
        port = int(sys.argv[sys.argv.index("--server_port") + 1])
    demo.launch(server_port=port, ssr_mode=True, share=True)