import gradio as gr
from autogen import AssistantAgent, UserProxyAgent, GroupChatManager
from transformers import pipeline

# Initialize sentiment analysis pipeline
sentiment_analyzer = pipeline("sentiment-analysis")

# Define specialized agents
class ThreatIntelAgent(AssistantAgent):
    def __init__(self):
        super().__init__(name="ThreatIntelAgent", system_message="You are a Threat Intelligence Agent. Analyze IoCs, malware, and threat feeds.")

    def analyze_ioc(self, ioc):
        # Placeholder for threat intelligence analysis
        return f"Analyzed IoC: {ioc}. This matches known ransomware patterns."

class ComplianceAgent(AssistantAgent):
    def __init__(self):
        super().__init__(name="ComplianceAgent", system_message="You are a Compliance Agent. Handle regulatory queries and audit support.")

    def check_compliance(self, query):
        # Placeholder for compliance checks
        return f"Compliance check for: {query}. This aligns with GDPR guidelines."

# Central AutoGen Manager
class AutoGenManager:
    def __init__(self):
        self.threat_intel_agent = ThreatIntelAgent()
        self.compliance_agent = ComplianceAgent()

    def route_request(self, user_input, user_role):
        # Sentiment analysis
        sentiment = sentiment_analyzer(user_input)[0]
        print(f"Sentiment: {sentiment}")

        # Route based on role and input
        if user_role == "Threat Analyst":
            return self.threat_intel_agent.analyze_ioc(user_input)
        elif user_role == "Compliance Officer":
            return self.compliance_agent.check_compliance(user_input)
        else:
            return "Request routed to general support."

# Gradio Interface
def chat_interface(user_input, user_role):
    manager = AutoGenManager()
    response = manager.route_request(user_input, user_role)
    return response

# Create Gradio app
with gr.Blocks() as demo:
    gr.Markdown("## AI-Powered Cybersecurity Assistant")
    with gr.Row():
        user_input = gr.Textbox(label="Your Request", placeholder="Enter your request here...")
        user_role = gr.Dropdown(
            label="Your Role",
            choices=["Threat Analyst", "Compliance Officer", "Security Engineer", "General User"],
            value="General User"
        )
    output = gr.Textbox(label="Response")
    submit_btn = gr.Button("Submit")
    submit_btn.click(fn=chat_interface, inputs=[user_input, user_role], outputs=output)

# Launch the app
if __name__ == "__main__":
    demo.launch()