import gradio as gr
from models import text_model_1, text_label_map_1, text_model_2, text_label_map_2
from models import image_model_1, image_label_map_1, image_model_2, image_label_map_2
from models import video_model

# Text classification function (combines two text models)
def classify_text(input_text):
    results = []
    for model, label_map in [(text_model_1, text_label_map_1), (text_model_2, text_label_map_2)]:
        res = model(input_text)[0]
        label = label_map.get(res['label'], res['label'])
        score = res['score'] * 100
        results.append(f"{label} ({score:.2f}%)")
    return "\n".join(results)

# Image classification function (combines two image models)
def classify_image(image):
    results = []
    for model, label_map in [(image_model_1, image_label_map_1), (image_model_2, image_label_map_2)]:
        res = model(image)[0]
        label = label_map.get(res['label'], res['label'])
        score = res['score'] * 100
        results.append(f"{label} ({score:.2f}%)")
    return "\n".join(results)

# Video classification placeholder
def classify_video(video_file):
    return video_model(video_file)

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## Multi-Modal AI Detection Tool")

    with gr.Tab("Text Detection"):
        text_input = gr.Textbox(label="Enter text")
        text_output = gr.Textbox(label="Predictions")
        text_btn = gr.Button("Classify Text")
        text_btn.click(classify_text, inputs=text_input, outputs=text_output)

    with gr.Tab("Image Detection"):
        image_input = gr.Image(type="pil")
        image_output = gr.Textbox(label="Predictions")
        image_btn = gr.Button("Classify Image")
        image_btn.click(classify_image, inputs=image_input, outputs=image_output)

    with gr.Tab("Video Detection"):
        video_input = gr.File(label="Upload video")
        video_output = gr.Textbox(label="Prediction")
        video_btn = gr.Button("Classify Video")
        video_btn.click(classify_video, inputs=video_input, outputs=video_output)

demo.launch()