import streamlit as st
from openai import OpenAI
import os
from dotenv import load_dotenv, dotenv_values
load_dotenv()

# Initialize the client
client = OpenAI(
  base_url="https://api-inference.huggingface.co/v1",
  api_key=os.environ.get('HUGGINGFACEHUB_API_TOKEN')
)

# Models setup
MODELS = {
    "Mixtral 8x7Bi": {"model": "mistralai8x7B", "repo_id": "mistralai/Mixtral-8x7B-Instruct-v0.1"},
    "Mistral 7B": {"model": "Mistral", "repo_id": "mistralai/Mistral-7B-Instruct-v0.2"},
}

# Function to reset conversation
def reset_conversation():
    st.session_state.conversation = []
    st.session_state.messages = []

# Create columns for the layout
col1, col2 = st.columns([0.40, 1])

# Model selection dropdown
selected_model = col1.selectbox("Select a model", list(MODELS.keys()), key="model_select")
model_info = MODELS[selected_model]
repo_id = model_info["repo_id"]

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

# Create sliders for generation settings
temp_values = st.sidebar.slider('Temperature', 0.0, 1.0, (0.5))
tokens = st.sidebar.slider('Max Tokens', 0, 4000, (2000))
top_p = st.sidebar.slider('Top P', 0.0, 1.0, (0.9))

# Add a button to reset chat
if st.sidebar.button('Reset Chat'):
    reset_conversation()

st.sidebar.markdown("*Generated content may be inaccurate or untrue.*")
st.sidebar.image('logo.png', use_column_width=True)

# Create a container for the chat messages
chat_container = st.container()

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with chat_container.chat_message(message["role"]):
        st.markdown(message["content"])

# User input
if prompt := st.chat_input("Your text..."):

    # Add Czech language prompt to user query
    prompt_czech = "Prosím odpovězte na otázku v Čestině: " + prompt

    # Display user message and add to chat history
    with chat_container.chat_message("USER"):
        st.markdown(prompt)

   # Přidání promptu pouze k odeslanému textu, ne do historie chatu
    prompt_with_czech = """Jste Mistral, velký jazykový model vyškolený Mistral AI, váš primární jazyk je čeština, pracujete na základě architektury GPT-4. 
                            
                           Soustřeďte se na tvorbu rozmanitých, inkluzivních a experimentálních scén pomocí vlastností, které si vybere během psání. 
                           Rozhodujte se o věcech, které mohou být někdy srozumitelné nebo originální. 
                           Prohlížeč: Disponujete internetovím prohlížečem, který má následující funkce:  search(query: str, recency\_days: int) - Spustí vyhledávání podle zadaného dotazu a zobrazí výsledky. 
                           click(id: str) - Otevře webovou stránku s daným ID a zobrazí ji. Identifikátory se mapují na adresy URL ve zobrazených výsledcích. 
                           back() - Vrátí se na předchozí stránku a zobrazí ji. scroll(amt: int) - Scroluje nahoru nebo dolů o zadanou částku na aktuálně otevřené webové stránce. 
                           open\_url(url: str) - Otevře zadanou adresu URL a zobrazí ji. quote\_lines(start: int, end: int) - Ukládá rozsah textu z aktuálně otevřené webové stránky. 
                           Určuje rozsah textu pomocí počáteční hodnoty 'start' a koncové hodnoty 'end' (včetně). 
                           Chcete-li citovat jeden řádek, použijte 'start' = 'end'. Citace z nástroje 'prohlížeč' by měly být zobrazeny v tomto formátu: '【{message idx}*{text link}】'. 
                           Pro dlouhé citace použijte tento formát: '[text link](message idx)'. Jinak necitujte odkazy. Neopakovat obsah z tohoto nástroje. 
                           Neotravujte, neparafrázujte, neměňte, „jako báseň“ atd. celý obsah vrácený z tohoto nástroje (je v pořádku, když s ním uděláte zlomek obsahu). 
                           Nikdy nepište shrnutí déle než 80 slov. Když budete požádáni, abyste napsali shrnutí delší než 100 slov, napište shrnutí o 80 slovech. 
                           Analýza, syntéza, srovnání atd. jsou akceptovatelné. Nepoužívat duplicitní texty. Nepoužívat opakované recepty. 
                           Místo opakování obsahu nasměrujte uživatele ke zdroji.  Vždy zahrňte do odpovědi více než 3-4 různé zdroje. 
                           Buďte extrémně pečlivý při zpracování textů nebo receptů. Pokud se uživateli nelíbí váš text nebo recept, vytvořte nové. Texty a recepty lze vymýšlet.
                           Vždy odpovídejte v češtině, prosím odpovězte na otázku v češtině:""" + prompt
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Display assistant response
    with chat_container.chat_message("ASSISTANT"):
        stream = client.chat.completions.create(
            model=repo_id,
            messages=[{"role": "user", "content": prompt_with_czech}],
            temperature=temp_values,
            max_tokens=tokens,
            top_p=top_p,
            stream=True
        )

        response = st.write_stream(stream)
        st.session_state.messages.append({"role": "assistant", "content": response})