import os

import streamlit as st
from dotenv import load_dotenv
from langchain_community.llms import VLLMOpenAI

from preprocessing import preprocess_pipeline

load_dotenv()


def disable_input():
    st.session_state.translating = True


def translate(prompt, top_p, temperature, repetition_penalty):
    llm = VLLMOpenAI(
        openai_api_key="EMPTY",
        openai_api_base=os.getenv("VISTRAL_7B_MT_SERVER"),
        model_name="nguyen1207/Vistral-7B-MT-0.1",
        temperature=temperature,
        top_p=top_p,
        frequency_penalty=repetition_penalty,
    )

    stream = llm.stream(
        prompt,
    )

    count = 0

    for response in stream:
        if count < 3:
            count += 1
            yield ""
        else:
            yield response


st.set_page_config(page_title="Vietnamese to English Translation")

st.title(
    "🇻🇳 Vietnamese to 🇺🇸 English Translation but with Teencode and Slang understanding 🤯"
)

st.sidebar.header("Translation Parameters")
top_p = st.sidebar.slider("Top p", min_value=0.0, max_value=1.0, value=0.95)
temperature = st.sidebar.slider("Temperature", min_value=0.0, max_value=2.0, value=0.3)
repetition_penalty = st.sidebar.slider(
    "Repetition Penalty", min_value=1.0, max_value=3.0, value=1.05
)

if "messages" not in st.session_state:
    st.session_state.messages = []
    st.session_state.translating = False

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])


if user_input := st.chat_input(
    "Vietnamese text goes here... 🇻🇳",
    disabled=st.session_state.translating,
    on_submit=disable_input,
):
    if user_input.strip() != "":
        st.session_state.translating = True
        preprocessed_input = preprocess_pipeline(user_input)

        st.session_state.messages.append({"role": "user", "content": user_input})
        with st.chat_message("user"):
            st.markdown(user_input)

        with st.chat_message("assistant"):
            prompt_template = """<s> [INST] Dịch câu sau từ tiếng Việt sang tiếng Anh:

Tiếng Việt: {} [/INST] """

            prompt = prompt_template.format(preprocessed_input)

            stream = translate(prompt, top_p, temperature, repetition_penalty)

            translation = st.write_stream(stream)
            st.markdown(translation)

        st.session_state.messages.append({"role": "assistant", "content": translation})

        # Reset the input field
        st.session_state.translating = False
        st.rerun()