import streamlit as st
import os
import openai
import json
import bcrypt
import os
from groq import Groq
import re

import warnings
warnings.filterwarnings('ignore')

client = Groq(
    api_key=os.environ.get("GROQ_API_KEY"),
)

# hiding streamlit header and menus
hide_menu = """
                <style>
                #MainMenu {visibiliyt: hidden;}
                footer {visibility: hidden;}
                header {visibility: hidden;}
                </style>"""

st.markdown(hide_menu, unsafe_allow_html = True)

# function to remove the streamlit tag from webpage title
def set_page_title(title):
    st.markdown(unsafe_allow_html=True, body=f"""
        <iframe height=0 srcdoc="<script>
            const title = window.parent.document.querySelector('title') \
                
            const oldObserver = window.parent.titleObserver
            if (oldObserver) {{
                oldObserver.disconnect()
            }} \

            const newObserver = new MutationObserver(function(mutations) {{
                const target = mutations[0].target
                if (target.text !== '{title}') {{
                    target.text = '{title}'
                }}
            }}) \

            newObserver.observe(title, {{ childList: true }})
            window.parent.titleObserver = newObserver \

            title.text = '{title}'
        </script>" />
    """)




# Define a Streamlit key for storing whether the user is authenticated.
if 'authenticated' not in st.session_state:
    st.session_state.authenticated = False

# loading user logins from config
config_file_path = "config.json"  # Path to your configuration file

with open(config_file_path) as config_file:
    credentials = json.load(config_file)


# Function to verify credentials
def check_credentials(password):
    # Check if the entered credentials match the hashed passwords in the config file
    # if username in credentials and bcrypt.checkpw(password.encode("utf-8"), credentials[username].encode("utf-8")):

    if password=='123':
        return True
    return False


# Main login page
def login_page():
    
    set_page_title("TGPT")

    st.markdown("<h1 style='text-align: center;'>Welcome to Team-GPT</h1>", unsafe_allow_html=True)
    st.markdown('\n')


    # Login form
    with st.form(key='login_form'):
        username = st.text_input("Username")
        password = st.text_input("Password", type="password")
        
        # Add some CSS styling to center-align the button
        login_button = st.form_submit_button(label='Login', type = 'primary', use_container_width = True)
        
        st.session_state.username = username

        if login_button:
            if check_credentials(password):
                pattern = r"^[a-zA-Z]+$"
                if re.fullmatch(pattern, username):
                    st.session_state['authenticated'] = True  # Set the session state
                    st.rerun()  # Rerun the app to update the state
                else:
                    st.error("Username can't contain numbers or special characters!")
            else:
                st.error("The password is incorrect!")
    
    st.write('Password: 123')

    st.markdown('\n')
    st.markdown("<span style='color:#2180b9; text-align:left; font-weight:bold; font-size:20px;'>Instructions :</span>", unsafe_allow_html=True)
    st.markdown("""
    <div style='text-align: left; padding-left: 20px;'>
    <ul>
        <li>Username can't have numbers or special characters.</li>
        <li>By default, every new conversation is Private.</li>
        <li>To make the conversation Public, select 'Shared' before submitting first prompt.</li>
        <li style='color:red;'>Once the first prompt is submitted, you can't change the chat visibility.</li> 
    </ul>
    </div>
    """, unsafe_allow_html=True)
        


# function to create new chat
def create_new_chat(st):

    st.session_state.messages = []
    
    
# function to update chat
def update_chat(st):
    
    # getting file messages
    file_var = open(st.session_state.file_name, encoding='utf-8')
    file_content = file_var.read()

    file_messages = file_content.split('-end-\n')
    file_messages = [file_messages[i].replace('-end-','') for i in range(len(file_messages))]

    # getting session_messages
    session_messages = st.session_state.messages

    # computing the message difference
    diff = len(file_messages) - len(session_messages)

    # appending the messages to update session state
    new_messages = []
    
    if diff>0:
        for i in range(len(session_messages), len(file_messages)):
            session_messages.append(file_messages[i])
            new_messages.append(file_messages[i])
            
    # displaying new messages
    for message in new_messages:
        print(message)
        with st.chat_message(message["added_by"]):
            st.markdown(message["content"])


