import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain.llms.base import LLM  # Basis-Klasse, um einen Wrapper zu erstellen
from transformers import pipeline

# LeoLM-Wrapper-Klasse, die das LeoLM-Modell via Transformers-Pipeline nutzt
class LeoLM(LLM):
    def __init__(self, max_new_tokens=512, temperature=0.5):
        self.pipeline = pipeline("text-generation", model="LeoLM/leo-mistral-hessianai-7b")
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature

    def _call(self, prompt, stop=None):
        result = self.pipeline(prompt, max_length=self.max_new_tokens, do_sample=True, temperature=self.temperature)
        return result[0]["generated_text"]

    @property
    def _identifying_params(self):
        return {"model": "LeoLM/leo-mistral-hessianai-7b"}

# PDF wird beim Start automatisch geladen
loader = PyPDFLoader("TrendingMedia_ChatbotBasis_FINAL.pdf")
documents = loader.load()
splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
texts = splitter.split_documents(documents)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
db = FAISS.from_documents(texts, embeddings)
retriever = db.as_retriever(search_kwargs={"k": 2})

# Verwende den neuen LeoLM Wrapper als LLM
llm = LeoLM(max_new_tokens=512, temperature=0.5)

qa_chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever)

def ask_question(user_input):
    if user_input.lower() in ["start", "hallo", "hi", "hey"]:
        return "👋 Willkommen bei Trending Media! Wie kann ich dir behilflich sein?"

    response = qa_chain.run(user_input)

    if response.strip() == "" or "I'm sorry" in response or len(response.split()) < 5:
        if "kontakt" in user_input.lower() or "erreichen" in user_input.lower():
            return (
                "📬 Du kannst uns direkt über dieses Formular kontaktieren:\n\n"
                "**Vorname & Nachname:**\n[_________]\n\n"
                "**E-Mail:**\n[_________]\n\n"
                "**Nachricht:**\n[__________________________]\n\n"
                "*Oder direkt über:* [📨 Kontaktformular](https://trendingmedia.ch/kontakt)"
            )
        else:
            return "❌ Das kann ich dir leider nicht beantworten. Ich bin auf Informationen aus unserem PDF spezialisiert."

    return response

with gr.Blocks() as demo:
    gr.Markdown("## 🤖 TrendingBot\nWillkommen bei Trending Media! Stelle mir deine Frage.")
    user_input = gr.Textbox(label="Deine Frage", placeholder="Frag mich etwas über unsere Lösungen...")
    bot_response = gr.Textbox(label="TrendingBot antwortet")
    user_input.submit(ask_question, inputs=user_input, outputs=bot_response)

demo.launch()