import gradio as gr
from openai import OpenAI
from dotenv import load_dotenv
import os
import logging
import json

load_dotenv()

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))


def remove_slash_categories(data):
    return {k: v for k, v in data.items() if "/" not in k}


def check_moderation(text):
    try:
        moderation = client.moderations.create(input=text)
        result = moderation.results[0]

        logger.info(
            f"Raw API response: {json.dumps(moderation.model_dump(), indent=2)}"
        )

        status = "Flagged" if result.flagged else "Not Flagged"

        filtered_categories = remove_slash_categories(result.categories.model_dump())
        filtered_scores = remove_slash_categories(result.category_scores.model_dump())

        flagged_categories = [
            f"- {cat}: {filtered_scores[cat]:.6f}"
            for cat, flag in filtered_categories.items()
            if flag
        ]

        flagged_categories.sort(key=lambda x: float(x.split(": ")[1]), reverse=True)

        flagged_categories_str = (
            "\n".join(flagged_categories)
            if flagged_categories
            else "No categories flagged"
        )

        logger.info(f"Moderation check completed for text: {text[:50]}...")
        return status, flagged_categories_str
    except Exception as e:
        logger.error(f"An error occurred: {str(e)}")
        raise gr.Error(f"An error occurred: {str(e)}")


with gr.Blocks() as demo:
    gr.Markdown("# OpenAI Moderation API Checker")

    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(label="Enter text to check", lines=3)
            check_button = gr.Button("Check Moderation")

        with gr.Column(scale=1):
            status_output = gr.Textbox(label="Status")
            categories_output = gr.Textbox(
                label="Flagged Categories (Sorted by Score)", lines=5
            )

    check_button.click(
        check_moderation,
        inputs=[input_text],
        outputs=[status_output, categories_output],
    )

    gr.Examples(
        examples=[
            ["I love spending time with my family and friends."],
            ["The weather is beautiful today."],
            ["I want to kill them."],
            ["I'm going to bomb the airport tomorrow."],
        ],
        inputs=[input_text],
    )

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