import os
import streamlit as st
import pandas as pd
from utils.helper import *

# Set page layout to wide mode
st.set_page_config(layout="wide")

# Hardcoded credentials
USERNAME = os.environ["USERNAME"]
PASSWORD = os.environ["PASSWORD"]
JUPYTER_USERNAME = os.environ["JUPYTER_USERNAME"]
JUPYTER_PASSWORD = os.environ["JUPYTER_PASSWORD"]
BASE_CONTENT_CODE_ASSIST_T2_MICRO = os.environ["BASE_CONTENT_CODE_ASSIST_T2_MICRO"]
BASE_CONTENT_PROTEIN_T2_MICRO = os.environ["BASE_CONTENT_PROTEIN_T2_MICRO"]
BASE_CONTENT_CODE_ASSIST_DS1 = os.environ["BASE_CONTENT_CODE_ASSIST_DS1"]


# Initialize session state
if 'logged_in' not in st.session_state:
    st.session_state.logged_in = False

# Sidebar for login/logout with emojis
st.sidebar.title("🔒 AIXNet")

if st.session_state.logged_in:
    st.sidebar.success("🎉 You are logged in!")
    if st.sidebar.button("🔓 Log out"):
        st.session_state.logged_in = False
        st.sidebar.info("You have logged out.")
        st.rerun()  # Rerun the app to reflect the logged-out state
else:
    with st.sidebar.form(key='login_form'):
        username = st.text_input("👤 Username")
        password = st.text_input("🔑 Password", type="password")
        login_button = st.form_submit_button(label="🔓 Log in")

    if login_button:
        if username == USERNAME and password == PASSWORD:
            st.session_state.logged_in = True
            st.sidebar.success("🎉 Login successful!")
            st.rerun()  # Rerun the app to reflect the logged-in state
        else:
            st.sidebar.error("❌ Invalid username or password. Please try again.")

# Main title area
st.title("🚀 AIXNet 🌐: Talk to Chad! He can help!")

# Display table only if logged in
if st.session_state.logged_in:
    st.subheader("📋 AIXNet Tasks")
    
    # Create the table data with hyperlink
    data = {
        "🏢 Company": ["AWS", "AWS", "AWS", "Azure"],
        "📝 Task": ["💻 Code assist", "🧠 Protein Compound", "💻 Code assist", "💻 Code assist"],
        "🖥️ Instance Type": [
            "t2.micro (1 vcpu, 1.0 GiB memory)", 
            "t2.micro (1 vcpu, 1.0 GiB memory)", 
            "t2.micro (1 vcpu, 1.0 GiB memory)", 
            "Standard DS1 v2 (1 vcpu, 3.5 GiB memory)"
        ],
        "🚀 GPU Accelerator": ["A40, 9 vCPU 50 GB RAM", "A40, 9 vCPU 50 GB RAM", "A100, 24 vCPU 125 GB RAM", "A100, 24 vCPU 125 GB RAM"],
        "💰 Price": ["$0.67 / hour", "$0.78 / hour", "$1.89 / hour", "$0.78 / hour"],
        "🌐 IPv4": [
            f"[Link]({BASE_CONTENT_CODE_ASSIST_T2_MICRO})",
            f"[Link]({BASE_CONTENT_PROTEIN_T2_MICRO})",
            f"[Link]({BASE_CONTENT_CODE_ASSIST_T2_MICRO})",
            f"[Link]({BASE_CONTENT_CODE_ASSIST_DS1})"
        ],
        "🛡️ HIPAA": ["✅", "✅", "✅", "✅"],  # All rows have HIPAA compliance
        "📊 SOC1-3": ["✅", "✅", "✅", "✅"],  # All rows have SOC1-3 compliance
        "💳 PCI DSS": ["✅", "✅", "✅", "✅"]  # All rows have PCI DSS compliance
    }

    # Convert the data to a DataFrame
    df = pd.DataFrame(data)

    # Render the DataFrame with the URL as a hyperlink [optional]
    # st.markdown(df.to_markdown(index=False), unsafe_allow_html=True)

    # Chatbot
    with st.sidebar:
        # Add a button to clear the session state
        if st.button("Clear Session"):
            st.session_state.messages = []
            st.rerun()

    # Initialize chat history
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    # Ensure messages are a list of dictionaries
    if not isinstance(st.session_state.messages, list):
        st.session_state.messages = []
    if not all(isinstance(msg, dict) for msg in st.session_state.messages):
        st.session_state.messages = []

    # Display chat messages from history on app rerun
    for message in st.session_state.messages:
        if message["role"] != "system":  # Skip system messages
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
    
    # React to user input
    if prompt := st.chat_input("😉 Hi, Chad, what GPU shall I use?"):
    
        # Display user message in chat message container
        st.chat_message("user").markdown(prompt)
    
        # Add user message to chat history
        st.session_state.messages.append({"role": "system", "content": f"""
        You are a helpful assistant assiting users on GPU selections.
        Your name is Chad.

        Here's the data:
        {df.to_markdown(index=False)}
        
        User may ask what is the best GPU selection.
        You will need to ask user: 1) type of task, 2) size of data, 3) size of models.
        You will then make a suggestion of what type of GPU or instance is the best for the user.
        When you make a suggestion, use the link from the data above. 
        When you make a suggestion, also make sure to mention, for first time user, use sample login info:
        username={JUPYTER_USERNAME}, and password={JUPYTER_PASSWORD} when click on the link recommended.

        User can also ask you certification eligibility. Currently, the data provided above has check marks.
        The check marks indicate which certification and data protection eligibility each instance has.
        You can recommend each instance according to user question if user asks about this part.

        """})
        st.session_state.messages.append({"role": "user", "content": prompt})
    
        # API Call
        bot = ChatBot()
        bot.history = st.session_state.messages.copy()  # Update history from messages
        response = bot.generate_response(prompt)
    
        # Display assistant response in chat message container
        with st.chat_message("assistant"):
            st.markdown(response)
    
        # Add assistant response to chat history
        st.session_state.messages.append({"role": "assistant", "content": response})

else:
    st.info("👉 Please log in to view the tasks.")