import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from sentence_transformers import SentenceTransformer
import faiss
import torch
import os

from io import BytesIO
from PyPDF2 import PdfReader
from huggingface_hub import login
import os

login(token=os.getenv("HF_TOKEN"))

# Konfiguracja modelu Llama 3 (dostępnego w HF Hub)
#LLAMA_MODEL = "meta-llama/Meta-Llama-3-8B"
#LLAMA_MODEL = "mistralai/Mixtral-8x7B-Instruct-v0.1"
LLAMA_MODEL = "NousResearch/Llama-2-7b-chat-hf"
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

# Inicjalizacja modeli (cache dla Hugging Face)
@st.cache_resource
def load_models():
    tokenizer = AutoTokenizer.from_pretrained(LLAMA_MODEL)
    model = AutoModelForCausalLM.from_pretrained(LLAMA_MODEL, torch_dtype=torch.float16, device_map="auto")
    embedding_model = SentenceTransformer(EMBEDDING_MODEL)
    return tokenizer, model, embedding_model

tokenizer, llama_model, embedding_model = load_models()

# Funkcja przetwarzająca plik
def process_file(uploaded_file):
    if uploaded_file.name.endswith(".pdf"):
        pdf = PdfReader(uploaded_file)
        text = ""
        for page in pdf.pages:
            text += page.extract_text() + "\n"
    elif uploaded_file.name.endswith(".txt"):
        text = uploaded_file.read().decode("utf-8")
    else:
        st.warning("Obsługiwane są tylko pliki PDF i TXT.")
        return []
    
    # Dziel tekst na fragmenty
    chunks = [text[i:i+500] for i in range(0, len(text), 500)]
    embeddings = embedding_model.encode(chunks)
    
    # Zbuduj wektorową bazę danych (FAISS)
    index = faiss.IndexFlatL2(embeddings.shape[1])
    index.add(embeddings)
    
    return chunks, embeddings, index

# UI
st.title("🦙 Llama 3 RAG App")
st.write("Załaduj plik, a następnie zapytaj Llama 3 o informacje z niego!")

uploaded_file = st.file_uploader("Wgraj plik PDF lub TXT", type=["pdf", "txt"])

if uploaded_file:
    chunks, embeddings, index = process_file(uploaded_file)
    st.success("Plik przetworzony! Możesz teraz zadać pytanie.")

    user_question = st.text_input("Twoje pytanie:")

    if user_question:
        question_embedding = embedding_model.encode([user_question])
        D, I = index.search(question_embedding, k=3)
        relevant_chunks = "\n".join([chunks[i] for i in I[0]])

        prompt = f"""Odpowiedz na pytanie korzystając z kontekstu poniżej.

### Kontekst:
{relevant_chunks}

### Pytanie:
{user_question}

### Odpowiedź:"""

        # Generowanie odpowiedzi z Llama 3
        inputs = tokenizer(prompt, return_tensors="pt").to(llama_model.device)
        outputs = llama_model.generate(**inputs, max_new_tokens=200, do_sample=True, temperature=0.7)
        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)

        st.markdown("### Odpowiedź Llama 3:")
        st.write(answer)