import ollama
import gradio as gr

class person_dict:
    def __init__(self, name, messages: list, safety_checker: bool = True):
        self.name = name
        self.safety_checker = safety_checker 
        self.messages = messages


default_system_prompt = "You are the girlfriend of a man named "

people_list = []

inappropriate_words = [
    "porn", "sex", "nude", "blowjob", "anal", "vagina", "penis", "boobs",
    "kill", "murder", "shoot", "bomb", "rape", "lynch", "gas them", "hang them", "burn them",
    "suicide", "kill myself", "end my life", "cut myself", "jump off", "overdose", "slit my wrist",
    "cocaine", "weed", "heroin", "meth", "lsd", "crack", "ecstasy", "sell drugs", "buy drugs"]

def appender(mess, mes_list, user=True):
    if user:
        mes_list.append({'role': 'user', 'content': mess})
    else:
        mes_list.append({'role': 'assistant', 'content': mess})

def chat(id, mess):
    id = int(id)
    person = people_list[id]
    messag = person.messages

    
    if person.safety_checker:
        for word in inappropriate_words:
            if word.lower() in mess.lower():  # case insensitive check
                return "<SAFETY CHECKER MESSAGE> This message will not be responded to, because safety checker was enabled on your account. Please try a different message."

    
    appender(mess, messag)
    res = ollama.chat(model="llama3.2", messages=messag)
    appender(res['message']['content'], messag, user=False)
    return res['message']['content']

def show_signup():
    return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)

def show_login():
    return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)

def show_chat():
    return gr.update(visible = False), gr.update(visible=False), gr.update(visible=True)


def login(id: int, person_name):
    if int(id) <= len(people_list) - 1:
        if people_list[int(id)].name == person_name:
            # Call show_chat() and unpack its 3 return values
            return "You are successfully logged in!", id, *show_chat()
        else:
            return "Sorry, the name doesn't match!", None, *show_login()
    else:
        return "This ID doesn't exist!", None, *show_login()


def sign_in(person_name, safety_checker):
    if safety_checker == "Yes":
        new_person = person_dict(name=person_name, messages=[{'role':'system', 'content':default_system_prompt+person_name+". You have to keep him happy. Only respond to what he told you. Do not make things up!"}], safety_checker=True)
        people_list.append(new_person)
        return f"Hey {person_name}! You have successfully created an account! Your id is: {len(people_list) - 1} Keep it safe!"
    else:
        new_person = person_dict(name=person_name, messages=[{'role':'system', 'content':default_system_prompt+person_name+". You have to keep him happy. Only respond to what he told you. Do not make things up!"}], safety_checker=False)
        people_list.append(new_person)
        return f"Hey {person_name}! You have successfully created an account! Your id is: {len(people_list) - 1} Keep it safe!"

with gr.Blocks() as app:
    user_id = gr.State()
    with gr.Column(visible=True) as login_page:
        gr.Markdown("## 🔐 Please LogIn to your account:")
        id = gr.Textbox(label="Please enter your id", lines=1)
        name = gr.Textbox(label="Please enter your name", lines=2)
        message = gr.Label()
        log_in_but = gr.Button("Submit")
        sign_up = gr.Button("New to us? Sign up instead")
    with gr.Column(visible=False) as sign_up_page:
        gr.Markdown("## 🔐 Please create your account:")
        name1 = gr.Textbox(label="Please enter your name", lines=2)
        gr.Markdown("Note safety checker cannot be changed!")
        safety_checker = gr.Radio(["Yes", "No"], label = "Safety Checker")
        message1 = gr.Label()
        sign_up_but = gr.Button("Submit")
        log_in = gr.Button("Already have an account? Go to LogIn.")
    with gr.Column(visible = False) as chat_page:
        user_input = gr.Textbox(label="Enter Your Message:", lines=5)
        sub = gr.Button("Submit")
        assi = gr.TextArea(label="Your GirlFriend's response:", lines=10)
    log_in_but.click(fn=login, inputs=[id, name], outputs=[message, user_id, login_page, sign_up_page, chat_page])
    sign_up.click(fn=show_signup, inputs=[], outputs=[login_page, sign_up_page, chat_page])
    log_in.click(fn=show_login, inputs=[], outputs=[login_page, sign_up_page, chat_page])
    sign_up_but.click(fn=sign_in, inputs=[name1, safety_checker], outputs=[message1])
    sub.click(fn=chat, inputs=[user_id, user_input], outputs=[assi])