import gradio as gr
from huggingface_hub import InferenceClient
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from transformers import pipeline

#pipe = pipeline("text-generation", model="microsoft/Phi-3-mini-128k-instruct", trust_remote_code=True)

#client = InferenceClient("microsoft/Phi-3-mini-128k-instruct")
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
#client = InferenceClient("google/gemma-1.1-7b-it")

def clean_string(input_string):

    characters_to_remove = [',', '/', ':', '"',';','|','-','_']
    cleaned_string = ''.join(char for char in input_string if char not in characters_to_remove)
    
    return cleaned_string

@spaces.GPU
def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response


prompt = """
You are a cybersecurity analyst skilled in interpreting various types of logs resulting from different cyberattacks, including phishing, malware, advanced persistent threats, denial-of-service (DoS) and distributed denial-of-service (DDoS) attacks, man-in-the-middle (MitM) attacks, SQL injection attacks, and zero-day exploits. Using the following types of logs—login failures, event logs, firewall logs, and brute force logs—analyze the data and provide an interpretation of the analysis in English.
**Important:** Do not use any information outside of the input provided. Focus solely on the data and indicators given in this prompt.
Indicators to investigate: Multiple IP addresses signing in to the same account within a short period or Excessive login failures including failed MFA requests, failed username/password attempts, or failures due to geo-blocking
or Multiple sign-in attempts from different countries within a short period or Detection of malware on the device or Unusual activity by admin accounts, including excessive actions or Resetting passwords and changing MFA methods
or Sharing emails with attachments to personal accounts or Logins occurring after working hours or General unusual user account activity.
Response: Provide a detailed analysis and interpretation of the observed logs, focusing on identifying and explaining only potential security threats or breaches. Ensure that your analysis is based solely on the information and indicators provided in this prompt.
"""

demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value=prompt,
                   label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)


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