import os
import requests
import streamlit as st

# Get the Hugging Face API Token from environment variables
HF_API_TOKEN = os.getenv("HF_API_KEY")
if not HF_API_TOKEN:
    raise ValueError("Hugging Face API Token is not set in the environment variables.")

# Hugging Face API URLs and headers for models
MISTRAL_API_URL = "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1"
MINICHAT_API_URL = "https://api-inference.huggingface.co/models/GeneZC/MiniChat-2-3B"
DIALOGPT_API_URL = "https://api-inference.huggingface.co/models/microsoft/DialoGPT-large"
PHI3_API_URL = "https://api-inference.huggingface.co/models/microsoft/Phi-3-mini-4k-instruct"
GEMMA_API_URL = "https://api-inference.huggingface.co/models/google/gemma-1.1-7b-it"
GEMMA_2B_API_URL = "https://api-inference.huggingface.co/models/google/gemma-1.1-2b-it"
META_LLAMA_70B_API_URL = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-70B-Instruct"
META_LLAMA_8B_API_URL = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B-Instruct"
GEMMA_27B_API_URL = "https://api-inference.huggingface.co/models/google/gemma-2-27b"
GEMMA_27B_IT_API_URL = "https://api-inference.huggingface.co/models/google/gemma-2-27b-it"

HEADERS = {"Authorization": f"Bearer {HF_API_TOKEN}"}

def query_model(api_url, payload):
    response = requests.post(api_url, headers=HEADERS, json=payload)
    return response.json()

def add_message_to_conversation(user_message, bot_message, model_name):
    st.session_state.conversation.append((user_message, bot_message, model_name))

# Streamlit app
st.set_page_config(page_title="Multi-LLM Chatbot Interface", layout="wide")
st.title("Multi-LLM Chatbot Interface")
st.write("Multi LLM-Chatbot Interface")

# Initialize session state for conversation and model history
if "conversation" not in st.session_state:
    st.session_state.conversation = []
if "model_history" not in st.session_state:
    st.session_state.model_history = {model: [] for model in [
        "Mistral-8x7B", "MiniChat-2-3B", "DialoGPT (GPT-2-1.5B)", "Phi-3-mini-4k-instruct",
        "Gemma-1.1-7B", "Gemma-1.1-2B", "Meta-Llama-3-70B-Instruct", "Meta-Llama-3-8B-Instruct",
        "Gemma-2-27B", "Gemma-2-27B-IT"
    ]}

# Dropdown for LLM selection
llm_selection = st.selectbox("Select Language Model", [
    "Mistral-8x7B", "MiniChat-2-3B", "DialoGPT (GPT-2-1.5B)", "Phi-3-mini-4k-instruct",
    "Gemma-1.1-7B", "Gemma-1.1-2B", "Meta-Llama-3-70B-Instruct", "Meta-Llama-3-8B-Instruct",
    "Gemma-2-27B", "Gemma-2-27B-IT"
])

# User input for question
question = st.text_input("Question", placeholder="Enter your question here...")

# Handle user input and LLM response
if st.button("Send") and question:
    try:
        with st.spinner("Waiting for the model to respond..."):
            chat_history = " ".join(st.session_state.model_history[llm_selection]) + f"User: {question}\n"
            if llm_selection == "Mistral-8x7B":
                response = query_model(MISTRAL_API_URL, {"inputs": chat_history})
                answer = response[0].get("generated_text", "No response") if isinstance(response, list) else "No response"
            elif llm_selection == "MiniChat-2-3B":
                response = query_model(MINICHAT_API_URL, {"inputs": chat_history})
                if "error" in response and "is currently loading" in response["error"]:
                    answer = f"Model is loading, please wait {response['estimated_time']} seconds."
                else:
                    answer = response[0].get("generated_text", "No response") if isinstance(response, list) else "No response"
            elif llm_selection == "DialoGPT (GPT-2-1.5B)":
                response = query_model(DIALOGPT_API_URL, {"inputs": chat_history})
                answer = response.get("generated_text", "No response") if isinstance(response, dict) else response[0].get("generated_text", "No response") if isinstance(response, list) else "No response"
            elif llm_selection == "Phi-3-mini-4k-instruct":
                response = query_model(PHI3_API_URL, {"inputs": chat_history})
                answer = response[0].get("generated_text", "No response") if isinstance(response, list) else "No response"
            elif llm_selection == "Gemma-1.1-7B":
                response = query_model(GEMMA_API_URL, {"inputs": chat_history})
                answer = response.get("generated_text", "No response") if isinstance(response, dict) else response[0].get("generated_text", "No response") if isinstance(response, list) else "No response"
            elif llm_selection == "Gemma-1.1-2B":
                response = query_model(GEMMA_2B_API_URL, {"inputs": chat_history})
                answer = response.get("generated_text", "No response") if isinstance(response, dict) else response[0].get("generated_text", "No response") if isinstance(response, list) else "No response"
            elif llm_selection == "Meta-Llama-3-70B-Instruct":
                response = query_model(META_LLAMA_70B_API_URL, {"inputs": chat_history})
                answer = response.get("generated_text", "No response") if isinstance(response, dict) else response[0].get("generated_text", "No response") if isinstance(response, list) else "No response"
            elif llm_selection == "Meta-Llama-3-8B-Instruct":
                response = query_model(META_LLAMA_8B_API_URL, {"inputs": chat_history})
                answer = response.get("generated_text", "No response") if isinstance(response, dict) else response[0].get("generated_text", "No response") if isinstance(response, list) else "No response"
            elif llm_selection == "Gemma-2-27B":
                response = query_model(GEMMA_27B_API_URL, {"inputs": chat_history})
                answer = response.get("generated_text", "No response") if isinstance(response, dict) else response[0].get("generated_text", "No response") if isinstance(response, list) else "No response"
            elif llm_selection == "Gemma-2-27B-IT":
                response = query_model(GEMMA_27B_IT_API_URL, {"inputs": chat_history})
                answer = response.get("generated_text", "No response") if isinstance(response, dict) else response[0].get("generated_text", "No response") if isinstance(response, list) else "No response"

            add_message_to_conversation(question, answer, llm_selection)
            st.session_state.model_history[llm_selection].append(f"User: {question}\n{llm_selection}: {answer}\n")
    except ValueError as e:
        st.error(str(e))

# Custom CSS for chat bubbles
st.markdown(
    """
    <style>
    .chat-bubble {
        padding: 10px 14px;
        border-radius: 14px;
        margin-bottom: 10px;
        display: inline-block;
        max-width: 80%;
        color: black;
    }
    .chat-bubble.user {
        background-color: #dcf8c6;
        align-self: flex-end;
    }
    .chat-bubble.bot {
        background-color: #fff;
        align-self: flex-start;
    }
    .chat-container {
        display: flex;
        flex-direction: column;
        gap: 10px;
        margin-top: 20px;
    }
    </style>
    """,
    unsafe_allow_html=True
)

# Display the conversation
st.write('<div class="chat-container">', unsafe_allow_html=True)
for user_message, bot_message, model_name in st.session_state.conversation:
    st.write(f'<div class="chat-bubble user">You: {user_message}</div>', unsafe_allow_html=True)
    st.write(f'<div class="chat-bubble bot">{model_name}: {bot_message}</div>', unsafe_allow_html=True)
st.write('</div>', unsafe_allow_html=True)