def run_chatbot(st):
    
    set_page_title("TGPT")

    # Initialize chat history
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # default chat is New
    if "chat_type" not in st.session_state:
        st.session_state.chat_type = 'New'

    # replace current chat variable
    if "replace_current_chat" not in st.session_state:
        st.session_state.replace_current_chat = False

    # creating file name variable in session_state
    if "file_name" not in st.session_state:
        st.session_state.file_name = ''
        
    # initializing visibility as shared
    if "visibility" not in st.session_state:
        st.session_state.visibility = 'Shared👥'
        

    # sidebar
    with st.sidebar:
        if st.button("Logout", type = "secondary", key = 'logout'):
            st.session_state.clear()  # Clear all session state variables
            st.rerun()

        if st.button("New Chat ➕", type = "primary", key = 'new_chat'):
            st.session_state.messages = []
            st.session_state.file_name = ''
            st.session_state.chat_type = 'New'
            st.session_state.replace_current_chat = False
            st.session_state.visibility = 'Private🔒'
            

        # displaying chat history from history files
        st.header('Chat History')

        history_files = os.listdir('chat_history/')

        # Sort the files based on their creation time in descending order
        complete_files = sorted(history_files, key=lambda x: os.path.getctime(os.path.join('chat_history/', x)),
                              reverse=True)
        
        # removing the username and visibility from chat heading
        usersnames = [f.split('_')[0] for f in complete_files]
        visibility_status = [f.split('_')[1] for f in complete_files]
        files = [f.split('_')[2] for f in complete_files]
        

        # iterating over each history file to add as a button
        for i in range(len(complete_files)):
            
            add_chat = False
            
            # based on user who created file and visibility, deciding if chat should be visible to a user or not
            if (visibility_status[i]=='Private🔒') & (usersnames[i]==st.session_state.username):
                add_chat = True
            elif visibility_status[i]=='Shared👥':
                add_chat = True
                
            if add_chat:
            
                unique_key = f"{i}_{files[i][:-4]}"

                if st.button(files[i][:-4], key=unique_key):
                    # print('Loading : ', chat)

                    st.session_state.replace_current_chat = True
                    st.session_state.chat_type = 'Old'
                    st.session_state.file_name = 'chat_history/' + complete_files[i]
            
            else:
                continue
            


    # replacing current chat with old
    if st.session_state.replace_current_chat:

        st.session_state.messages = []
        st.session_state.visibility = st.session_state.file_name.split('_')[2]

        prev_file = open(st.session_state.file_name, encoding='utf-8')
        file_content = prev_file.read()

        msgs = file_content.split('-end-\n')

        for i in range(len(msgs)):

            msgs[i] = msgs[i].replace('-end-','')
            
            # text format -> user:username:msg (for AI, the user and username are same : assistant)
            # first word is reserved for the user type : user, assistant
            # second word is for username
            # third is text
            
            msg_content = msgs[i].split(':',2) # splitting the text into 3 parts [user, username, msg]
            user_type = msg_content[0]
            from_user = msg_content[1]
            msg_text = msg_content[2]
            
            st.session_state.messages.append({"role": user_type, "added_by" : from_user, "content": msg_text})
            
            
    # creating columns to seperate title and visibility buttons        
    header1, header2 = st.columns(2)        
    
    with header1:
        # Setting title for chat window
        st.markdown("<h4 style='font-family: Arial;'>Team-GPT v0.2</h4>", unsafe_allow_html=True)
       
    
    # Shared-private option through radio button   
    with header2:     
        
        # css for button positioning
        st.markdown(
            """
            <style>
            .stRadio > div {
                flex-direction: row;
                justify-content: flex-end;
            }
            .st-Radio input[type="radio"] + label::before {
                width: 8px;
                height: 8px;
                margin-right: 4px;
            }
            </style>
            """,
            unsafe_allow_html=True
        )
        
        # for the past chats, we just display the visibility status earlier selected
        if st.session_state.chat_type=='Old':
            option = st.radio(label = '', options = (st.session_state.visibility,))
        else:
            option = st.radio(label = '', options = ('Private🔒', 'Shared👥'))
            st.session_state.visibility = option
            
    
    # displaying all the messages
    for message in st.session_state.messages:
        with st.chat_message(message["added_by"]):
            st.markdown(message["content"])


    # user input
    prompt = st.chat_input("Type here...")

    if prompt:

        if st.session_state.chat_type=='New':
            # making sure all the characters are alpha-numeric
            file_name = "".join(c if c.isalnum() or c in {'_', ' ', '.'} else ' ' for c in prompt)
            file_name = " ".join(file_name.split()[:5])
            file_name = file_name.strip()  # removing any trailing spaces

            # naming file after user and its visibility
            file_name = "chat_history/" + st.session_state.username + '_' +  st.session_state.visibility + '_' + file_name + ".txt"    
            
            st.session_state.file_name = file_name


        
        with st.chat_message(st.session_state.username):
            st.markdown(prompt)

        # add user msg to chat history
        st.session_state.messages.append({"role":"user", "added_by":st.session_state.username, "content":prompt})

        # writing prompt to file
        # print("file : ", st.session_state.file_name)
        file = open(st.session_state.file_name, "a", encoding='utf-8')

        if st.session_state.chat_type=='New':
            file.write('user:' + st.session_state.username + ':' + prompt + "-end-")
            st.session_state.chat_type = 'Old'

        elif st.session_state.chat_type=='Old':
            file.write("\n" + 'user:' +  st.session_state.username + ':'  + prompt + "-end-")



        with st.chat_message("assistant"):

            message_placeholder = st.empty()
            full_response = ""


            response = client.chat.completions.create(
                messages = [{"role":m["role"], "content":m["content"]} for m in st.session_state.messages[-10:]],
                model= 'llama-3.1-8b-instant',
                temperature = 1, 
                max_tokens = 3000
            )

            full_response = response.choices[0].message.content

            # sending just last 10 messages with the prompt to create some context to prompt
            # for response in openai.ChatCompletion.create(model = "gpt-3.5-turbo",
            #                                              messages = [
            #                                                  {"role":m["role"], "content":m["content"]} for m in st.session_state.messages[-10:]
            #                                              ], stream = True):
            #     print(response,'\n\n')
            #     full_response += response.choices[0].delta.get("content", "")
            #     message_placeholder.markdown(full_response + "| ")

            message_placeholder.markdown(full_response)
            file.write("\n" + 'assistant:assistant:' + full_response + "-end-")

        st.session_state.messages.append({"role":"assistant", "added_by":"assistant", "content":full_response})
        print(st.session_state.messages,'\n')
        update_chat(st)

if not st.session_state['authenticated']:
    login_page()
        
else:
    run_chatbot(st)