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"]
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"]


# 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 🌐")

# Display table only if logged in
if st.session_state.logged_in:
    st.subheader("📋 AIXNet Tasks")
    
    # Create the table data with hyperlink
    data = {
        "📝 Task": ["💻 Code assist", "🧠 Protein Compound"],
        "🖥️ Instance Type": ["t2.micro", "t2.micro"],
        "🚀 GPU Accelerator": ["A40, 9 vCPU 50 GB RAM", "A40, 9 vCPU 50 GB RAM"],
        "💰 Price": ["$0.67 / hour", "$0.78 / hour"],
        "🌐 IPv4": [
            f"[Link]({BASE_CONTENT_CODE_ASSIST_T2_MICRO})",
            f"[Link]({BASE_CONTENT_PROTEIN_T2_MICRO})"]
    }

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

    # Render the DataFrame with the URL as a hyperlink
    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.experimental_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:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    # React to user input
    if prompt := st.chat_input("😉 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()
        st.session_state.messages.append({"role": "user", "content": prompt})
    
        # API Call
        bot = ChatBot(
            protocol={"role": "system", "content": f"""
        You are a helpful assistant assiting users on GPU selections.
        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.
        """}
        )
        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.")