import gradio as gr
from train_abuse_model import (
    run_training,
    evaluate_saved_model,
    push_model_to_hub
)
from predict_pipeline import run_prediction_pipeline


with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Abuse Detection App")
    gr.Markdown("⚠️ Keep this tab open while training or evaluating.")

    with gr.Tab("🧪 Train / Evaluate"):
        with gr.Row():
            start_btn = gr.Button("🚀 Start Training")
            eval_btn = gr.Button("🔍 Evaluate Trained Model")
            push_btn = gr.Button("📤 Push Model to Hub")
        output_box = gr.Textbox(label="Logs", lines=25, interactive=False)
        start_btn.click(fn=run_training, outputs=output_box)
        eval_btn.click(fn=evaluate_saved_model, outputs=output_box)
        push_btn.click(fn=push_model_to_hub, outputs=output_box)

    with gr.Tab("🔮 Abuse Detection"):
        desc_input = gr.Textbox(label="📝 Relationship Description", lines=5, placeholder="Write a relationship story here...")
        chat_upload = gr.File(label="📁 Optional: WhatsApp Chat ZIP (.zip)", file_types=[".zip"])
        predict_btn = gr.Button("Run Prediction")

        enriched_output = gr.Textbox(label="📎 Enriched Input (Used for Prediction)", lines=8, interactive=False)
        label_output = gr.Textbox(label="🏷️ Predicted Labels", lines=2, interactive=False)

        predict_btn.click(
            fn=run_prediction_pipeline,
            inputs=[desc_input, chat_upload],
            outputs=[enriched_output, label_output]
        )

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