import gradio as gr
import sys
from pathlib import Path
import os
from utils.data_processing import detect_nsfw
# Import YOLO-related modules
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))





# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# NSFW Content Detection")
    with gr.Row():
        detection_mode = gr.Radio(["Simple Check", "Detailed Analysis"], label="Detection Mode", value="Simple Check")
    with gr.Row():
        conf_threshold = gr.Slider(0, 1, value=0.3, label="Confidence Threshold", visible=False)
        iou_threshold = gr.Slider(0, 1, value=0.45, label="Overlap Threshold", visible=False)
        label_mode = gr.Dropdown(["Draw box", "Draw Label", "Draw Confidence", "Censor Predictions"], label="Label Display Mode", value="Draw box", visible=False)
    
    with gr.Row():
        input_image = gr.Image(type="numpy", label="Upload an image or enter a URL")
        output_text = gr.Textbox(label="Detection Result")
    with gr.Row():
        output_image = gr.Image(type="numpy", label="Processed Image (for detailed analysis)", visible=False)

    detect_button = gr.Button("Detect")
    
    def update_visibility(mode):
        return [gr.update(visible=(mode == "Detailed Analysis"))] * 4
    
    detection_mode.change(update_visibility, inputs=[detection_mode], outputs=[conf_threshold, iou_threshold, label_mode, output_image])
    
    detect_button.click(
        detect_nsfw,
        inputs=[input_image, detection_mode, conf_threshold, iou_threshold, label_mode],
        outputs=[output_text, output_image]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0")