import base64

import streamlit as st
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

from model.funcs import execution_time


def get_base64(file_path):
    with open(file_path, "rb") as file:
        base64_bytes = base64.b64encode(file.read())
        base64_string = base64_bytes.decode("utf-8")
    return base64_string


def set_background(png_file):
    bin_str = get_base64(png_file)
    page_bg_img = (
        """
    <style>
    .stApp {
    background-image: url("data:image/png;base64,%s");
    background-size: cover;
    }
    </style>
    """
        % bin_str
    )
    st.markdown(page_bg_img, unsafe_allow_html=True)


set_background("text_generation.png")


@st.cache_data
def load_model():
    model_path = "17/"
    model_name = "sberbank-ai/rugpt3small_based_on_gpt2"
    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
    model = GPT2LMHeadModel.from_pretrained(model_path)
    return tokenizer, model


tokenizer, model = load_model()


@execution_time
def generate_text(
    prompt, num_beams=2, temperature=1.5, top_p=0.9, top_k=3, max_length=150
):
    prompt = tokenizer.encode(prompt, return_tensors="pt")
    model.eval()
    with torch.no_grad():
        out = model.generate(
            prompt,
            do_sample=True,
            num_beams=num_beams,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            max_length=max_length,
        )
    out = list(map(tokenizer.decode, out))[0]
    return out


with st.sidebar:
    num_beams = st.slider("Number of Beams", min_value=1, max_value=5, value=2)
    temperature = st.slider("Temperature", min_value=0.1, max_value=2.0, value=1.5)
    top_p = st.slider("Top-p", min_value=0.1, max_value=1.0, value=0.9)
    top_k = st.slider("Top-k", min_value=1, max_value=10, value=3)
    max_length = st.slider("Maximum Length", min_value=20, max_value=300, value=150)

styled_text = """
<style>
.styled-text {
    font-size: 30px;
    text-shadow: -2px -2px 4px #000000;
    color: #FFFFFF;
    -webkit-text-stroke-width: 1px;
    -webkit-text-stroke-color: #000000;
}
</style>
"""

st.markdown(styled_text, unsafe_allow_html=True)

prompt = st.text_input(
    "Ask a question",
    key="question_input",
    placeholder="Type here...",
    type="default",
    value="",
)
generate = st.button("Generate", key="generate_button")

if generate:
    if not prompt:
        st.write("42")
    else:
        generated_text = generate_text(
            prompt, num_beams, temperature, top_p, top_k, max_length
        )
        paragraphs = generated_text.split("\n")
        styled_paragraphs = [
            f'<div class="styled-text">{paragraph}</div>' for paragraph in paragraphs
        ]
        styled_generated_text = " ".join(styled_paragraphs)
        st.markdown(styled_generated_text, unsafe_allow_html=True)