import gradio as gr
import spaces
from api.app import app as fastapi_app
import api.app as api_module
from inference.predictor import CommentModerator

# Ensure model is initialized for both API and UI
if api_module.moderator is None:
    api_module.moderator = CommentModerator()
moderator = api_module.moderator

# ZeroGPU decorator tells Hugging Face this space uses GPU
@spaces.GPU
def test_moderation(comment: str):
    if not comment or not comment.strip():
        return "Please enter a comment.", {}
    detail = moderator.predict(comment)
    status = "🚨 FLAGGED (Harmful)" if detail["flagged"] else "✅ SAFE"
    scores = detail["categories"]
    return (
        f"Status: {status} (Top: {detail['top_category']}, Confidence: {detail['confidence']:.2f})",
        scores,
    )

# Interactive UI demo
with gr.Blocks(title="Comment Moderation API") as demo:
    gr.Markdown("# 🛡️ Comment Moderation System (FastAPI + ONNX)")
    gr.Markdown(
        "Direct REST API endpoints: `POST /moderate`, `POST /moderate/detail`, and `GET /health`."
    )
    
    with gr.Row():
        input_text = gr.Textbox(lines=2, placeholder="Type a comment to test...", label="Input Comment")
    btn = gr.Button("Moderate Comment", variant="primary")
    
    with gr.Row():
        status_output = gr.Textbox(label="Verdict")
        scores_output = gr.Label(label="Category Probabilities")
        
    btn.click(fn=test_moderation, inputs=input_text, outputs=[status_output, scores_output])

# Mount all FastAPI routes (/moderate, /moderate/detail, /health) onto Gradio's server
demo.app.include_router(fastapi_app.router)

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