import gradio as gr
from gradio_client import Client
import matplotlib.pyplot as plt
import numpy as np
import io
from PIL import Image
import base64
import os
import requests
from tenacity import retry, wait_exponential, stop_after_attempt

# Load Hugging Face token from environment variable
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("Hugging Face token not found in environment variables.")

@retry(wait=wait_exponential(multiplier=1, min=4, max=10), stop=stop_after_attempt(5))
def get_dynamic_endpoint():
    """
    Fetch the dynamic endpoint using the Hugging Face API.
    Returns:
        str: The current dynamic endpoint.
    """
    api_url = "https://api.huggingface.co/space/duchaba/friendly-text-moderation"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}

    response = requests.get(api_url, headers=headers)
    response.raise_for_status()  # Raise an error for bad status codes

    # Extract the endpoint from the response
    data = response.json()
    endpoint = data.get("url")
    if not endpoint:
        raise ValueError("Endpoint URL not found in the response.")
    return endpoint

def moderate_text(text, safer_value):
    """
    Moderate the given text using the Hugging Face API.
    
    Args:
        text (str): The text to be moderated.
        safer_value (float): The safer value to be used for moderation.
    
    Returns:
        tuple: A tuple containing the moderation result and the generated image.
    """
    try:
        # Fetch the dynamic endpoint
        dynamic_endpoint = get_dynamic_endpoint()

        # Initialize the client with the dynamic endpoint
        client = Client(dynamic_endpoint, hf_token=HF_TOKEN)
        result = client.predict(
            text,
            safer_value,
            api_name="/censor_me"
        )

        # Ensure the result contains the expected data
        base64_data = result.get('plot', '').split(',')[1] if 'plot' in result else None
        if not base64_data:
            raise ValueError("Expected plot data not found in the result.")

        # Decode base64 to bytes
        img_data = base64.b64decode(base64_data)

        # Convert bytes to PIL Image
        img = Image.open(io.BytesIO(img_data))

        return result, img
    
    except Exception as e:
        # Log the error for debugging purposes
        print(f"Error occurred: {e}")
        return str(e), None

# Define the Gradio interface
demo = gr.Interface(
    fn=moderate_text,
    inputs=[
        gr.Textbox(label="Enter Text:", placeholder="Type your text here...", lines=5),
        gr.Slider(minimum=0.005, maximum=0.1, value=0.005, label="Personalize Safer Value: (larger value is less safe)")
    ],
    outputs=[gr.Textbox(label="Moderated Text:", lines=5), gr.Image(type="pil", label="Moderation Pie Chart")],
    title="Friendly Text Moderator",
    description="Enter text to be moderated and adjust the safer value to see how it affects the moderation."
)

# Customize the CSS
custom_css = """
body {
    background-color: #f5f5f5;
    font-family: Arial, sans-serif;
}
.gradio-container {
    max-width: 800px;
    margin: auto;
    padding: 20px;
    background-color: white;
    border: 1px solid #ddd;
    border-radius: 8px;
    box-shadow: 0 2px 10px rgba(0, 0, 0, 0.1);
}
.gr-button {
    background-color: #4CAF50;
    color: white;
}
.gr-button:hover {
    background-color: #45a049;
}
"""

# Add the custom CSS to the Gradio app
demo.css = custom_css

# Launch the app
demo.launch()