import gradio as gr
import json
import os
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage
import re


client = MistralClient(api_key=os.environ["MISTRAL_API_KEY"])
model = "open-mixtral-8x7b"

messages = {
    'English': {
        'system_message': "You are a moderation assistant. Given a {} in the context of {}, determine if it should be flagged as spam, not spam, or possibly spam. Respond with a JSON object containing a 'spam' field with a string value of 'true', 'false', or 'possible'.",
        'user_message': "Content: {}\\nContext: {}\\n\\nShould this content be flagged as spam in the given context? Respond with a JSON object."
    },
    'Spanish': {
        'system_message': "Eres un asistente de moderación. Dado un {} en el contexto de {}, determina si debe marcarse como spam, no spam o posiblemente spam. Responde con un objeto JSON que contenga un campo 'spam' con un valor de cadena de 'true', 'false' o 'possible'.",
        'user_message': "Contenido: {}\\nContexto: {}\\n\\n¿Se debe marcar este contenido como spam en el contexto dado? Responde con un objeto JSON."
    },
    'French': {
        'system_message': "Vous êtes un assistant de modération. Étant donné un {} dans le contexte de {}, déterminez s'il doit être signalé comme spam, non spam ou éventuellement spam. Répondez avec un objet JSON contenant un champ 'spam' avec une valeur de chaîne 'true', 'false' ou 'possible'.",
        'user_message': "Contenu : {}\\nContexte : {}\\n\\nCe contenu doit-il être signalé comme spam dans le contexte donné ? Répondez avec un objet JSON."
    }
}


def make_request(client, model, system_message, user_message, temperature=0.7, top_p=0.6, max_tokens=50):
    print(f"Making request with system message: {system_message}")
    print(f"Making request with user message: {user_message}")
    messages = [
        ChatMessage(role="system", content=system_message),
        ChatMessage(role="user", content=user_message)
    ]
    response = client.chat(model=model, messages=messages, temperature=temperature, top_p=top_p, max_tokens=max_tokens)
    print(f"Received response from the model: {response}")
    return response.choices[0].message.content if response.choices else ""

def moderate_content(content, content_context, moderation_rules, language):
    print(f"Moderating content: {content}")
    print(f"Content context: {content_context}")
    print(f"Moderation rules: {moderation_rules}")
    print(f"Language: {language}")
    
    system_message = messages[language]['system_message'].format("content", content_context)
    user_message = messages[language]['user_message'].format(content, content_context)
    
    response = make_request(client, model, system_message, user_message, temperature=0.3, top_p=0.6, max_tokens=50)
    print(f"Received response: {response}")
    
    # Extract the JSON part of the response using regular expressions
    json_match = re.search(r'{.*}', response, re.DOTALL)
    if json_match:
        json_response = json_match.group()
        print(f"Extracted JSON response: {json_response}")
        try:
            response_json = json.loads(json_response)
            print(f"Parsed JSON response: {response_json}")
            if 'spam' in response_json:
                spam_status = response_json['spam'].lower()
                print(f"Spam status: {spam_status}")
                if spam_status == 'true':
                    return "Spam"
                elif spam_status == 'false':
                    return "Not Spam"
                elif spam_status == 'possible':
                    return "Possibly Spam"
        except (json.JSONDecodeError, KeyError, ValueError) as e:
            print(f"Error occurred while parsing JSON: {str(e)}")
    else:
        print("No JSON found in the response.")
    
    return "Unable to determine spam status"

def format_spam_status(spam_status):
    print(f"Formatting spam status: {spam_status}")
    if spam_status == "Spam":
        return "<div style='background-color: #ff4d4d; color: white; padding: 10px; border-radius: 5px;'><strong>Spam</strong></div>"
    elif spam_status == "Not Spam":
        return "<div style='background-color: #66bb6a; color: white; padding: 10px; border-radius: 5px;'><strong>Not Spam</strong></div>"
    elif spam_status == "Possibly Spam":
        return "<div style='background-color: #ffa726; color: white; padding: 10px; border-radius: 5px;'><strong>Possibly Spam</strong></div>"
    else:
        return spam_status


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🚨 Content Moderation Pilot 🚨
        
        This is a pilot application for content moderation using the Mistral AI API. The intended use of this moderation system is via an API integration. 🌐
        
        Please enter the content, context, moderation rules, and select the language to determine if the content is spam. 🕵️‍♀️
        """
    )
    
    with gr.Row():
        with gr.Column():
            content = gr.Textbox(label="Content", placeholder="Enter the content to moderate...")
            content_context = gr.Textbox(label="What is the context of this content?", placeholder="Provide the context...")
        with gr.Column():    
            moderation_rules = gr.Textbox(label="Moderation Rules", placeholder="Enter the moderation rules...")
            language = gr.Dropdown(["English", "Spanish", "French"], label="Language", value="English")
    
    submit_button = gr.Button("🚀 Moderate Content")
    
    gr.Markdown("---")
    
    with gr.Row():
        spam_status = gr.HTML(label="Spam Status")

    submit_button.click(
        fn=moderate_content,
        inputs=[content, content_context, moderation_rules, language],
        outputs=spam_status
    )
    
    gr.Markdown(
        """
        ## 📚 API Tutorial
        
        To use the content moderation system programmatically, you can follow these steps:
        
        1. Make a POST request to the `/call/moderate` endpoint with the following JSON payload:
        
        ```json
        {
            "data": [
                "The content to moderate",
                "The context of the content",
                "The moderation rules to apply",
                "The language of the content (English, Spanish, or French)"
            ]
        }
        ```
        
        2. The POST request will return a unique `event_id` that you can use to retrieve the results.
        
        3. Make a GET request to the `/call/moderate/<event_id>` endpoint to stream the results using server-sent events.
        
        Here's an example using `curl`:
        
        1. POST request to submit the moderation task:
        
        ```bash
        curl -X POST -H "Content-Type: application/json" -d '{
            "data": [
                "This is the content to moderate",
                "Social media post",
                "Detect spam and offensive language",
                "English"
            ]
        }' https://huggingface.co/spaces/monsimas/SpamOrNot/call/moderate
        ```
        
        2. GET request to retrieve the results:
        
        ```bash
        curl -N https://huggingface.co/spaces/monsimas/SpamOrNot/call/moderate/<event_id>
        ```
        
        Replace `<event_id>` with the actual `event_id` returned from the POST request.
        
        The GET request will stream the moderation results using server-sent events. The events will have the following structure:
        
        ```
        event: completed
        data: {"spam_status": "Spam" | "Not Spam" | "Possibly Spam"}
        ```
        
        You can integrate this API into your content moderation workflow to automatically detect and flag spam content. 🚀
        
        Note: If you're using a programming language like Python, you can use libraries like `requests` to make the API requests and handle the server-sent events.
        """
    )

demo.launch()