import streamlit as st
import google.generativeai as genai
import os
import PyPDF2 as pdf
from dotenv import load_dotenv

load_dotenv()

genai.configure(api_key=("AIzaSyDziGvuT1woHnH4_S3L_zQZV55Yj-123A8"))

#alternative key
#genai.configure(api_key=("AIzaSyAr3d_7fp0wMxuUrnf_tATknu_TRPKDdxg"))

# gemini function for general content generation
def get_gemini_response(input):
    model = genai.GenerativeModel('gemini-pro')
    response = model.generate_content(input)
    return response

# convert pdf to text
def input_pdf_text(uploaded_file):
    reader = pdf.PdfReader(uploaded_file)
    text = ""
    for page in range(len(reader.pages)):
        page = reader.pages[page]
        text += str(page.extract_text())
    return text

# malware detection function
def detect_malware(input_text):
    malware_prompt = f"""
    ### As a cybersecurity expert, your task is to analyze the following text for any indications of malware.
    ### Text:
    {input_text}
    ### Analysis Output:
    1. Identify any potential malware-related content.
    2. Explain the reasoning behind your identification.
    3. Provide recommendations for mitigating any identified risks.
    """
    response = get_gemini_response(malware_prompt)
    return response

# chatbot function
def chatbot_response(user_input):
    chatbot_prompt = f"""
    ### You are an intelligent and friendly chatbot. Engage in a meaningful conversation with the user.
    ### User Input:
    {user_input}
    ### Chatbot Response:
    """
    response = get_gemini_response(chatbot_prompt)
    return response

# Function to parse and display response content
def display_response_content(response):
    st.subheader("Response Output")
    if response and response.candidates:
        response_content = response.candidates[0].content.parts[0].text if response.candidates[0].content.parts else ""
        sections = response_content.split('###')
        for section in sections:
            if section.strip():
                section_lines = section.split('\n')
                section_title = section_lines[0].strip()
                section_body = '\n'.join(line.strip() for line in section_lines[1:] if line.strip())
                if section_title:
                    st.markdown(f"**{section_title}**")
                if section_body:
                    st.write(section_body)
    else:
        st.write("No response received from the model.")

## Streamlit App
st.title("AI-Powered Security and Chatbot System")
st.text("Use the AI system for malware detection and Awaring yourself.")

# Tabs for different functionalities
tab1, tab2 = st.tabs(["Malware Detection", "Chatbot"])

with tab1:
    st.header("Malware Detection")
    uploaded_file = st.file_uploader("Upload a file for malware detection", type="pdf", help="Please upload a PDF file.")
    submit_malware = st.button('Analyze for Malware')
    
    if submit_malware:
        if uploaded_file is not None:
            text = input_pdf_text(uploaded_file)
            response = detect_malware(text)
            
            # Parse and display response in a structured way
            display_response_content(response)

with tab2:
    st.header("Chatbot")
    user_input = st.text_input("Type your message here")
    submit_chat = st.button('Send')

    if submit_chat:
        if user_input:
            response = chatbot_response(user_input)
            
            # Parse and display response in a structured way
            display_response_content(response)