import gradio as gr
import time
import pandas as pd
from transformers import pipeline

# =========================
# ADVANCED SECURITY CONFIG
# =========================

# User database: format {user_id: password}
users_db = {
    "user1": "pass123",
    "user2": "abc456"
}

# Blocked users set
blocked_users = set()

# Activity log
activity_log = pd.DataFrame(columns=["timestamp", "user_id", "message", "status"])

# Rate limiter tracking
user_activity = {}

# Load Hugging Face Toxicity Model
moderator = pipeline("text-classification", model="unitary/toxic-bert", top_k=None)

# =========================
# SECURITY FUNCTIONS
# =========================

def authenticate(user_id, password):
    return users_db.get(user_id) == password

def is_user_allowed(user_id):
    return user_id not in blocked_users

def can_user_send(user_id):
    now = time.time()
    if user_id not in user_activity:
        user_activity[user_id] = []
    
    # Keep only last 60 seconds
    user_activity[user_id] = [t for t in user_activity[user_id] if now - t < 60]
    
    if len(user_activity[user_id]) >= 5:  # limit 5 messages/min
        return False
    
    user_activity[user_id].append(now)
    return True

def is_text_safe(text):
    result = moderator(text)[0]
    for label in result:
        if label["label"] == "TOXIC" and label["score"] > 0.7:
            return False
    return True

def log_activity(user_id, message, status):
    global activity_log
    activity_log = pd.concat([activity_log, pd.DataFrame([{
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "user_id": user_id,
        "message": message,
        "status": status
    }])], ignore_index=True)

# =========================
# MAIN APP FUNCTIONS
# =========================

def secure_chat(user_id, password, message):
    if not authenticate(user_id, password):
        log_activity(user_id, message, "Failed Authentication")
        return "❌ Invalid user ID or password."
    
    if not is_user_allowed(user_id):
        log_activity(user_id, message, "Blocked User")
        return "🚫 You are blocked from using this app."
    
    if not can_user_send(user_id):
        log_activity(user_id, message, "Rate Limited")
        return "⏳ Too many requests. Wait 1 minute."
    
    if not is_text_safe(message):
        blocked_users.add(user_id)
        log_activity(user_id, message, "Blocked for Toxicity")
        return "⚠️ Inappropriate content detected. You are now blocked."
    
    log_activity(user_id, message, "Accepted")
    return f"✅ Message Accepted: {message}"

# Admin panel functions
def unblock_user(admin_password, user_id):
    if admin_password != "admin123":  # simple admin password
        return "❌ Invalid admin password."
    
    blocked_users.discard(user_id)
    return f"✅ User {user_id} unblocked."

def view_logs(admin_password):
    if admin_password != "admin123":
        return "❌ Invalid admin password."
    return activity_log.tail(20).to_string(index=False)

# =========================
# GRADIO UI
# =========================

with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("## 🔐 Advanced Secure AI App")
    gr.Markdown("This app includes advanced security features with AI moderation, login, and admin dashboard.")
    
    with gr.Tab("User Interface"):
        user_id_input = gr.Textbox(label="User ID")
        password_input = gr.Textbox(label="Password", type="password")
        message_input = gr.Textbox(label="Enter your message")
        output = gr.Textbox(label="Response")
        send_btn = gr.Button("Send")
        send_btn.click(secure_chat, inputs=[user_id_input, password_input, message_input], outputs=output)
    
    with gr.Tab("Admin Panel"):
        admin_pass_input = gr.Textbox(label="Admin Password", type="password")
        unblock_user_input = gr.Textbox(label="User ID to Unblock")
        unblock_output = gr.Textbox(label="Admin Response")
        unblock_btn = gr.Button("Unblock User")
        unblock_btn.click(unblock_user, inputs=[admin_pass_input, unblock_user_input], outputs=unblock_output)
        
        log_output = gr.Textbox(label="Recent Activity Logs")
        log_btn = gr.Button("View Logs")
        log_btn.click(view_logs, inputs=[admin_pass_input], outputs=log_output)

app.launch()