# Standard library imports
import os
import gradio as gr
import torch.nn.functional as F
import torch.nn as nn

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

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 Component
        face_image_output = gr.Image(label="Detected Faces")

        # 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,
        )

        # Link process button to the face detection function
        # The face_detection function will need to be updated to handle these inputs
        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,
        )
    # 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 Component
        age_text_output = gr.Textbox(label="Estimated Age")

        # 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,
        )

        # Link process button to the age estimation function
        # The age_estimation function will need to be updated to handle these inputs
        age_process_btn.click(
            fn=age_estimation,
            inputs=[age_input_type, age_img_upload, age_url_input, age_base64_input],
            outputs=age_text_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 Component
        obj_image_output = gr.Image(label="Detected Objects")

        # 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,
        )

        # Link process button to the object detection function
        # The object_detection function will need to be updated to handle these inputs
        obj_process_btn.click(
            fn=object_detection,
            inputs=[obj_input_type, obj_img_upload, obj_url_input, obj_base64_input],
            outputs=obj_image_output,
        )

    # 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)