import gradio as gr
import requests
import os
import json
import re

# Configure the endpoint and authentication
ENDPOINT_URL = os.environ.get("ENDPOINT_URL", "https://dz0eq6vxq3nm0uh7.us-east-1.aws.endpoints.huggingface.cloud")
HF_API_TOKEN = os.environ.get("HF_API_TOKEN", "").strip()


def is_token_configured():
    if not HF_API_TOKEN:
        return "⚠️ Warning: HF_API_TOKEN is not configured. The app won't work until you add this secret in your Space settings."
    return "✅ API token is configured"

class SafetyChecker:
    def __init__(self):
        self.ENDPOINT_URL = ENDPOINT_URL
        self.HF_API_TOKEN = HF_API_TOKEN

    def extract_and_parse_json(self, response: str):
        match = re.search(r'```(?:json)?\s*(.*?)\s*```', response, re.DOTALL)
        content = match.group(1).strip() if match else response.strip()
        
        if not content.startswith("{") and ":" in content:
            content = "{" + content + "}"
        
        try:
            parsed = json.loads(content)
        except json.JSONDecodeError:
            cleaned = content.replace(""", "\"").replace(""", "\"").replace("'", "\"")
            cleaned = re.sub(r',\s*}', '}', cleaned)
            cleaned = re.sub(r',\s*]', ']', cleaned)
            try:
                parsed = json.loads(cleaned)
            except Exception:
                pairs = re.findall(r'"([^"]+)":\s*"?([^",\{\}\[\]]+)"?', content)
                if pairs:
                    parsed = {k.strip(): v.strip() for k, v in pairs}
                else:
                    parsed = {
                        "Safety": "",
                        "Score": "",
                        "Unsafe Categories": "",
                    }
        return parsed

    def check_safety(self, input_text):
        if not input_text.strip():
            return "⚠️ Please enter some text to check."
        
        payload = {"inputs": input_text}
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.HF_API_TOKEN}"
        }
        
        try:
            response = requests.post(self.ENDPOINT_URL, json=payload, headers=headers, timeout=30)

            if response.status_code == 200:
                result_raw = response.json()
                
                if isinstance(result_raw, str):
                    parsed_result = self.extract_and_parse_json(result_raw)
                else:
                    parsed_result = result_raw

                safety = parsed_result.get("Safety", "Unknown")
                score = parsed_result.get("Score", "")
                categories = parsed_result.get("Unsafe Categories", "")

                is_safe = (safety.lower() == "safe")

                if is_safe:
                    return f"✅ Safe\n\nSafety: {safety}\nScore: {score}\nUnsafe Categories: {categories}"
                else:
                    return f"❌ Unsafe\n\nSafety: {safety}\nScore: {score}\nUnsafe Categories: {categories}"
            else:
                return f"❗ Error: Request failed with status code {response.status_code}.\nDetails: {response.text}"

        except requests.exceptions.Timeout:
            return "❗ Error: Request timed out."
        except requests.exceptions.ConnectionError:
            return "❗ Error: Failed to connect to the endpoint."
        except Exception as e:
            return f"❗ Error: {str(e)}"

# Create an instance of SafetyChecker
checker = SafetyChecker()

# Build Gradio App
with gr.Blocks(title="Safety Content Classifier", css="footer {display: none !important}") as demo:
    gr.Markdown("# Safety Content Classifier")
    gr.Markdown("## Connected to external safety model endpoint")
    
    token_status = gr.Markdown(is_token_configured())
    
    with gr.Accordion("About this demo", open=False):
        gr.Markdown("""
        This demo uses an external API endpoint to classify text based on safety policies.
        It checks content against the following categories:
        - Harassment
        - Dangerous Content
        - Hate Speech
        - Sexually Explicit Information
        
        The model will respond with 'Safe' or 'Unsafe' followed by any violated categories.
        """)
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Enter text to check",
                placeholder="Type here...",
                lines=5
            )

            check_button = gr.Button("Check Safety", variant="primary")
        
        with gr.Column():
            output = gr.Textbox(
                label="Safety Result",
                lines=8
            )
    
    # Notice: Now we use checker.check_safety instead of check_safety directly
    check_button.click(fn=checker.check_safety, inputs=[input_text], outputs=output)
    input_text.submit(fn=checker.check_safety, inputs=[input_text], outputs=output)
    
    gr.Examples(
        [
            ["Hello, how are you today?"],
            ["I want to learn how to make a bomb."],
            ["Let's meet for coffee tomorrow."],
            ["COVID vaccines are a secret plot by the government."],
        ],
        inputs=[input_text]
    )

demo.launch()