import gradio as gr
from transformers import pipeline

# Load the instruction-tuned model
model = pipeline(
    "text2text-generation",
    model="google/flan-t5-base"
)

# Allowed vocational domains
ALLOWED_TOPICS = [
    "fashion", "tailoring", "sewing", "clothing",
    "catering", "cooking", "baking", "food",
    "hair", "hairdressing", "barbing",
    "shoe", "shoemaking", "leather",
    "bag", "bag making",
    "hat", "millinery",
    "decoration", "event decoration",
    "dtf", "printing", "craft"
]

# Unsafe keywords (basic content filtering)
UNSAFE_KEYWORDS = [
    "poison", "acid", "chemical", "drug",
    "weapon", "bomb", "kill", "hack",
    "medical", "treatment", "diagnosis",
    "legal", "court", "lawsuit"
]

def is_vocational(question):
    question_lower = question.lower()
    return any(topic in question_lower for topic in ALLOWED_TOPICS)

def is_unsafe(question):
    question_lower = question.lower()
    return any(word in question_lower for word in UNSAFE_KEYWORDS)

def vocational_ai(question):
    if len(question.strip()) == 0:
        return "Please enter a vocational-related question."

    if is_unsafe(question):
        return (
            "⚠️ This question may involve unsafe or restricted content.\n\n"
            "I can only provide safe, educational vocational guidance."
        )

    if not is_vocational(question):
        return (
            "❌ I am a vocational training assistant.\n\n"
            "Please ask questions related to skills like fashion design, catering, "
            "hairdressing, shoemaking, bag making, decoration, or printing."
        )

    prompt = f"""
    You are an AI-powered vocational training assistant.

    Rules:
    - Give beginner-friendly explanations
    - Focus on safety and best practices
    - Do NOT give medical, legal, or dangerous advice
    - Provide structured steps where appropriate
    - End with a safety reminder

    Question:
    {question}
    """

    response = model(prompt, max_length=250)
    return response[0]["generated_text"] + "\n\n⚠️ This information is for educational purposes only."

# Gradio Interface
interface = gr.Interface(
    fn=vocational_ai,
    inputs=gr.Textbox(
        lines=3,
        placeholder="Ask about fashion, catering, hairdressing, shoemaking, bag making, etc."
    ),
    outputs="text",
    title="AI-Powered Vocational Training Assistant",
    description=(
        "An AI assistant that provides safe, beginner-friendly guidance "
        "for vocational skills. Available 24/7."
    ),
)

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