import streamlit as st
import pdfplumber
import docx
from transformers import AutoTokenizer, AutoModelForCausalLM

st.title("Jogi Dokumentumok AI Válaszgenerálás")

# Előző beszélgetések tárolása session state-ben
if "previous_conversations" not in st.session_state:
    st.session_state.previous_conversations = []

# Korábbi beszélgetések törlése
if st.checkbox("Korábbi Beszélgetések Törlése"):
    st.session_state.previous_conversations = []

# Tokenizáló modell és válaszgeneráló modell választása
tokenizer_name = st.selectbox("Tokenizáló Modell", ["bert-base-uncased", "bert-base-multilingual-cased"])
generator_name = st.selectbox("Válaszgeneráló Modell", ["microsoft/phi-2", "gpt2"])

tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
model = AutoModelForCausalLM.from_pretrained(generator_name, trust_remote_code=True)

# Dokumentum feltöltése drag and drop segítségével
document_file = st.file_uploader("Húzd ide a dokumentumot vagy kattints a feltöltéshez", type=["pdf", "docx", "doc"])

if document_file is not None:
    # Dokumentum méretének ellenőrzése
    if len(document_file.getvalue()) > 10000000:  # 10 MB-nál nagyobb dokumentumok esetén
        st.error("A dokumentum mérete túl nagy. Kérlek válassz egy kisebb méretű dokumentumot!")
    else:
        if document_file.type == "application/pdf":
            with pdfplumber.open(document_file) as pdf:
                document_text = ""
                for page in pdf.pages:
                    document_text += page.extract_text()
        elif document_file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
            docx_file = docx.Document(document_file)
            document_text = ""
            for paragraph in docx_file.paragraphs:
                document_text += paragraph.text
        elif document_file.type == "application/msword":
            doc_file = docx.Document(document_file)
            document_text = ""
            for paragraph in doc_file.paragraphs:
                document_text += paragraph.text
        else:
            st.error("A fájltípus nem támogatott. Kérlek válassz ki egy PDF, DOCX vagy DOC fájlt!")

        # Előző beszélgetésekhez csatolható kontextus
        context = st.text_area("Korábbi Beszélgetéshez Tartozó Kontextus", "")

        # Kérdés mező hozzáadása
        question = st.text_input("Kérdés a Dokumentumból", "")

        # Válaszgenerálás
        if st.button("Generálj Választ"):
            if context:
                input_text = f"{context} {document_text}"
            else:
                input_text = document_text

            # Ha van kérdés, azt is hozzáadjuk a bemenethez
            if question:
                input_text += f" {question}"

            # Model használata a válasz generálásához
            response = generate_response(input_text, tokenizer, model)

            # Válasz megjelenítése
            st.subheader("Generált Válasz:")
            st.write(response)

            # Aktuális beszélgetés hozzáadása az előző beszélgetésekhez
            st.session_state.previous_conversations.append({"input_text": input_text, "response": response})
else:
    st.warning("Kérlek válassz ki egy dokumentumfájlt!")