import gradio as gr
from transformers import pipeline

# Use lightweight model suitable for free-tier HF Space
llm = pipeline("text-generation", model="sshleifer/tiny-gpt2", max_new_tokens=100)

# Context for each role
roles = {
    "Blue Team": "You are a cybersecurity analyst defending a system.",
    "Red Team": "You are simulating a cyber attack for penetration testing.",
    "GRC": "You are an AI compliance auditor evaluating AI systems."
}

# Response generation function
def generate_response(prompt, role):
    context = roles.get(role, "")
    full_prompt = f"{context}\n{prompt}"
    try:
        result = llm(full_prompt)[0]['generated_text']
    except Exception as e:
        result = f"Error: {str(e)}"
    return result

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("📄 [View full README](https://huggingface.co/spaces/sabriallani/WorkshopCyberAI/blob/main/README.md)  \n"
                "_Learn more about the use cases, models, and architecture._")
    gr.Markdown("# 🛡️ CyberAI Prompt Generator")
    gr.Markdown("Select your role and input a prompt to get a contextual AI response.")

    with gr.Row():
        role = gr.Dropdown(choices=list(roles.keys()), label="Role", value="Blue Team")
        prompt = gr.Textbox(label="Your prompt", placeholder="e.g. Generate a threat model for a hospital IoT network")

    output = gr.Textbox(label="AI Response")

    btn = gr.Button("Generate")
    btn.click(fn=generate_response, inputs=[prompt, role], outputs=output)

demo.launch()