import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
import fitz  # PyMuPDF PDF olvasáshoz
from docx import Document  # DOCX olvasáshoz

# Fordító modell betöltése (M2M100 fordítás magyar és angol között)
translator_tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M")
translator_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/m2m100_418M")
translation_pipeline = pipeline("translation", model=translator_model, tokenizer=translator_tokenizer)

# Generatív modell betöltése (angol nyelven működik)
gpt_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
gpt_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
generate_answer = pipeline("text2text-generation", model=gpt_model, tokenizer=gpt_tokenizer)

# Szöveg feldarabolása max hosszal
def chunk_text(text, max_length=512):
    words = text.split()
    chunks = []
    while len(words) > max_length:
        chunk = " ".join(words[:max_length])
        chunks.append(chunk)
        words = words[max_length:]
    chunks.append(" ".join(words))
    return chunks

# Fordítás funkciók
def translate_to_english(text):
    chunks = chunk_text(text, max_length=512)
    translated_text = ""
    for chunk in chunks:
        translation = translation_pipeline(chunk, src_lang="hu", tgt_lang="en", max_length=512)[0]['translation_text']
        translated_text += " " + translation
    return translated_text.strip()

def translate_to_hungarian(text):
    chunks = chunk_text(text, max_length=512)
    translated_text = ""
    for chunk in chunks:
        translation = translation_pipeline(chunk, src_lang="en", tgt_lang="hu", max_length=512)[0]['translation_text']
        translated_text += " " + translation
    return translated_text.strip()

# PDF szöveg kinyerése
def extract_text_from_pdf(pdf_file):
    text = ""
    with fitz.open("pdf", pdf_file.read()) as doc:
        for page in doc:
            text += page.get_text()
    return text

# DOCX szöveg kinyerése
def extract_text_from_docx(docx_file):
    doc = Document(docx_file)
    text = "\n".join([paragraph.text for paragraph in doc.paragraphs])
    return text

# Streamlit alkalmazás
st.title("AI Jogi Chatbot - Fordítás Alapú Megközelítéssel")

# Dokumentum feltöltése
uploaded_file = st.file_uploader("Töltsön fel egy dokumentumot", type=["txt", "pdf", "docx"])

if uploaded_file:
    file_type = uploaded_file.type

    # Dokumentum típus szerinti beolvasás
    if file_type == "application/pdf":
        context = extract_text_from_pdf(uploaded_file)
    elif file_type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
        context = extract_text_from_docx(uploaded_file)
    else:
        # TXT fájl feltöltése esetén
        try:
            context = uploaded_file.read().decode("utf-8")
        except UnicodeDecodeError:
            context = uploaded_file.read().decode("ISO-8859-1")

    # Teljes dokumentum megjelenítése
    st.write("### Feltöltött dokumentum tartalma:")
    st.write(context)
    
    # Dokumentum fordítása angolra
    with st.spinner('Fordítás angolra...'):
        context_english = translate_to_english(context)
    
    # Chatelés a dokumentum kontextusában
    st.write("### Chat")
    question = st.text_input("Írja be a kérdését a dokumentumról:")
    
    if question:
        with st.spinner('Válasz generálása...'):
            # Kérdés fordítása angolra
            question_english = translate_to_english(question)
            
            # Generatív válaszadás angol nyelven
            input_text = f"{context_english}\n\nQuestion: {question_english}\nAnswer:"
            response_english = generate_answer(input_text, max_length=150)[0]["generated_text"]
            
            # Válasz fordítása magyarra
            response_hungarian = translate_to_hungarian(response_english)
        
        st.write("Válasz:", response_hungarian)