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

st.title("JAI")

# 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 = AutoTokenizer.from_pretrained("pile-of-law/legalbert-large-1.7M-2")
model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2", trust_remote_code=True)

# Dokumentumfeldolgozó függvény
def process_document(document_file):
    document_text = ""
    if document_file.type == "application/pdf":
        with pdfplumber.open(document_file) as pdf:
            for page in pdf.pages:
                text = page.extract_text()
                if text:
                    document_text += text.strip() + "\n\n"
    elif document_file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
        docx_file = docx.Document(document_file)
        for paragraph in docx_file.paragraphs:
            text = paragraph.text
            if text:
                document_text += text.strip() + "\n\n"
    elif document_file.type == "application/msword":
        doc_file = docx.Document(document_file)
        for paragraph in doc_file.paragraphs:
            text = paragraph.text
            if text:
                document_text += text.strip() + "\n\n"
    else:
        st.error("A fájltípus nem támogatott. Kérlek válassz ki egy PDF, DOCX vagy DOC fájlt!")
    return document_text

# Válaszgeneráló függvény
def generate_response(input_text):
    # Bemeneti szöveg tisztítása
    cleaned_input = input_text.strip()  

    # Tokenizálás
    inputs = tokenizer(cleaned_input, return_tensors="pt", max_length=1024, truncation=True)

    # Modell használata a válasz generálásához
    outputs = model(**inputs)

    # Válasz visszaadása
    response = tokenizer.decode(outputs.logits.argmax(dim=1)[0])
    return response

# 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:
    document_text = process_document(document_file)

    # 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", "")

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

    # Válasz generálása csak akkor, ha a felhasználó elküldi a promptot
    if input_text.strip() != "":
        response = generate_response(input_text)
        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})

        # FAISS használata az előző beszélgetések indexelésére és keresésére
        # Először hozz létre egy üres FAISS indexet
        index = faiss.IndexFlatL2(embedding_dimension)  # Az embedding_dimension a vektormérete az előző beszélgetések reprezentálásához

        # Az előző beszélgetéseket reprezentáló vektorokat hozzáadhatod az indexhez
        for conversation_vector in previous_conversation_vectors:
            index.add(conversation_vector)

        # Keresés egy adott lekérdezés vektor alapján
        query_vector = generate_vector_from_text(input_text)  # A lekérdezés szövegből létrehozott vektor
        k = 5  # A legjobb 5 hasonlóságú beszélgetés visszaadása
        D, I = index.search(query_vector, k)  # D: távolságok, I: indexek

        # Azon beszélgetések visszaadása, amelyek a legjobban illeszkednek a lekérdezéshez
        relevant_conversations = [st.session_state.previous_conversations[i] for i in I]