import gradio as gr
import os

from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import FAISS
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain_huggingface.llms import HuggingFacePipeline
from langchain.chains import RetrievalQA
from transformers import pipeline

# 1) Ladda och indexera alla PDF-filer
def load_vectorstore():
    all_docs, loaded_files = [], []
    for fn in os.listdir("document"):
        if fn.lower().endswith(".pdf"):
            path = os.path.join("document", fn)
            loader = PyPDFLoader(path)
            docs = loader.load_and_split()
            all_docs.extend(docs)
            loaded_files.append(fn)
    embedding = HuggingFaceEmbeddings(model_name="KBLab/sentence-bert-swedish-cased")
    vectorstore = FAISS.from_documents(all_docs, embedding)
    return vectorstore, loaded_files

# 2) Bygg UI + logik i Gradio
with gr.Blocks() as demo:
    # A) Status‐meddelande under uppstart
    status = gr.Markdown("🔄 Laddar dokument och modell, vänta…", elem_id="status-text")

    # B) Börja indexera och initiera modell
    vectorstore, files = load_vectorstore()
    llm_pipe = pipeline("text-generation", model="tiiuae/falcon-rw-1b", device=-1)
    llm = HuggingFacePipeline(
        pipeline=llm_pipe,
        model_kwargs={"temperature": 0.3, "max_new_tokens": 512}
    )
    qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=vectorstore.as_retriever())

    # C) Dölj status‐text och visa PDF‐listan när klart
    status.visible = False
    file_list_md = "\n".join(f"- {f}" for f in files)
    gr.Markdown(
        f"✅ Klar! Du kan nu ställa frågor om dokumenten nedan:\n\n{file_list_md}",
        elem_id="status-text"
    )

    # D) Temperature‐slider och Chatbot
    temp_slider = gr.Slider(
        minimum=0.0, maximum=1.0, value=0.3, step=0.05,
        label="Temperatur (kreativitetsgrad)"
    )

    def chat_fn(message, temp, history):
        # Skydd mot alltför långa frågor
        if len(message) > 1000:
            return [], [{"role":"assistant","content":
                f"⚠️ Din fråga är för lång ({len(message)} tecken). Försök korta ner den."}]
        # Uppdatera temperatur dynamiskt
        llm.model_kwargs["temperature"] = temp
        try:
            resp = qa_chain.invoke({"query": message})
            assistant_msg = resp["result"]
        except Exception as e:
            assistant_msg = f"Ett fel uppstod: {e}"
        # Returnera hela historiken i OpenAI-stil
        history = history or []
        history.append({"role":"user","content":message})
        history.append({"role":"assistant","content":assistant_msg})
        return history, history

    gr.ChatInterface(
        fn=chat_fn,
        inputs=[gr.Textbox(label="Skriv din fråga här:"), temp_slider],
        title="🌟 Dokumentassistent (Svenska)",
        description="Hej! Ställ en fråga baserat på dina PDF-dokument.",
        chatbot=gr.Chatbot(type="messages")
    )

demo.launch()